From f53bb63b2d93c3fae9433bd408decb1f4c657fb5 Mon Sep 17 00:00:00 2001 From: Adrien Raffin-Caboisse Date: Thu, 21 Apr 2022 23:43:20 +0200 Subject: [PATCH] fix: move tag command to subcommand of nodes --- cmd/headscale/cli/nodes.go | 186 +++++++++++++++++++++++++++++++++ cmd/headscale/cli/tags.go | 208 ------------------------------------- cmd/headscale/cli/utils.go | 10 ++ 3 files changed, 196 insertions(+), 208 deletions(-) delete mode 100644 cmd/headscale/cli/tags.go diff --git a/cmd/headscale/cli/nodes.go b/cmd/headscale/cli/nodes.go index 5b5e1e23..143e5a4c 100644 --- a/cmd/headscale/cli/nodes.go +++ b/cmd/headscale/cli/nodes.go @@ -46,6 +46,25 @@ func init() { log.Fatalf(err.Error()) } nodeCmd.AddCommand(deleteNodeCmd) + + nodeCmd.AddCommand(tagCmd) + addTagCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)") + err = addTagCmd.MarkFlagRequired("identifier") + if err != nil { + log.Fatalf(err.Error()) + } + addTagCmd.Flags(). + StringSliceP("tags", "t", []string{}, "List of tags to add to the node") + tagCmd.AddCommand(addTagCmd) + + delTagCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)") + err = delTagCmd.MarkFlagRequired("identifier") + if err != nil { + log.Fatalf(err.Error()) + } + delTagCmd.Flags(). + StringSliceP("tags", "t", []string{}, "List of tags to remove from the node") + tagCmd.AddCommand(delTagCmd) } var nodeCmd = &cobra.Command{ @@ -399,3 +418,170 @@ func nodesToPtables( return tableData, nil } + +var tagCmd = &cobra.Command{ + Use: "tags", + Short: "Manage the tags of Headscale", + Aliases: []string{"t", "tag"}, +} + +var addTagCmd = &cobra.Command{ + Use: "add", + Short: "Add tags to a node in your network", + Run: func(cmd *cobra.Command, args []string) { + output, _ := cmd.Flags().GetString("output") + ctx, client, conn, cancel := getHeadscaleCLIClient() + defer cancel() + defer conn.Close() + + // retrieve flags from CLI + identifier, err := cmd.Flags().GetUint64("identifier") + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error converting ID to integer: %s", err), + output, + ) + + return + } + tagsToAdd, err := cmd.Flags().GetStringSlice("tags") + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error retrieving list of tags to add to machine, %v", err), + output, + ) + + return + } + + // retrieve machine informations + request := &v1.GetMachineRequest{ + MachineId: identifier, + } + resp, err := client.GetMachine(ctx, request) + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error retrieving machine: %s", err), + output, + ) + } + + // update machine + mergedTags := resp.Machine.GetForcedTags() + for _, tag := range tagsToAdd { + if !containsString(mergedTags, tag) { + mergedTags = append(mergedTags, tag) + } + } + + machine := resp.GetMachine() + machine.ForcedTags = mergedTags + + updateReq := &v1.UpdateMachineRequest{ + Machine: machine, + } + + // send updated machine upstream + updateResponse, err := client.UpdateMachine(ctx, updateReq) + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error while updating machine: %s", err), + output, + ) + } + + if updateResponse != nil { + SuccessOutput( + updateResponse.GetMachine(), + "Machine updated", + output, + ) + } + }, +} + +var delTagCmd = &cobra.Command{ + Use: "del", + Short: "remove tags to a node in your network", + Aliases: []string{"remove", "rm"}, + Run: func(cmd *cobra.Command, args []string) { + output, _ := cmd.Flags().GetString("output") + ctx, client, conn, cancel := getHeadscaleCLIClient() + defer cancel() + defer conn.Close() + + // retrieve flags from CLI + identifier, err := cmd.Flags().GetUint64("identifier") + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error converting ID to integer: %s", err), + output, + ) + + return + } + tagsToRemove, err := cmd.Flags().GetStringSlice("tags") + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error retrieving list of tags to add to machine: %v", err), + output, + ) + + return + } + + // retrieve machine informations + request := &v1.GetMachineRequest{ + MachineId: identifier, + } + resp, err := client.GetMachine(ctx, request) + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error retrieving machine: %s", err), + output, + ) + } + + // update machine + keepTags := resp.Machine.GetForcedTags() + for _, tag := range tagsToRemove { + for i, t := range keepTags { + if t == tag { + keepTags = append(keepTags[:i], keepTags[i+1:]...) + } + } + } + + machine := resp.GetMachine() + machine.ForcedTags = keepTags + + updateReq := &v1.UpdateMachineRequest{ + Machine: machine, + } + + // send updated machine upstream + updateResponse, err := client.UpdateMachine(ctx, updateReq) + if err != nil { + ErrorOutput( + err, + fmt.Sprintf("Error while updating machine: %s", err), + output, + ) + } + + if updateResponse != nil { + SuccessOutput( + updateResponse.GetMachine(), + "Machine updated", + output, + ) + } + }, +} diff --git a/cmd/headscale/cli/tags.go b/cmd/headscale/cli/tags.go deleted file mode 100644 index ec42c9a9..00000000 --- a/cmd/headscale/cli/tags.go +++ /dev/null @@ -1,208 +0,0 @@ -package cli - -import ( - "fmt" - "log" - - v1 "github.com/juanfont/headscale/gen/go/headscale/v1" - "github.com/spf13/cobra" -) - -func init() { - rootCmd.AddCommand(tagCmd) - - addTagCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)") - err := addTagCmd.MarkFlagRequired("identifier") - if err != nil { - log.Fatalf(err.Error()) - } - addTagCmd.Flags(). - StringSliceP("tags", "t", []string{}, "List of tags to add to the node") - tagCmd.AddCommand(addTagCmd) - - delTagCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)") - err = delTagCmd.MarkFlagRequired("identifier") - if err != nil { - log.Fatalf(err.Error()) - } - delTagCmd.Flags(). - StringSliceP("tags", "t", []string{}, "List of tags to remove from the node") - tagCmd.AddCommand(delTagCmd) -} - -var tagCmd = &cobra.Command{ - Use: "tags", - Short: "Manage the tags of Headscale", - Aliases: []string{"t", "tag"}, -} - -var addTagCmd = &cobra.Command{ - Use: "add", - Short: "Add tags to a node in your network", - Run: func(cmd *cobra.Command, args []string) { - output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := getHeadscaleCLIClient() - defer cancel() - defer conn.Close() - - // retrieve flags from CLI - identifier, err := cmd.Flags().GetUint64("identifier") - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error converting ID to integer: %s", err), - output, - ) - - return - } - tagsToAdd, err := cmd.Flags().GetStringSlice("tags") - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error retrieving list of tags to add to machine, %v", err), - output, - ) - - return - } - - // retrieve machine informations - request := &v1.GetMachineRequest{ - MachineId: identifier, - } - resp, err := client.GetMachine(ctx, request) - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error retrieving machine: %s", err), - output, - ) - } - - // update machine - mergedTags := resp.Machine.GetForcedTags() - for _, tag := range tagsToAdd { - if !containsString(mergedTags, tag) { - mergedTags = append(mergedTags, tag) - } - } - - machine := resp.GetMachine() - machine.ForcedTags = mergedTags - - updateReq := &v1.UpdateMachineRequest{ - Machine: machine, - } - - // send updated machine upstream - updateResponse, err := client.UpdateMachine(ctx, updateReq) - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error while updating machine: %s", err), - output, - ) - } - - if updateResponse != nil { - SuccessOutput( - updateResponse.GetMachine(), - "Machine updated", - output, - ) - } - }, -} - -var delTagCmd = &cobra.Command{ - Use: "del", - Short: "remove tags to a node in your network", - Aliases: []string{"remove", "rm"}, - Run: func(cmd *cobra.Command, args []string) { - output, _ := cmd.Flags().GetString("output") - ctx, client, conn, cancel := getHeadscaleCLIClient() - defer cancel() - defer conn.Close() - - // retrieve flags from CLI - identifier, err := cmd.Flags().GetUint64("identifier") - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error converting ID to integer: %s", err), - output, - ) - - return - } - tagsToRemove, err := cmd.Flags().GetStringSlice("tags") - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error retrieving list of tags to add to machine: %v", err), - output, - ) - - return - } - - // retrieve machine informations - request := &v1.GetMachineRequest{ - MachineId: identifier, - } - resp, err := client.GetMachine(ctx, request) - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error retrieving machine: %s", err), - output, - ) - } - - // update machine - keepTags := resp.Machine.GetForcedTags() - for _, tag := range tagsToRemove { - for i, t := range keepTags { - if t == tag { - keepTags = append(keepTags[:i], keepTags[i+1:]...) - } - } - } - - machine := resp.GetMachine() - machine.ForcedTags = keepTags - - updateReq := &v1.UpdateMachineRequest{ - Machine: machine, - } - - // send updated machine upstream - updateResponse, err := client.UpdateMachine(ctx, updateReq) - if err != nil { - ErrorOutput( - err, - fmt.Sprintf("Error while updating machine: %s", err), - output, - ) - } - - if updateResponse != nil { - SuccessOutput( - updateResponse.GetMachine(), - "Machine updated", - output, - ) - } - }, -} - -func containsString(ss []string, s string) bool { - for _, v := range ss { - if v == s { - return true - } - } - - return false -} diff --git a/cmd/headscale/cli/utils.go b/cmd/headscale/cli/utils.go index 5c724591..de369baa 100644 --- a/cmd/headscale/cli/utils.go +++ b/cmd/headscale/cli/utils.go @@ -564,3 +564,13 @@ func GetFileMode(key string) fs.FileMode { return fs.FileMode(mode) } + +func containsString(ss []string, s string) bool { + for _, v := range ss { + if v == s { + return true + } + } + + return false +}