Adds hosts autocompletion for "connect" command!

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 <TAB>
+# or
+$aws-ssh connect -i profile- <TAB>
+
+```
+
 ### 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.

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)
+}

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)
 }

cmd/reconf.go

@@ -8,7 +8,10 @@ import (
 )
 
 var reconfCmd = &cobra.Command{
-	Use:   "reconf <filename>",
+	Use: "reconf <filename>",
+	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,

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.

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)

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,

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)
 }

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}