1
0
mirror of https://github.com/juanfont/headscale.git synced 2025-09-02 13:47:00 +02:00

This commit adds the "findSingleUser" helper function that centralizes and standardizes user lookup operations across CLI commands. The function:

Takes username and ID flags from commands
Handles error checking and validation
Ensures exactly one user is found
Provides consistent error messages
The function eliminates duplicate code across multiple command files (users, nodes, preauthkeys).
This commit is contained in:
Tevin Flores 2025-05-29 15:18:19 -04:00
parent b8044c29dd
commit a32b156d25
4 changed files with 163 additions and 162 deletions

View File

@ -21,84 +21,58 @@ import (
func init() {
rootCmd.AddCommand(nodeCmd)
listNodesCmd.Flags().StringP("user", "u", "", "Filter by user")
listNodesCmd.Flags().BoolP("tags", "t", false, "Show tags")
listNodesCmd.Flags().StringP("namespace", "n", "", "User")
listNodesNamespaceFlag := listNodesCmd.Flags().Lookup("namespace")
listNodesNamespaceFlag.Deprecated = deprecateNamespaceMessage
listNodesNamespaceFlag.Hidden = true
usernameAndIDFlag(listNodesCmd)
nodeCmd.AddCommand(listNodesCmd)
listNodeRoutesCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
listNodeRoutesCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
nodeCmd.AddCommand(listNodeRoutesCmd)
registerNodeCmd.Flags().StringP("user", "u", "", "User")
registerNodeCmd.Flags().StringP("namespace", "n", "", "User")
registerNodeNamespaceFlag := registerNodeCmd.Flags().Lookup("namespace")
registerNodeNamespaceFlag.Deprecated = deprecateNamespaceMessage
registerNodeNamespaceFlag.Hidden = true
err := registerNodeCmd.MarkFlagRequired("user")
if err != nil {
log.Fatal(err.Error())
}
usernameAndIDFlag(registerNodeCmd)
registerNodeCmd.Flags().StringP("key", "k", "", "Key")
err = registerNodeCmd.MarkFlagRequired("key")
err := registerNodeCmd.MarkFlagRequired("key")
if err != nil {
log.Fatal(err.Error())
}
nodeCmd.AddCommand(registerNodeCmd)
expireNodeCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
err = expireNodeCmd.MarkFlagRequired("identifier")
expireNodeCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
err = expireNodeCmd.MarkFlagRequired("node-id")
if err != nil {
log.Fatal(err.Error())
}
nodeCmd.AddCommand(expireNodeCmd)
renameNodeCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
err = renameNodeCmd.MarkFlagRequired("identifier")
renameNodeCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
err = renameNodeCmd.MarkFlagRequired("node-id")
if err != nil {
log.Fatal(err.Error())
}
nodeCmd.AddCommand(renameNodeCmd)
deleteNodeCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
err = deleteNodeCmd.MarkFlagRequired("identifier")
deleteNodeCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
err = deleteNodeCmd.MarkFlagRequired("node-id")
if err != nil {
log.Fatal(err.Error())
}
nodeCmd.AddCommand(deleteNodeCmd)
moveNodeCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
moveNodeCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
err = moveNodeCmd.MarkFlagRequired("identifier")
err = moveNodeCmd.MarkFlagRequired("node-id")
if err != nil {
log.Fatal(err.Error())
}
moveNodeCmd.Flags().Uint64P("user", "u", 0, "New user")
usernameAndIDFlag(moveNodeCmd, "Target user ID to move the node to", "Target username to move the node to")
moveNodeCmd.Flags().StringP("namespace", "n", "", "User")
moveNodeNamespaceFlag := moveNodeCmd.Flags().Lookup("namespace")
moveNodeNamespaceFlag.Deprecated = deprecateNamespaceMessage
moveNodeNamespaceFlag.Hidden = true
err = moveNodeCmd.MarkFlagRequired("user")
if err != nil {
log.Fatal(err.Error())
}
nodeCmd.AddCommand(moveNodeCmd)
tagCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
tagCmd.MarkFlagRequired("identifier")
tagCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
tagCmd.MarkFlagRequired("node-id")
tagCmd.Flags().StringSliceP("tags", "t", []string{}, "List of tags to add to the node")
nodeCmd.AddCommand(tagCmd)
approveRoutesCmd.Flags().Uint64P("identifier", "i", 0, "Node identifier (ID)")
approveRoutesCmd.MarkFlagRequired("identifier")
approveRoutesCmd.Flags().Uint64P("node-id", "N", 0, "Node identifier (ID)")
approveRoutesCmd.MarkFlagRequired("node-id")
approveRoutesCmd.Flags().StringSliceP("routes", "r", []string{}, `List of routes that will be approved (comma-separated, e.g. "10.0.0.0/8,192.168.0.0/24" or empty string to remove all approved routes)`)
nodeCmd.AddCommand(approveRoutesCmd)
@ -116,15 +90,17 @@ var registerNodeCmd = &cobra.Command{
Short: "Registers a node to your network",
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetString("user")
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
}
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
user, err := findSingleUser(ctx, client, cmd, "register Node", output)
if err != nil {
// The helper already calls ErrorOutput, so we can just return
return
}
registrationID, err := cmd.Flags().GetString("key")
if err != nil {
ErrorOutput(
@ -136,7 +112,7 @@ var registerNodeCmd = &cobra.Command{
request := &v1.RegisterNodeRequest{
Key: registrationID,
User: user,
User: user.GetName(),
}
response, err := client.RegisterNode(ctx, request)
@ -163,10 +139,6 @@ var listNodesCmd = &cobra.Command{
Aliases: []string{"ls", "show"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetString("user")
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
}
showTags, err := cmd.Flags().GetBool("tags")
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting tags flag: %s", err), output)
@ -176,8 +148,20 @@ var listNodesCmd = &cobra.Command{
defer cancel()
defer conn.Close()
request := &v1.ListNodesRequest{
User: user,
// Check if user identifier flags are provided
id, _ := cmd.Flags().GetInt64("identifier")
username, _ := cmd.Flags().GetString("name")
request := &v1.ListNodesRequest{}
// Only filter by user if user flags are provided
if id >= 0 || username != "" {
user, err := findSingleUser(ctx, client, cmd, "list", output)
if err != nil {
// The helper already calls ErrorOutput, so we can just return
return
}
request.User = user.GetName()
}
response, err := client.ListNodes(ctx, request)
@ -193,7 +177,7 @@ var listNodesCmd = &cobra.Command{
SuccessOutput(response.GetNodes(), "", output)
}
tableData, err := nodesToPtables(user, showTags, response.GetNodes())
tableData, err := nodesToPtables(request.User, showTags, response.GetNodes())
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error converting to table: %s", err), output)
}
@ -464,7 +448,7 @@ var moveNodeCmd = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
identifier, err := cmd.Flags().GetUint64("identifier")
identifier, err := cmd.Flags().GetUint64("node-id")
if err != nil {
ErrorOutput(
err,
@ -475,21 +459,16 @@ var moveNodeCmd = &cobra.Command{
return
}
user, err := cmd.Flags().GetUint64("user")
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error getting user: %s", err),
output,
)
return
}
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
user, err := findSingleUser(ctx, client, cmd, "move Node", output)
if err != nil {
// The helper already calls ErrorOutput, so we can just return
return
}
getRequest := &v1.GetNodeRequest{
NodeId: identifier,
}
@ -510,7 +489,7 @@ var moveNodeCmd = &cobra.Command{
moveRequest := &v1.MoveNodeRequest{
NodeId: identifier,
User: user,
User: user.GetId(),
}
moveResponse, err := client.MoveNode(ctx, moveRequest)

View File

@ -20,20 +20,20 @@ const (
func init() {
rootCmd.AddCommand(preauthkeysCmd)
preauthkeysCmd.PersistentFlags().Uint64P("user", "u", 0, "User identifier (ID)")
preauthkeysCmd.PersistentFlags().StringP("namespace", "n", "", "User")
preauthkeysCmd.PersistentFlags().String("namespace", "", "User")
pakNamespaceFlag := preauthkeysCmd.PersistentFlags().Lookup("namespace")
pakNamespaceFlag.Deprecated = deprecateNamespaceMessage
pakNamespaceFlag.Hidden = true
err := preauthkeysCmd.MarkPersistentFlagRequired("user")
if err != nil {
log.Fatal().Err(err).Msg("")
}
preauthkeysCmd.AddCommand(listPreAuthKeys)
preauthkeysCmd.AddCommand(createPreAuthKeyCmd)
preauthkeysCmd.AddCommand(expirePreAuthKeyCmd)
usernameAndIDFlag(listPreAuthKeys)
usernameAndIDFlag(createPreAuthKeyCmd)
usernameAndIDFlag(expirePreAuthKeyCmd)
createPreAuthKeyCmd.PersistentFlags().
Bool("reusable", false, "Make the preauthkey reusable")
createPreAuthKeyCmd.PersistentFlags().
@ -57,17 +57,17 @@ var listPreAuthKeys = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetUint64("user")
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
}
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
user, err := findSingleUser(ctx, client, cmd, "list", output)
if err != nil {
return
}
request := &v1.ListPreAuthKeysRequest{
User: user,
User: user.GetId(),
}
response, err := client.ListPreAuthKeys(ctx, request)
@ -141,9 +141,13 @@ var createPreAuthKeyCmd = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetUint64("user")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
user, err := findSingleUser(ctx, client, cmd, "list", output)
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
return
}
reusable, _ := cmd.Flags().GetBool("reusable")
@ -151,7 +155,7 @@ var createPreAuthKeyCmd = &cobra.Command{
tags, _ := cmd.Flags().GetStringSlice("tags")
request := &v1.CreatePreAuthKeyRequest{
User: user,
User: user.GetId(),
Reusable: reusable,
Ephemeral: ephemeral,
AclTags: tags,
@ -176,10 +180,6 @@ var createPreAuthKeyCmd = &cobra.Command{
request.Expiration = timestamppb.New(expiration)
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
response, err := client.CreatePreAuthKey(ctx, request)
if err != nil {
ErrorOutput(
@ -206,17 +206,18 @@ var expirePreAuthKeyCmd = &cobra.Command{
},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
user, err := cmd.Flags().GetUint64("user")
if err != nil {
ErrorOutput(err, fmt.Sprintf("Error getting user: %s", err), output)
}
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
user, err := findSingleUser(ctx, client, cmd, "list", output)
if err != nil {
return
}
request := &v1.ExpirePreAuthKeyRequest{
User: user,
User: user.GetId(),
Key: args[0],
}

View File

@ -0,0 +1,84 @@
package cli
import (
"context"
"errors"
"fmt"
v1 "github.com/juanfont/headscale/gen/go/headscale/v1"
"github.com/spf13/cobra"
"google.golang.org/grpc/status"
)
// usernameAndIDFlag adds the common user identification flags to a command
func usernameAndIDFlag(cmd *cobra.Command, opts ...string) {
idHelp := "User identifier (ID)"
nameHelp := "Username"
if len(opts) > 0 && opts[0] != "" {
idHelp = opts[0]
}
if len(opts) > 1 && opts[1] != "" {
nameHelp = opts[1]
}
cmd.PersistentFlags().Int64P("identifier", "i", -1, idHelp)
cmd.PersistentFlags().StringP("name", "n", "", nameHelp)
}
// usernameAndIDFromFlag returns the username and ID from the flags of the command.
// If both are empty, it will exit the program with an error.
func usernameAndIDFromFlag(cmd *cobra.Command) (uint64, string) {
username, _ := cmd.Flags().GetString("name")
identifier, _ := cmd.Flags().GetInt64("identifier")
if username == "" && identifier < 0 {
err := errors.New("--name or --identifier flag is required")
ErrorOutput(
err,
fmt.Sprintf(
"User identification error: %s",
status.Convert(err).Message(),
),
"",
)
}
return uint64(identifier), username
}
// findSingleUser takes command flags and returns a single user
// It handles all error checking and ensures exactly one user is found
func findSingleUser(
ctx context.Context,
client v1.HeadscaleServiceClient,
cmd *cobra.Command,
operationName string,
output string,
) (*v1.User, error) {
id, username := usernameAndIDFromFlag(cmd)
listReq := &v1.ListUsersRequest{
Name: username,
Id: id,
}
users, err := client.ListUsers(ctx, listReq)
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
return nil, err
}
if len(users.GetUsers()) != 1 {
err := fmt.Errorf("Unable to determine user to %s, query returned multiple users, use ID", operationName)
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
return nil, err
}
return users.GetUsers()[0], nil
}

View File

@ -13,31 +13,6 @@ import (
"google.golang.org/grpc/status"
)
func usernameAndIDFlag(cmd *cobra.Command) {
cmd.Flags().Int64P("identifier", "i", -1, "User identifier (ID)")
cmd.Flags().StringP("name", "n", "", "Username")
}
// usernameAndIDFromFlag returns the username and ID from the flags of the command.
// If both are empty, it will exit the program with an error.
func usernameAndIDFromFlag(cmd *cobra.Command) (uint64, string) {
username, _ := cmd.Flags().GetString("name")
identifier, _ := cmd.Flags().GetInt64("identifier")
if username == "" && identifier < 0 {
err := errors.New("--name or --identifier flag is required")
ErrorOutput(
err,
fmt.Sprintf(
"Cannot rename user: %s",
status.Convert(err).Message(),
),
"",
)
}
return uint64(identifier), username
}
func init() {
rootCmd.AddCommand(userCmd)
userCmd.AddCommand(createUserCmd)
@ -133,36 +108,16 @@ var destroyUserCmd = &cobra.Command{
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
id, username := usernameAndIDFromFlag(cmd)
request := &v1.ListUsersRequest{
Name: username,
Id: id,
}
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
users, err := client.ListUsers(ctx, request)
user, err := findSingleUser(ctx, client, cmd, "rename", output)
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
// The helper already calls ErrorOutput, so we can just return
return
}
if len(users.GetUsers()) != 1 {
err := fmt.Errorf("Unable to determine user to delete, query returned multiple users, use ID")
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
}
user := users.GetUsers()[0]
confirm := false
force, _ := cmd.Flags().GetBool("force")
if !force {
@ -277,34 +232,16 @@ var renameUserCmd = &cobra.Command{
defer cancel()
defer conn.Close()
id, username := usernameAndIDFromFlag(cmd)
listReq := &v1.ListUsersRequest{
Name: username,
Id: id,
}
users, err := client.ListUsers(ctx, listReq)
user, err := findSingleUser(ctx, client, cmd, "rename", output)
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
}
if len(users.GetUsers()) != 1 {
err := fmt.Errorf("Unable to determine user to delete, query returned multiple users, use ID")
ErrorOutput(
err,
fmt.Sprintf("Error: %s", status.Convert(err).Message()),
output,
)
// The helper already calls ErrorOutput, so we can just return
return
}
newName, _ := cmd.Flags().GetString("new-name")
renameReq := &v1.RenameUserRequest{
OldId: id,
OldId: user.GetId(),
NewName: newName,
}