From 7646bd8fd4de084ad17a6803e40f363ce3f16068 Mon Sep 17 00:00:00 2001
From: Juan Font Alonso <juanfontalonso@gmail.com>
Date: Wed, 28 Apr 2021 16:55:29 +0200
Subject: [PATCH] CLI stuff should be in the cmd/ path

---
 cli.go                      | 64 ----------------------------------
 cmd/headscale/cli/routes.go |  3 +-
 routes.go                   | 69 +++++++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 65 deletions(-)
 create mode 100644 routes.go

diff --git a/cli.go b/cli.go
index 6615d3d9..0cb53336 100644
--- a/cli.go
+++ b/cli.go
@@ -1,13 +1,9 @@
 package headscale
 
 import (
-	"encoding/json"
-	"errors"
 	"fmt"
 	"log"
 
-	"github.com/jinzhu/gorm/dialects/postgres"
-	"inet.af/netaddr"
 	"tailscale.com/wgengine/wgcfg"
 )
 
@@ -51,63 +47,3 @@ func (h *Headscale) RegisterMachine(key string, namespace string) error {
 	fmt.Println("Machine registered 🎉")
 	return nil
 }
-
-// ListNodeRoutes prints the subnet routes advertised by a node (identified by
-// namespace and node name)
-func (h *Headscale) ListNodeRoutes(namespace string, nodeName string) error {
-	m, err := h.GetMachine(namespace, nodeName)
-	if err != nil {
-		return err
-	}
-
-	hi, err := m.GetHostInfo()
-	if err != nil {
-		return err
-	}
-	fmt.Println(hi.RoutableIPs)
-	return nil
-}
-
-// EnableNodeRoute enables a subnet route advertised by a node (identified by
-// namespace and node name)
-func (h *Headscale) EnableNodeRoute(namespace string, nodeName string, routeStr string) error {
-	m, err := h.GetMachine(namespace, nodeName)
-	if err != nil {
-		return err
-	}
-	hi, err := m.GetHostInfo()
-	if err != nil {
-		return err
-	}
-	route, err := netaddr.ParseIPPrefix(routeStr)
-	if err != nil {
-		return err
-	}
-
-	for _, rIP := range hi.RoutableIPs {
-		if rIP == route {
-			db, err := h.db()
-			if err != nil {
-				log.Printf("Cannot open DB: %s", err)
-				return err
-			}
-
-			routes, _ := json.Marshal([]string{routeStr}) // TODO: only one for the time being, so overwriting the rest
-			m.EnabledRoutes = postgres.Jsonb{RawMessage: json.RawMessage(routes)}
-			db.Save(&m)
-			db.Close()
-
-			peers, _ := h.getPeers(*m)
-			h.pollMu.Lock()
-			for _, p := range *peers {
-				if pUp, ok := h.clientsPolling[uint64(p.ID)]; ok {
-					pUp <- []byte{}
-				}
-			}
-			h.pollMu.Unlock()
-			return nil
-		}
-	}
-
-	return errors.New("could not find routable range")
-}
diff --git a/cmd/headscale/cli/routes.go b/cmd/headscale/cli/routes.go
index d2f58817..c14956f6 100644
--- a/cmd/headscale/cli/routes.go
+++ b/cmd/headscale/cli/routes.go
@@ -21,11 +21,12 @@ var ListRoutesCmd = &cobra.Command{
 		if err != nil {
 			log.Fatalf("Error initializing: %s", err)
 		}
-		err = h.ListNodeRoutes(args[0], args[1])
+		routes, err := h.GetNodeRoutes(args[0], args[1])
 		if err != nil {
 			fmt.Println(err)
 			return
 		}
+		fmt.Println(routes)
 	},
 }
 
diff --git a/routes.go b/routes.go
new file mode 100644
index 00000000..2894366a
--- /dev/null
+++ b/routes.go
@@ -0,0 +1,69 @@
+package headscale
+
+import (
+	"encoding/json"
+	"errors"
+	"log"
+
+	"github.com/jinzhu/gorm/dialects/postgres"
+	"inet.af/netaddr"
+)
+
+// GetNodeRoutes returns the subnet routes advertised by a node (identified by
+// namespace and node name)
+func (h *Headscale) GetNodeRoutes(namespace string, nodeName string) (*[]netaddr.IPPrefix, error) {
+	m, err := h.GetMachine(namespace, nodeName)
+	if err != nil {
+		return nil, err
+	}
+
+	hi, err := m.GetHostInfo()
+	if err != nil {
+		return nil, err
+	}
+	return &hi.RoutableIPs, nil
+}
+
+// EnableNodeRoute enables a subnet route advertised by a node (identified by
+// namespace and node name)
+func (h *Headscale) EnableNodeRoute(namespace string, nodeName string, routeStr string) error {
+	m, err := h.GetMachine(namespace, nodeName)
+	if err != nil {
+		return err
+	}
+	hi, err := m.GetHostInfo()
+	if err != nil {
+		return err
+	}
+	route, err := netaddr.ParseIPPrefix(routeStr)
+	if err != nil {
+		return err
+	}
+
+	for _, rIP := range hi.RoutableIPs {
+		if rIP == route {
+			db, err := h.db()
+			if err != nil {
+				log.Printf("Cannot open DB: %s", err)
+				return err
+			}
+
+			routes, _ := json.Marshal([]string{routeStr}) // TODO: only one for the time being, so overwriting the rest
+			m.EnabledRoutes = postgres.Jsonb{RawMessage: json.RawMessage(routes)}
+			db.Save(&m)
+			db.Close()
+
+			peers, _ := h.getPeers(*m)
+			h.pollMu.Lock()
+			for _, p := range *peers {
+				if pUp, ok := h.clientsPolling[uint64(p.ID)]; ok {
+					pUp <- []byte{}
+				}
+			}
+			h.pollMu.Unlock()
+			return nil
+		}
+	}
+
+	return errors.New("could not find routable range")
+}