From b1d06f3ffd60d6fa08dca284df4b3eac091aee55 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Sun, 21 Feb 2021 01:30:03 +0100
Subject: [PATCH] headscale now has a CLI - registration of machines occurs
 there

---
 app.go                     | 36 +++++++++++++++++++
 cmd/headscale/headscale.go | 72 ++++++++++++++++++++++++++++++++++++--
 go.mod                     |  1 +
 go.sum                     |  2 ++
 handlers.go                | 51 ++++++++++-----------------
 5 files changed, 127 insertions(+), 35 deletions(-)

diff --git a/app.go b/app.go
index 3545cd83..b08e3bda 100644
--- a/app.go
+++ b/app.go
@@ -3,6 +3,7 @@ package headscale
 import (
 	"fmt"
 	"io/ioutil"
+	"log"
 
 	"github.com/gin-gonic/gin"
 	"tailscale.com/tailcfg"
@@ -65,3 +66,38 @@ func (h *Headscale) Serve() error {
 	err := r.Run(h.cfg.Addr)
 	return err
 }
+
+func (h *Headscale) RegisterMachine(key string) error {
+	mKey, err := wgcfg.ParseHexKey(key)
+	if err != nil {
+		log.Printf("Cannot parse client key: %s", err)
+		return err
+	}
+	db, err := h.db()
+	if err != nil {
+		log.Printf("Cannot open DB: %s", err)
+		return err
+	}
+	defer db.Close()
+	m := Machine{}
+	if db.First(&m, "machine_key = ?", mKey.HexString()).RecordNotFound() {
+		log.Printf("Cannot find machine with machine key: %s", mKey.Base64())
+		return err
+	}
+
+	if m.isAlreadyRegistered() {
+		fmt.Println("This machine already registered")
+		return nil
+	}
+
+	ip, err := h.getAvailableIP()
+	if err != nil {
+		log.Println(err)
+		return err
+	}
+	m.IPAddress = ip.String()
+	m.Registered = true
+	db.Save(&m)
+	fmt.Println("Machine registered 🎉")
+	return nil
+}
diff --git a/cmd/headscale/headscale.go b/cmd/headscale/headscale.go
index ef895bf9..3b3ea123 100644
--- a/cmd/headscale/headscale.go
+++ b/cmd/headscale/headscale.go
@@ -1,16 +1,72 @@
 package main
 
 import (
+	"fmt"
 	"io"
 	"log"
 	"os"
 
 	"github.com/juanfont/headscale"
+	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 	"gopkg.in/yaml.v2"
 	"tailscale.com/tailcfg"
 )
 
+const version = "0.1"
+
+var versionCmd = &cobra.Command{
+	Use:   "version",
+	Short: "Print the version.",
+	Long:  "The version of headscale.",
+	Run: func(cmd *cobra.Command, args []string) {
+		fmt.Println(version)
+	},
+}
+
+var headscaleCmd = &cobra.Command{
+	Use:   "headscale",
+	Short: "headscale - a Tailscale control server",
+	Long: fmt.Sprintf(`
+headscale is an open source implementation of the Tailscale control server
+
+Juan Font Alonso <juanfontalonso@gmail.com> - 2021
+https://gitlab.com/juanfont/headscale`),
+}
+
+var serveCmd = &cobra.Command{
+	Use:   "serve",
+	Short: "Launches the headscale server",
+	Args: func(cmd *cobra.Command, args []string) error {
+		return nil
+	},
+	Run: func(cmd *cobra.Command, args []string) {
+		h, err := getHeadscaleApp()
+		if err != nil {
+			log.Fatalf("Error initializing: %s", err)
+		}
+		h.Serve()
+	},
+}
+
+var registerCmd = &cobra.Command{
+	Use:   "register machineID",
+	Short: "Registers a machine to your network",
+	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)
+		}
+		h.RegisterMachine(args[0])
+	},
+}
+
 func main() {
 	viper.SetConfigName("config")
 	viper.AddConfigPath(".")
@@ -20,6 +76,18 @@ func main() {
 		log.Fatalf("Fatal error config file: %s \n", err)
 	}
 
+	headscaleCmd.AddCommand(versionCmd)
+	headscaleCmd.AddCommand(serveCmd)
+	headscaleCmd.AddCommand(registerCmd)
+
+	if err := headscaleCmd.Execute(); err != nil {
+		fmt.Println(err)
+		os.Exit(-1)
+	}
+
+}
+
+func getHeadscaleApp() (*headscale.Headscale, error) {
 	derpMap, err := loadDerpMap(viper.GetString("derp_map_path"))
 	if err != nil {
 		log.Printf("Could not load DERP servers map file: %s", err)
@@ -39,9 +107,9 @@ func main() {
 	}
 	h, err := headscale.NewHeadscale(cfg)
 	if err != nil {
-		log.Fatalln(err)
+		return nil, err
 	}
-	h.Serve()
+	return h, nil
 }
 
 func loadDerpMap(path string) (*tailcfg.DERPMap, error) {
diff --git a/go.mod b/go.mod
index 2be6955f..1a054379 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/gin-gonic/gin v1.6.3
 	github.com/jinzhu/gorm v1.9.16
 	github.com/klauspost/compress v1.11.7
+	github.com/spf13/cobra v0.0.2-0.20171109065643-2da4a54c5cee
 	github.com/spf13/viper v1.7.1
 	golang.org/x/crypto v0.0.0-20210220033148-5ea612d1eb83
 	gopkg.in/yaml.v2 v2.2.8
diff --git a/go.sum b/go.sum
index c4bb0cd5..c2c463c3 100644
--- a/go.sum
+++ b/go.sum
@@ -232,6 +232,7 @@ github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpO
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/imdario/mergo v0.3.5/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
 github.com/imdario/mergo v0.3.8/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJh5FfA=
+github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
 github.com/jinzhu/gorm v1.9.16 h1:+IyIjPEABKRpsu/F8OvDPy9fyQlgsg2luMV2ZIH5i5o=
 github.com/jinzhu/gorm v1.9.16/go.mod h1:G3LB3wezTOWM2ITLzPxEXgSkOXAntiLHS7UdBefADcs=
@@ -380,6 +381,7 @@ github.com/spf13/afero v1.2.2 h1:5jhuqJyZCZf2JRofRvN/nIFgIWNzPa3/Vz8mYylgbWc=
 github.com/spf13/afero v1.2.2/go.mod h1:9ZxEEn6pIJ8Rxe320qSDBk6AsU0r9pR7Q4OcevTdifk=
 github.com/spf13/cast v1.3.0 h1:oget//CVOEoFewqQxwr0Ej5yjygnqGkvggSE/gB35Q8=
 github.com/spf13/cast v1.3.0/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE=
+github.com/spf13/cobra v0.0.2-0.20171109065643-2da4a54c5cee h1:GQkkv3XSnxhAMjdq2wLfEnptEVr+2BNvmHizILHn+d4=
 github.com/spf13/cobra v0.0.2-0.20171109065643-2da4a54c5cee/go.mod h1:1l0Ry5zgKvJasoi3XT1TypsSe7PqH0Sj9dhYf7v3XqQ=
 github.com/spf13/jwalterweatherman v1.0.0 h1:XHEdyB+EcvlqZamSM4ZOMGlc93t6AcsBEu9Gc1vn7yk=
 github.com/spf13/jwalterweatherman v1.0.0/go.mod h1:cQK4TGJAtQXfYWX+Ddv3mKDzgVb68N+wFjFa4jdeBTo=
diff --git a/handlers.go b/handlers.go
index a7036041..d805ec2d 100644
--- a/handlers.go
+++ b/handlers.go
@@ -262,41 +262,26 @@ func (h *Headscale) RegisterWebAPI(c *gin.Context) {
 		c.String(http.StatusBadRequest, "Wrong params")
 		return
 	}
-	mKey, err := wgcfg.ParseHexKey(mKeyStr)
-	if err != nil {
-		log.Printf("Cannot parse client key: %s", err)
-		c.String(http.StatusInternalServerError, "Sad!")
-		return
-	}
-	db, err := h.db()
-	if err != nil {
-		log.Printf("Cannot open DB: %s", err)
-		c.String(http.StatusInternalServerError, ":(")
-		return
-	}
-	defer db.Close()
-	m := Machine{}
-	if db.First(&m, "machine_key = ?", mKey.HexString()).RecordNotFound() {
-		log.Printf("Cannot find machine with machine key: %s", mKey.Base64())
-		c.String(http.StatusNotFound, "Sad!")
-		return
-	}
 
-	if !m.isAlreadyRegistered() {
-		ip, err := h.getAvailableIP()
-		if err != nil {
-			log.Println(err)
-			c.String(http.StatusInternalServerError, "Upsy dupsy")
-			return
-		}
-		m.IPAddress = ip.String()
-		m.Registered = true // very naive 😱
-		db.Save(&m)
+	c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(fmt.Sprintf(`
+	<html>
+	<body>
+	<h1>headscale</h1>
+	<p>
+		Run the command below in the headscale server to add this machine to your network:
+	</p>
 
-		c.JSON(http.StatusOK, gin.H{"msg": "Ook"})
-		return
-	}
-	c.JSON(http.StatusOK, gin.H{"msg": "Eek"})
+	<p>
+		<code>
+			<b>headscale register %s</b>
+		</code>
+	</p>
+
+	</body>
+	</html>
+	
+	`, mKeyStr)))
+	return
 }
 
 func (h *Headscale) handleNewServer(c *gin.Context, db *gorm.DB, idKey wgcfg.Key, req tailcfg.RegisterRequest) {