From 176eea4a845c16a542898a9e793f6ffee83dee8d Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Fri, 23 Apr 2021 00:25:01 +0200
Subject: [PATCH] WIP on PreAuthKeys

---
 cmd/headscale/headscale.go | 43 ++++++++++++++++++++-
 db.go                      |  1 +
 preauth_keys.go            | 78 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 121 insertions(+), 1 deletion(-)
 create mode 100644 preauth_keys.go

diff --git a/cmd/headscale/headscale.go b/cmd/headscale/headscale.go
index ebad844..fe656a0 100644
--- a/cmd/headscale/headscale.go
+++ b/cmd/headscale/headscale.go
@@ -169,6 +169,43 @@ var enableRouteCmd = &cobra.Command{
 	},
 }
 
+var preauthkeysCmd = &cobra.Command{
+	Use:   "preauthkey",
+	Short: "Handle the preauthkeys in Headscale",
+}
+
+var listPreAuthKeys = &cobra.Command{
+	Use:   "list NAMESPACE",
+	Short: "List the preauthkeys for this namespace",
+	Args: func(cmd *cobra.Command, args []string) error {
+		if len(args) < 1 {
+			return fmt.Errorf("Missing parameters")
+		}
+		return nil
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		h, err := getHeadscaleApp()
+		if err != nil {
+			log.Fatalf("Error initializing: %s", err)
+		}
+		keys, err := h.GetPreAuthKeys(args[0])
+		if err != nil {
+			fmt.Println(err)
+			return
+		}
+		for _, k := range *keys {
+			fmt.Printf(
+				"key: %s, namespace: %s, reusable: %v, expiration: %s, created_at: %s",
+				k.Key,
+				k.Namespace.Name,
+				k.Reusable,
+				k.Expiration.Format("2006-01-02 15:04:05"),
+				k.CreatedAt.Format("2006-01-02 15:04:05"),
+			)
+		}
+	},
+}
+
 func main() {
 	viper.SetConfigName("config")
 	viper.AddConfigPath("/etc/headscale/")
@@ -183,14 +220,18 @@ func main() {
 	headscaleCmd.AddCommand(versionCmd)
 	headscaleCmd.AddCommand(serveCmd)
 	headscaleCmd.AddCommand(registerCmd)
+	headscaleCmd.AddCommand(preauthkeysCmd)
 	headscaleCmd.AddCommand(namespaceCmd)
+	headscaleCmd.AddCommand(nodeCmd)
+
 	namespaceCmd.AddCommand(createNamespaceCmd)
 	namespaceCmd.AddCommand(listNamespacesCmd)
 
-	headscaleCmd.AddCommand(nodeCmd)
 	nodeCmd.AddCommand(listRoutesCmd)
 	nodeCmd.AddCommand(enableRouteCmd)
 
+	preauthkeysCmd.AddCommand(listPreAuthKeys)
+
 	if err := headscaleCmd.Execute(); err != nil {
 		fmt.Println(err)
 		os.Exit(-1)
diff --git a/db.go b/db.go
index 7e1e0cc..a100556 100644
--- a/db.go
+++ b/db.go
@@ -24,6 +24,7 @@ func (h *Headscale) initDB() error {
 	db.AutoMigrate(&Machine{})
 	db.AutoMigrate(&KV{})
 	db.AutoMigrate(&Namespace{})
+	db.AutoMigrate(&PreAuthKey{})
 	db.Close()
 
 	h.setValue("db_version", dbVersion)
diff --git a/preauth_keys.go b/preauth_keys.go
new file mode 100644
index 0000000..14fcdbb
--- /dev/null
+++ b/preauth_keys.go
@@ -0,0 +1,78 @@
+package headscale
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"log"
+	"time"
+)
+
+type PreAuthKey struct {
+	ID          uint64 `gorm:"primary_key"`
+	Key         string
+	NamespaceID uint
+	Namespace   Namespace
+	Reusable    bool
+
+	CreatedAt  *time.Time
+	Expiration *time.Time
+}
+
+func (h *Headscale) CreatePreAuthKey(namespaceName string, reusable bool, expiration *time.Time) (*PreAuthKey, error) {
+	n, err := h.GetNamespace(namespaceName)
+	if err != nil {
+		return nil, err
+	}
+
+	db, err := h.db()
+	if err != nil {
+		log.Printf("Cannot open DB: %s", err)
+		return nil, err
+	}
+	defer db.Close()
+
+	now := time.Now().UTC()
+	kstr, err := h.generateKey()
+	if err != nil {
+		return nil, err
+	}
+
+	k := PreAuthKey{
+		Key:         kstr,
+		NamespaceID: n.ID,
+		Reusable:    reusable,
+		CreatedAt:   &now,
+		Expiration:  expiration,
+	}
+	db.Save(&k)
+
+	return &k, nil
+}
+
+func (h *Headscale) GetPreAuthKeys(namespaceName string) (*[]PreAuthKey, error) {
+	n, err := h.GetNamespace(namespaceName)
+	if err != nil {
+		return nil, err
+	}
+	db, err := h.db()
+	if err != nil {
+		log.Printf("Cannot open DB: %s", err)
+		return nil, err
+	}
+	defer db.Close()
+
+	keys := []PreAuthKey{}
+	if err := db.Where(&PreAuthKey{NamespaceID: n.ID}).Find(&keys).Error; err != nil {
+		return nil, err
+	}
+	return &keys, nil
+}
+
+func (h *Headscale) generateKey() (string, error) {
+	size := 24
+	bytes := make([]byte, size)
+	if _, err := rand.Read(bytes); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(bytes), nil
+}