From 7b0d2dfb4ac22da6c31363eaf038bf4a302899e2 Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Thu, 4 Nov 2021 22:44:35 +0000 Subject: [PATCH] Convert nodes command to use gRPC --- cmd/headscale/cli/nodes.go | 300 +++++++++++++++++++++---------------- 1 file changed, 171 insertions(+), 129 deletions(-) diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go index cdf37efb..70fae835 100644 --- a/cmd/headscale/cli/nodes.go +++ b/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 }