Convert nodes command to use gRPC

cmd/headscale/cli/nodes.go

@@ -1,14 +1,14 @@
 package cli
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"strconv"
-	"strings"
 	"time"
 
 	survey "github.com/AlecAivazis/survey/v2"
-	"github.com/juanfont/headscale"
+	v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
 	"github.com/pterm/pterm"
 	"github.com/spf13/cobra"
 	"tailscale.com/tailcfg"
@@ -73,30 +73,37 @@ var registerNodeCmd = &cobra.Command{
 	Use:   "register",
 	Short: "Registers a machine to your network",
 	Run: func(cmd *cobra.Command, args []string) {
-		n, err := cmd.Flags().GetString("namespace")
+		output, _ := cmd.Flags().GetString("output")
+		namespace, err := cmd.Flags().GetString("namespace")
 		if err != nil {
-			log.Fatalf("Error getting namespace: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
+			return
 		}
-		o, _ := cmd.Flags().GetString("output")
 
-		h, err := getHeadscaleApp()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		machineKey, err := cmd.Flags().GetString("key")
 		if err != nil {
-			log.Fatalf("Error initializing: %s", err)
-		}
-		machineIDStr, err := cmd.Flags().GetString("key")
-		if err != nil {
-			log.Fatalf("Error getting machine ID: %s", err)
-		}
-		m, err := h.RegisterMachine(machineIDStr, n)
-		if strings.HasPrefix(o, "json") {
-			JsonOutput(m, err, o)
+			ErrorOutput(err, fmt.Sprintf("Error getting machine key from flag: %s", err), output)
 			return
 		}
+
+		request := &v1.RegisterMachineRequest{
+			Key:       machineKey,
+			Namespace: namespace,
+		}
+
+		response, err := client.RegisterMachine(ctx, request)
 		if err != nil {
-			fmt.Printf("Cannot register machine: %s\n", err)
+			ErrorOutput(err, fmt.Sprintf("Cannot register machine: %s\n", err), output)
 			return
 		}
-		fmt.Printf("Machine registered\n")
+
+		SuccessOutput(response.Machine, "Machine register", output)
 	},
 }
 
@@ -104,72 +111,44 @@ var listNodesCmd = &cobra.Command{
 	Use:   "list",
 	Short: "List nodes",
 	Run: func(cmd *cobra.Command, args []string) {
-		n, err := cmd.Flags().GetString("namespace")
+		output, _ := cmd.Flags().GetString("output")
+		namespace, err := cmd.Flags().GetString("namespace")
 		if err != nil {
-			log.Fatalf("Error getting namespace: %s", err)
-		}
-		o, _ := cmd.Flags().GetString("output")
-
-		h, err := getHeadscaleApp()
-		if err != nil {
-			log.Fatalf("Error initializing: %s", err)
-		}
-
-		var namespaces []headscale.Namespace
-		var namespace *headscale.Namespace
-		var sharedMachines *[]headscale.Machine
-		if len(n) == 0 {
-			// no namespace provided, list all
-			tmp, err := h.ListNamespaces()
-			if err != nil {
-				log.Fatalf("Error fetching namespace: %s", err)
-			}
-			namespaces = *tmp
-		} else {
-			namespace, err = h.GetNamespace(n)
-			if err != nil {
-				log.Fatalf("Error fetching namespace: %s", err)
-			}
-			namespaces = append(namespaces, *namespace)
-
-			sharedMachines, err = h.ListSharedMachinesInNamespace(n)
-			if err != nil {
-				log.Fatalf("Error fetching shared machines: %s", err)
-			}
-		}
-
-		var allMachines []headscale.Machine
-		for _, namespace := range namespaces {
-			machines, err := h.ListMachinesInNamespace(namespace.Name)
-			if err != nil {
-				log.Fatalf("Error fetching machines: %s", err)
-			}
-			allMachines = append(allMachines, *machines...)
-		}
-
-		// listing sharedMachines is only relevant when a particular namespace is
-		// requested
-		if sharedMachines != nil {
-			allMachines = append(allMachines, *sharedMachines...)
-		}
-
-		if strings.HasPrefix(o, "json") {
-			JsonOutput(allMachines, err, o)
+			ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
 			return
 		}
 
-		if err != nil {
-			log.Fatalf("Error getting nodes: %s", err)
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		request := &v1.ListMachinesRequest{
+			Namespace: namespace,
 		}
 
-		d, err := nodesToPtables(namespace, allMachines)
+		response, err := client.ListMachines(ctx, request)
 		if err != nil {
-			log.Fatalf("Error converting to table: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Cannot get nodes: %s", err), output)
+			return
+		}
+
+		if output != "" {
+			SuccessOutput(response.Machines, "", output)
+			return
+		}
+
+		d, err := nodesToPtables(namespace, response.Machines)
+		if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Error converting to table: %s", err), output)
+			return
 		}
 
 		err = pterm.DefaultTable.WithHasHeader().WithData(d).Render()
 		if err != nil {
-			log.Fatal(err)
+			ErrorOutput(err, fmt.Sprintf("Failed to render pterm table: %s", err), output)
+			return
 		}
 	},
 }
@@ -179,24 +158,38 @@ var deleteNodeCmd = &cobra.Command{
 	Short: "Delete a node",
 	Run: func(cmd *cobra.Command, args []string) {
 		output, _ := cmd.Flags().GetString("output")
-		h, err := getHeadscaleApp()
-		if err != nil {
-			log.Fatalf("Error initializing: %s", err)
-		}
+
 		id, err := cmd.Flags().GetInt("identifier")
 		if err != nil {
-			log.Fatalf("Error converting ID to integer: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error converting ID to integer: %s", err), output)
+			return
 		}
-		m, err := h.GetMachineByID(uint64(id))
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		getRequest := &v1.GetMachineRequest{
+			MachineId: uint64(id),
+		}
+
+		getResponse, err := client.GetMachine(ctx, getRequest)
 		if err != nil {
-			log.Fatalf("Error getting node: %s", err)
+			ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
+			return
+		}
+
+		deleteRequest := &v1.DeleteMachineRequest{
+			MachineId: uint64(id),
 		}
 
 		confirm := false
 		force, _ := cmd.Flags().GetBool("force")
 		if !force {
 			prompt := &survey.Confirm{
-				Message: fmt.Sprintf("Do you want to remove the node %s?", m.Name),
+				Message: fmt.Sprintf("Do you want to remove the node %s?", getResponse.GetMachine().Name),
 			}
 			err = survey.AskOne(prompt, &confirm)
 			if err != nil {
@@ -205,71 +198,96 @@ var deleteNodeCmd = &cobra.Command{
 		}
 
 		if confirm || force {
-			err = h.DeleteMachine(m)
-			if strings.HasPrefix(output, "json") {
-				JsonOutput(map[string]string{"Result": "Node deleted"}, err, output)
+			response, err := client.DeleteMachine(ctx, deleteRequest)
+			if output != "" {
+				SuccessOutput(response, "", output)
 				return
 			}
 			if err != nil {
-				log.Fatalf("Error deleting node: %s", err)
-			}
-			fmt.Printf("Node deleted\n")
-		} else {
-			if strings.HasPrefix(output, "json") {
-				JsonOutput(map[string]string{"Result": "Node not deleted"}, err, output)
+				ErrorOutput(err, fmt.Sprintf("Error deleting node: %s", err), output)
 				return
 			}
-			fmt.Printf("Node not deleted\n")
+			SuccessOutput(map[string]string{"Result": "Node deleted"}, "Node deleted", output)
+		} else {
+			SuccessOutput(map[string]string{"Result": "Node not deleted"}, "Node not deleted", output)
 		}
 	},
 }
 
-func sharingWorker(cmd *cobra.Command, args []string) (*headscale.Headscale, string, *headscale.Machine, *headscale.Namespace) {
+func sharingWorker(
+	cmd *cobra.Command,
+	args []string,
+) (string, *v1.Machine, *v1.Namespace, error) {
+	output, _ := cmd.Flags().GetString("output")
 	namespaceStr, err := cmd.Flags().GetString("namespace")
 	if err != nil {
-		log.Fatalf("Error getting namespace: %s", err)
+		ErrorOutput(err, fmt.Sprintf("Error getting namespace: %s", err), output)
+		return "", nil, nil, err
 	}
 
-	output, _ := cmd.Flags().GetString("output")
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 
-	h, err := getHeadscaleApp()
-	if err != nil {
-		log.Fatalf("Error initializing: %s", err)
-	}
-
-	namespace, err := h.GetNamespace(namespaceStr)
-	if err != nil {
-		log.Fatalf("Error fetching namespace %s: %s", namespaceStr, err)
-	}
+	client, conn := getHeadscaleGRPCClient(ctx)
+	defer conn.Close()
 
 	id, err := cmd.Flags().GetInt("identifier")
 	if err != nil {
-		log.Fatalf("Error converting ID to integer: %s", err)
-	}
-	machine, err := h.GetMachineByID(uint64(id))
-	if err != nil {
-		log.Fatalf("Error getting node: %s", err)
+		ErrorOutput(err, fmt.Sprintf("Error converting ID to integer: %s", err), output)
+		return "", nil, nil, err
 	}
 
-	return h, output, machine, namespace
+	machineRequest := &v1.GetMachineRequest{
+		MachineId: uint64(id),
+	}
+
+	machineResponse, err := client.GetMachine(ctx, machineRequest)
+	if err != nil {
+		ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
+		return "", nil, nil, err
+	}
+
+	namespaceRequest := &v1.GetNamespaceRequest{
+		Name: namespaceStr,
+	}
+
+	namespaceResponse, err := client.GetNamespace(ctx, namespaceRequest)
+	if err != nil {
+		ErrorOutput(err, fmt.Sprintf("Error getting node node: %s", err), output)
+		return "", nil, nil, err
+	}
+
+	return output, machineResponse.GetMachine(), namespaceResponse.GetNamespace(), nil
 }
 
 var shareMachineCmd = &cobra.Command{
 	Use:   "share",
 	Short: "Shares a node from the current namespace to the specified one",
 	Run: func(cmd *cobra.Command, args []string) {
-		h, output, machine, namespace := sharingWorker(cmd, args)
-		err := h.AddSharedMachineToNamespace(machine, namespace)
-		if strings.HasPrefix(output, "json") {
-			JsonOutput(map[string]string{"Result": "Node shared"}, err, output)
-			return
-		}
+		output, machine, namespace, err := sharingWorker(cmd, args)
 		if err != nil {
-			fmt.Printf("Error sharing node: %s\n", err)
+			ErrorOutput(err, fmt.Sprintf("Failed to fetch namespace or machine: %s", err), output)
 			return
 		}
 
-		fmt.Println("Node shared!")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		request := &v1.ShareMachineRequest{
+			MachineId: machine.Id,
+			Namespace: namespace.Name,
+		}
+
+		response, err := client.ShareMachine(ctx, request)
+		if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Error sharing node: %s", err), output)
+			return
+		}
+
+		SuccessOutput(response.Machine, "Node shared", output)
 	},
 }
 
@@ -277,33 +295,45 @@ var unshareMachineCmd = &cobra.Command{
 	Use:   "unshare",
 	Short: "Unshares a node from the specified namespace",
 	Run: func(cmd *cobra.Command, args []string) {
-		h, output, machine, namespace := sharingWorker(cmd, args)
-		err := h.RemoveSharedMachineFromNamespace(machine, namespace)
-		if strings.HasPrefix(output, "json") {
-			JsonOutput(map[string]string{"Result": "Node unshared"}, err, output)
-			return
-		}
+		output, machine, namespace, err := sharingWorker(cmd, args)
 		if err != nil {
-			fmt.Printf("Error unsharing node: %s\n", err)
+			ErrorOutput(err, fmt.Sprintf("Failed to fetch namespace or machine: %s", err), output)
 			return
 		}
 
-		fmt.Println("Node unshared!")
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		client, conn := getHeadscaleGRPCClient(ctx)
+		defer conn.Close()
+
+		request := &v1.UnshareMachineRequest{
+			MachineId: machine.Id,
+			Namespace: namespace.Name,
+		}
+
+		response, err := client.UnshareMachine(ctx, request)
+		if err != nil {
+			ErrorOutput(err, fmt.Sprintf("Error unsharing node: %s", err), output)
+			return
+		}
+
+		SuccessOutput(response.Machine, "Node shared", output)
 	},
 }
 
-func nodesToPtables(currentNamespace *headscale.Namespace, machines []headscale.Machine) (pterm.TableData, error) {
+func nodesToPtables(currentNamespace string, machines []*v1.Machine) (pterm.TableData, error) {
 	d := pterm.TableData{{"ID", "Name", "NodeKey", "Namespace", "IP address", "Ephemeral", "Last seen", "Online"}}
 
 	for _, machine := range machines {
 		var ephemeral bool
-		if machine.AuthKey != nil && machine.AuthKey.Ephemeral {
+		if machine.PreAuthKey != nil && machine.PreAuthKey.Ephemeral {
 			ephemeral = true
 		}
 		var lastSeen time.Time
 		var lastSeenTime string
 		if machine.LastSeen != nil {
-			lastSeen = *machine.LastSeen
+			lastSeen = machine.LastSeen.AsTime()
 			lastSeenTime = lastSeen.Format("2006-01-02 15:04:05")
 		}
 		nKey, err := wgkey.ParseHex(machine.NodeKey)
@@ -320,13 +350,25 @@ func nodesToPtables(currentNamespace *headscale.Namespace, machines []headscale.
 		}
 
 		var namespace string
-		if (currentNamespace == nil) || (currentNamespace.ID == machine.NamespaceID) {
+		if currentNamespace == machine.Namespace.Name {
 			namespace = pterm.LightMagenta(machine.Namespace.Name)
 		} else {
 			// Shared into this namespace
 			namespace = pterm.LightYellow(machine.Namespace.Name)
 		}
-		d = append(d, []string{strconv.FormatUint(machine.ID, 10), machine.Name, nodeKey.ShortString(), namespace, machine.IPAddress, strconv.FormatBool(ephemeral), lastSeenTime, online})
+		d = append(
+			d,
+			[]string{
+				strconv.FormatUint(machine.Id, 10),
+				machine.Name,
+				nodeKey.ShortString(),
+				namespace,
+				machine.IpAddress,
+				strconv.FormatBool(ephemeral),
+				lastSeenTime,
+				online,
+			},
+		)
 	}
 	return d, nil
 }