juanfont.headscale/cmd/headscale/cli/tags.go

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
}