-
Notifications
You must be signed in to change notification settings - Fork 79
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: allow setting default SSH keys for contexts
- Loading branch information
Showing
6 changed files
with
249 additions
and
8 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,212 @@ | ||
| package context | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "os" | ||
| "slices" | ||
| "strings" | ||
| "sync" | ||
|
|
||
| "github.com/spf13/cobra" | ||
|
|
||
| "github.com/hetznercloud/cli/internal/cmd/cmpl" | ||
| "github.com/hetznercloud/cli/internal/hcapi2" | ||
| "github.com/hetznercloud/cli/internal/state" | ||
| "github.com/hetznercloud/cli/internal/state/config" | ||
| "github.com/hetznercloud/hcloud-go/v2/hcloud" | ||
| ) | ||
|
|
||
| // TODO: make *state.Config mockable and write tests | ||
|
|
||
| func newSSHKeyCommand(s state.State) *cobra.Command { | ||
| cmd := &cobra.Command{ | ||
| Use: "ssh-key", | ||
| Short: "Manage a context's default SSH key", | ||
| Args: cobra.NoArgs, | ||
| TraverseChildren: true, | ||
| DisableFlagsInUseLine: true, | ||
| } | ||
| cmd.AddCommand(newSSHKeyAddCommand(s)) | ||
| cmd.AddCommand(newSSHKeyRemoveCommand(s)) | ||
| cmd.AddCommand(newSSHKeyListCommand(s)) | ||
| return cmd | ||
| } | ||
|
|
||
| func newSSHKeyAddCommand(s state.State) *cobra.Command { | ||
| cmd := &cobra.Command{ | ||
| Use: "add SSH-KEY...", | ||
| Short: "Add a default SSH key to the context", | ||
| TraverseChildren: true, | ||
| DisableFlagsInUseLine: true, | ||
| RunE: state.Wrap(s, runSSHKeyAdd), | ||
| } | ||
| cmd.Flags().String("context", "", "Name of the context to add the default SSH key(s) to") | ||
| _ = cmd.RegisterFlagCompletionFunc("context", cmpl.SuggestCandidates(s.Config().ContextNames()...)) | ||
|
|
||
| cmd.Flags().Bool("all", false, "Add all available SSH keys to the context") | ||
| _ = cmd.RegisterFlagCompletionFunc("all", cmpl.SuggestCandidates("true", "false")) | ||
| return cmd | ||
| } | ||
|
|
||
| func runSSHKeyAdd(s state.State, cmd *cobra.Command, args []string) error { | ||
|
|
||
| ctx, err := getContext(s, cmd) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| client := hcapi2.NewClient(hcloud.WithToken(ctx.Token)) | ||
| keys := args | ||
|
|
||
| all, _ := cmd.Flags().GetBool("all") | ||
| if all { | ||
| allKeys, err := client.SSHKey().All(s) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| if len(allKeys) == 0 { | ||
| return fmt.Errorf("no SSH keys available") | ||
| } | ||
| keys = make([]string, len(allKeys)) | ||
| for i, key := range allKeys { | ||
| keys[i] = fmt.Sprintf("%d", key.ID) | ||
| } | ||
| } else { | ||
| if len(keys) == 0 { | ||
| return fmt.Errorf("no SSH keys specified") | ||
| } | ||
| var ( | ||
| notExist []string | ||
| wg = &sync.WaitGroup{} | ||
| mu = &sync.Mutex{} | ||
| ) | ||
| wg.Add(len(keys)) | ||
| for _, key := range keys { | ||
| key := key | ||
| go func() { | ||
| k, _, _ := client.SSHKey().Get(s, key) | ||
| if k == nil { | ||
| mu.Lock() | ||
| notExist = append(notExist, key) | ||
| mu.Unlock() | ||
| } | ||
| wg.Done() | ||
| }() | ||
| } | ||
| wg.Wait() | ||
| _, _ = fmt.Fprintf(os.Stderr, "Warning: The given SSH key(s) %s do not exist in context \"%s\"\n", strings.Join(notExist, ", "), ctx.Name) | ||
| } | ||
|
|
||
| ctx.SSHKeys = append(ctx.SSHKeys, keys...) | ||
| // remove duplicates | ||
| slices.Sort(ctx.SSHKeys) | ||
| ctx.SSHKeys = slices.Compact(ctx.SSHKeys) | ||
|
|
||
| if err := s.Config().Write(); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| cmd.Printf("Added SSH key(s) %s to context \"%s\"\n", strings.Join(keys, ", "), ctx.Name) | ||
| return nil | ||
| } | ||
|
|
||
| func newSSHKeyRemoveCommand(s state.State) *cobra.Command { | ||
| cmd := &cobra.Command{ | ||
| Use: "remove SSH-KEY...", | ||
| Short: "Remove a default SSH key from the context", | ||
| TraverseChildren: true, | ||
| DisableFlagsInUseLine: true, | ||
| RunE: state.Wrap(s, runSSHKeyRemove), | ||
| } | ||
| cmd.Flags().String("context", "", "Name of the context to remove the default SSH key(s) from") | ||
| _ = cmd.RegisterFlagCompletionFunc("context", cmpl.SuggestCandidates(s.Config().ContextNames()...)) | ||
|
|
||
| cmd.Flags().Bool("all", false, "Remove all SSH keys from the context") | ||
| _ = cmd.RegisterFlagCompletionFunc("all", cmpl.SuggestCandidates("true", "false")) | ||
| return cmd | ||
| } | ||
|
|
||
| func runSSHKeyRemove(s state.State, cmd *cobra.Command, args []string) error { | ||
|
|
||
| ctx, err := getContext(s, cmd) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| keys := args | ||
| origLen := len(ctx.SSHKeys) | ||
|
|
||
| all, _ := cmd.Flags().GetBool("all") | ||
| if all { | ||
| ctx.SSHKeys = nil | ||
| } else { | ||
| var newKeys []string | ||
| for _, key := range ctx.SSHKeys { | ||
| if slices.Contains(keys, key) { | ||
| continue | ||
| } | ||
| newKeys = append(newKeys, key) | ||
| } | ||
| ctx.SSHKeys = newKeys | ||
| } | ||
|
|
||
| if err := s.Config().Write(); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| removed := origLen - len(ctx.SSHKeys) | ||
| cmd.Printf("Removed %d SSH key(s) from context \"%s\"\n", removed, ctx.Name) | ||
| return nil | ||
| } | ||
|
|
||
| func newSSHKeyListCommand(s state.State) *cobra.Command { | ||
| cmd := &cobra.Command{ | ||
| Use: "list", | ||
| Short: "List all default SSH keys of the context", | ||
| TraverseChildren: true, | ||
| DisableFlagsInUseLine: true, | ||
| RunE: state.Wrap(s, runSSHKeyList), | ||
| } | ||
| cmd.Flags().String("context", "", "Name of the context to list the default SSH keys from") | ||
| _ = cmd.RegisterFlagCompletionFunc("context", cmpl.SuggestCandidates(s.Config().ContextNames()...)) | ||
| return cmd | ||
| } | ||
|
|
||
| func runSSHKeyList(s state.State, cmd *cobra.Command, _ []string) error { | ||
|
|
||
| ctx, err := getContext(s, cmd) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| if len(ctx.SSHKeys) == 0 { | ||
| cmd.Printf("No SSH keys in context \"%s\"\n", ctx.Name) | ||
| return nil | ||
| } | ||
|
|
||
| cmd.Printf("SSH keys in context \"%s\":\n", ctx.Name) | ||
| for _, key := range ctx.SSHKeys { | ||
| cmd.Printf(" - %s\n", key) | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func getContext(s state.State, cmd *cobra.Command) (*config.Context, error) { | ||
|
|
||
| var ctx *config.Context | ||
|
|
||
| ctxName, _ := cmd.Flags().GetString("context") | ||
| if ctxName != "" { | ||
| ctx = s.Config().ContextByName(ctxName) | ||
| if ctx == nil { | ||
| return nil, fmt.Errorf("context not found: %v", ctxName) | ||
| } | ||
| } else { | ||
| ctx = s.Config().ActiveContext() | ||
| if ctx == nil { | ||
| return nil, fmt.Errorf("no active context") | ||
| } | ||
| } | ||
|
|
||
| return ctx, nil | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters