diff --git a/cmd/headscale/cli/tags.go b/cmd/headscale/cli/tags.go new file mode 100644 index 00000000..c40b6397 --- /dev/null +++ b/cmd/headscale/cli/tags.go @@ -0,0 +1,208 @@ +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", 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", 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 +}