Skip to content

Commit

Permalink
feat(cli): add tag subcommand to add and remove tags
Browse files Browse the repository at this point in the history
  • Loading branch information
restanrm committed Apr 15, 2022
1 parent 02f68eb commit 9de9bc2
Showing 1 changed file with 208 additions and 0 deletions.
208 changes: 208 additions & 0 deletions cmd/headscale/cli/tags.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 9de9bc2

Please sign in to comment.