From 77f5f8bd1cedd9c5bb3cbb20c8918b18ddfbaf08 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Thu, 4 Nov 2021 22:14:39 +0000
Subject: [PATCH] Simplify and streamline preauth commands for new cli/rpc/api

---
 preauth_keys.go      | 40 ++++++++++++++++++++++++++++++++--------
 preauth_keys_test.go | 10 +++++-----
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/preauth_keys.go b/preauth_keys.go
index de10cdb7..304e345b 100644
--- a/preauth_keys.go
+++ b/preauth_keys.go
@@ -4,14 +4,20 @@ import (
 	"crypto/rand"
 	"encoding/hex"
 	"errors"
+	"strconv"
 	"time"
 
+	"google.golang.org/protobuf/types/known/timestamppb"
 	"gorm.io/gorm"
+
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 )
 
-const errorAuthKeyNotFound = Error("AuthKey not found")
-const errorAuthKeyExpired = Error("AuthKey expired")
-const errSingleUseAuthKeyHasBeenUsed = Error("AuthKey has already been used")
+const (
+	errorAuthKeyNotFound           = Error("AuthKey not found")
+	errorAuthKeyExpired            = Error("AuthKey expired")
+	errSingleUseAuthKeyHasBeenUsed = Error("AuthKey has already been used")
+)
 
 // PreAuthKey describes a pre-authorization key usable in a particular namespace
 type PreAuthKey struct {
@@ -28,7 +34,12 @@ type PreAuthKey struct {
 }
 
 // CreatePreAuthKey creates a new PreAuthKey in a namespace, and returns it
-func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, ephemeral bool, expiration *time.Time) (*PreAuthKey, error) {
+func (h *Headscale) CreatePreAuthKey(
+	namespaceName string,
+	reusable bool,
+	ephemeral bool,
+	expiration *time.Time,
+) (*PreAuthKey, error) {
 	n, err := h.GetNamespace(namespaceName)
 	if err != nil {
 		return nil, err
@@ -54,8 +65,8 @@ func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, epheme
 	return &k, nil
 }
 
-// GetPreAuthKeys returns the list of PreAuthKeys for a namespace
-func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error) {
+// ListPreAuthKeys returns the list of PreAuthKeys for a namespace
+func (h *Headscale) ListPreAuthKeys(namespaceName string) ([]PreAuthKey, error) {
 	n, err := h.GetNamespace(namespaceName)
 	if err != nil {
 		return nil, err
@@ -65,7 +76,7 @@ func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error)
 	if err := h.db.Preload("Namespace").Where(&PreAuthKey{NamespaceID: n.ID}).Find(&keys).Error; err != nil {
 		return nil, err
 	}
-	return &keys, nil
+	return keys, nil
 }
 
 // GetPreAuthKey returns a PreAuthKey for a given key
@@ -83,7 +94,7 @@ func (h *Headscale) GetPreAuthKey(namespace string, key string) (*PreAuthKey, er
 }
 
 // MarkExpirePreAuthKey marks a PreAuthKey as expired
-func (h *Headscale) MarkExpirePreAuthKey(k *PreAuthKey) error {
+func (h *Headscale) ExpirePreAuthKey(k *PreAuthKey) error {
 	if err := h.db.Model(&k).Update("Expiration", time.Now()).Error; err != nil {
 		return err
 	}
@@ -126,3 +137,16 @@ func (h *Headscale) generateKey() (string, error) {
 	}
 	return hex.EncodeToString(bytes), nil
 }
+
+func (key *PreAuthKey) toProto() *v1.PreAuthKey {
+	return &v1.PreAuthKey{
+		Namespace:  key.Namespace.Name,
+		Id:         strconv.FormatUint(key.ID, 10),
+		Key:        key.Key,
+		Resuable:   key.Reusable,
+		Ephemeral:  key.Ephemeral,
+		Used:       key.Used,
+		Expiration: timestamppb.New(*key.Expiration),
+		CreatedAt:  timestamppb.New(*key.CreatedAt),
+	}
+}
diff --git a/preauth_keys_test.go b/preauth_keys_test.go
index f8973eaf..dceec00b 100644
--- a/preauth_keys_test.go
+++ b/preauth_keys_test.go
@@ -24,15 +24,15 @@ func (*Suite) TestCreatePreAuthKey(c *check.C) {
 	// Make sure the Namespace association is populated
 	c.Assert(k.Namespace.Name, check.Equals, n.Name)
 
-	_, err = h.GetPreAuthKeys("bogus")
+	_, err = h.ListPreAuthKeys("bogus")
 	c.Assert(err, check.NotNil)
 
-	keys, err := h.GetPreAuthKeys(n.Name)
+	keys, err := h.ListPreAuthKeys(n.Name)
 	c.Assert(err, check.IsNil)
-	c.Assert(len(*keys), check.Equals, 1)
+	c.Assert(len(keys), check.Equals, 1)
 
 	// Make sure the Namespace association is populated
-	c.Assert((*keys)[0].Namespace.Name, check.Equals, n.Name)
+	c.Assert((keys)[0].Namespace.Name, check.Equals, n.Name)
 }
 
 func (*Suite) TestExpiredPreAuthKey(c *check.C) {
@@ -172,7 +172,7 @@ func (*Suite) TestExpirePreauthKey(c *check.C) {
 	c.Assert(err, check.IsNil)
 	c.Assert(pak.Expiration, check.IsNil)
 
-	err = h.MarkExpirePreAuthKey(pak)
+	err = h.ExpirePreAuthKey(pak)
 	c.Assert(err, check.IsNil)
 	c.Assert(pak.Expiration, check.NotNil)