From 2b612711a44cea1218d7aac57c8c837618737a2e Mon Sep 17 00:00:00 2001 From: Eugene Dementyev Date: Fri, 8 Oct 2021 15:17:07 +1300 Subject: [PATCH] Adds hosts autocompletion for "connect" command! --- README.md | 11 +++++++ cmd/completion.go | 70 ++++++++++++++++++++++++++++++++++++++++++ cmd/connect.go | 18 ++++++++--- cmd/reconf.go | 5 ++- cmd/root.go | 5 +-- cmd/test.go | 5 ++- cmd/update.go | 5 ++- lib/cache/interface.go | 2 ++ lib/cache/yaml.go | 7 +++++ 9 files changed, 117 insertions(+), 11 deletions(-) create mode 100644 cmd/completion.go diff --git a/README.md b/README.md index 3a30bd3..6889e6e 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,17 @@ There are certain prerequisites: 4. But it's boring to look up the instance id every time so you can run `aws-ssh update` to generate cache of all EC2 instances across all available AWS profiles 5. Then just run `aws-ssh connect` to search for the right instance and press "Enter" +### ec2 connect with host autocompletion! + +You can also use hosts autocompletion! Refer to `aws-ssh completion -h` instructions how to set it up, then run like: + +```bash +$aws-ssh connect -i +# or +$aws-ssh connect -i profile- + +``` + ### Use reconf feature Instead of using EC2 connect, one can have their ssh keys directly on the instances, so for those cases there is `aws-ssh reconf` command which just generates ssh config to be included in the main one. diff --git a/cmd/completion.go b/cmd/completion.go new file mode 100644 index 0000000..c5bde04 --- /dev/null +++ b/cmd/completion.go @@ -0,0 +1,70 @@ +package cmd + +import ( + "os" + + "github.com/spf13/cobra" +) + +var completionCmd = &cobra.Command{ + Use: "completion [bash|zsh|fish|powershell]", + Short: "Generate completion script", + Long: `To load completions: + +Bash: + + $ source <(aws-ssh completion bash) + + # To load completions for each session, execute once: + # Linux: + $ aws-ssh completion bash > /etc/bash_completion.d/aws-ssh + # macOS: + $ aws-ssh completion bash > /usr/local/etc/bash_completion.d/aws-ssh + +Zsh: + + # If shell completion is not already enabled in your environment, + # you will need to enable it. You can execute the following once: + + $ echo "autoload -U compinit; compinit" >> ~/.zshrc + + # To load completions for each session, execute once: + $ aws-ssh completion zsh > "${fpath[1]}/_aws-ssh" + + # You will need to start a new shell for this setup to take effect. + +fish: + + $ aws-ssh completion fish | source + + # To load completions for each session, execute once: + $ aws-ssh completion fish > ~/.config/fish/completions/aws-ssh.fish + +PowerShell: + + PS> aws-ssh completion powershell | Out-String | Invoke-Expression + + # To load completions for every new session, run: + PS> aws-ssh completion powershell > aws-ssh.ps1 + # and source this file from your PowerShell profile. +`, + DisableFlagsInUseLine: true, + ValidArgs: []string{"bash", "zsh", "fish", "powershell"}, + Args: cobra.ExactValidArgs(1), + Run: func(cmd *cobra.Command, args []string) { + switch args[0] { + case "bash": + cmd.Root().GenBashCompletion(os.Stdout) + case "zsh": + cmd.Root().GenZshCompletion(os.Stdout) + case "fish": + cmd.Root().GenFishCompletion(os.Stdout, true) + case "powershell": + cmd.Root().GenPowerShellCompletionWithDesc(os.Stdout) + } + }, +} + +func init() { + rootCmd.AddCommand(completionCmd) +} diff --git a/cmd/connect.go b/cmd/connect.go index 7d9daa7..dbb0099 100644 --- a/cmd/connect.go +++ b/cmd/connect.go @@ -15,10 +15,6 @@ import ( var connectCmd = &cobra.Command{ Use: "connect [ssh command (ssh -tt {host})]", Short: "SSH into the EC2 instance using ec2 connect feature", - // override the default prerun - // as we don't need any profiles here - PersistentPreRun: func(cmd *cobra.Command, args []string) { - }, Long: `aws-ssh connects to the EC2 instance using ec2 connect feature. It makes a special API call to upload the first public key from your running ssh agent and then runs ssh command. @@ -99,10 +95,24 @@ func init() { defaultSSHConfigFile := path.Join(homeDir, ".ssh", "ec2_connect_config") connectCmd.Flags().StringP("instanceid", "i", "", "Instance ID to connect to") + connectCmd.Flags().StringP("user", "u", "", "Existing user on the instance") connectCmd.Flags().StringP("ssh-config-path", "c", defaultSSHConfigFile, "Path to the ssh config to generate") viper.BindPFlag("instanceid", connectCmd.Flags().Lookup("instanceid")) viper.BindPFlag("user", connectCmd.Flags().Lookup("user")) viper.BindPFlag("ssh-config-path", connectCmd.Flags().Lookup("ssh-config-path")) + + // custom completion for instances + connectCmd.RegisterFlagCompletionFunc("instanceid", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) { + cache := cache.NewYAMLCache(viper.GetString("cache-dir")) + names, err := cache.ListCanonicalNames() + if err != nil { + // can happen if there's no cache yet + // safe to just ignore it here + log.Fatalf("wtf") + } + return names, cobra.ShellCompDirectiveNoFileComp + }) + rootCmd.AddCommand(connectCmd) } diff --git a/cmd/reconf.go b/cmd/reconf.go index f8a3f4c..1418b60 100644 --- a/cmd/reconf.go +++ b/cmd/reconf.go @@ -8,7 +8,10 @@ import ( ) var reconfCmd = &cobra.Command{ - Use: "reconf ", + Use: "reconf ", + PersistentPreRun: func(cmd *cobra.Command, args []string) { + initConfig() + }, Args: cobra.ExactArgs(1), Short: "Creates a new ssh config", Long: `Reconfigures your ssh by creating a new config for it. Only one argument is required, diff --git a/cmd/root.go b/cmd/root.go index 27a3d66..ca43129 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -18,10 +18,7 @@ var cfgFile string // rootCmd represents the base command when called without any subcommands var rootCmd = &cobra.Command{ - Use: "aws-ssh", - PersistentPreRun: func(cmd *cobra.Command, args []string) { - initConfig() - }, + Use: "aws-ssh", Short: "Describe your AWS and get ssh config to connect to ec2 instances", Long: `This program goes through all available AWS accounts in parallel and determines IP addresses of ec2 instances. It also detects so-called "bastion" instances. diff --git a/cmd/test.go b/cmd/test.go index fecda3d..b130e74 100644 --- a/cmd/test.go +++ b/cmd/test.go @@ -9,7 +9,10 @@ import ( ) var testCmd = &cobra.Command{ - Use: "test", + Use: "test", + PersistentPreRun: func(cmd *cobra.Command, args []string) { + initConfig() + }, Short: "Tests the profiles", Long: `aws-ssh test tests the AWS profiles found in ~/.aws/config and ~/.aws/credentials (unless -p option is provided) diff --git a/cmd/update.go b/cmd/update.go index 3643ef2..8ba0c70 100644 --- a/cmd/update.go +++ b/cmd/update.go @@ -11,7 +11,10 @@ import ( // updateCmd represents the update command var updateCmd = &cobra.Command{ - Use: "update", + Use: "update", + PersistentPreRun: func(cmd *cobra.Command, args []string) { + initConfig() + }, Short: "Updates cache of ssh entries", Long: ` Cache is important for sophisticated behaviour of "connect" command, diff --git a/lib/cache/interface.go b/lib/cache/interface.go index 436e1bc..3f3b712 100644 --- a/lib/cache/interface.go +++ b/lib/cache/interface.go @@ -14,4 +14,6 @@ type Cache interface { // If name is empty or there is no exact match, // it switches to the fuzzy search mode Lookup(name string) (lib.SSHEntry, error) + // ListCanonicalNames() returns all known canonical host names from the cache + ListCanonicalNames() ([]string, error) } diff --git a/lib/cache/yaml.go b/lib/cache/yaml.go index 8dcd72f..3cf5cfc 100644 --- a/lib/cache/yaml.go +++ b/lib/cache/yaml.go @@ -168,6 +168,13 @@ func (y *YAMLCache) Lookup(name string) (lib.SSHEntry, error) { return entry, nil } +func (y *YAMLCache) ListCanonicalNames() ([]string, error) { + if err := y.loadIndex(); err != nil { + return []string{}, nil + } + return y.index.CanonicalNames, nil +} + func NewYAMLCache(basedir string) Cache { return &YAMLCache{basedir: basedir}