From e0254744c3581ef2c3d7c3d7f3c597b0aeab391e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Skytt=C3=A4?= <ville.skytta@upcloud.com>
Date: Wed, 27 Apr 2022 10:48:20 +0300
Subject: [PATCH] feat(client): complete sshi ssh on (unhashed) known hosts

bash-completion would have more extensive ssh completions to reuse in
shell function format, but cobra doesn't make integrating with it
feasible.

Implement basic known hosts completion ourselves. Requires unhashed
known hosts (`HashKnownHosts no`) to work.
---
 cliclient/cmd/ssh.go | 78 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 77 insertions(+), 1 deletion(-)

diff --git a/cliclient/cmd/ssh.go b/cliclient/cmd/ssh.go
index dedb2887..4dd4e69b 100644
--- a/cliclient/cmd/ssh.go
+++ b/cliclient/cmd/ssh.go
@@ -1,9 +1,85 @@
 package cmd
 
 import (
+	"bufio"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+
 	"github.com/spf13/cobra"
+	"golang.org/x/crypto/ssh"
 )
 
+func expandFilename(fn string) string {
+	if strings.HasPrefix(fn, "~/") {
+		if hd, err := os.UserHomeDir(); err == nil {
+			fn = filepath.Join(hd, fn[2:])
+		}
+	}
+	return fn
+}
+
+func sshKnownHostnames(in []byte) []string {
+	var hosts, ret []string
+	var err error
+	for {
+		_, hosts, _, _, in, err = ssh.ParseKnownHosts(in)
+		if err != nil {
+			break
+		}
+		for _, h := range hosts {
+			if i := strings.IndexRune(h, ']'); i != -1 && strings.HasPrefix(h, "[") {
+				h = h[1:i]
+			}
+			// Skip wildcards, negations, and hashed entries
+			if !strings.ContainsAny(h, "*?!|") {
+				ret = append(ret, h)
+			}
+		}
+	}
+	return ret
+}
+
+func sshHostnames(cmd *cobra.Command, args []string, toComplete string) (ret []string, compDir cobra.ShellCompDirective) {
+	compDir = cobra.ShellCompDirectiveNoFileComp
+	if len(args) != 0 {
+		return
+	}
+
+	ssh := exec.CommandContext(cmd.Context(), "ssh", "-G", toComplete)
+	stdout, err := ssh.StdoutPipe()
+	if err != nil || ssh.Start() != nil {
+		return
+	}
+	var fns []string
+	s := bufio.NewScanner(stdout)
+	for s.Scan() {
+		flds := strings.Fields(s.Text())
+		if len(flds) > 1 && (flds[0] == "globalknownhostsfile" || flds[0] == "userknownhostsfile") {
+			// These seem to be output space separated, with no escaping, so we're out of luck with filenames containing spaces
+			for _, fn := range flds[1:] {
+				if fn = expandFilename(fn); fn != "" {
+					fns = append(fns, fn)
+				}
+			}
+		}
+	}
+	_ = ssh.Wait()
+
+	buf := make([]byte, 1024*1024)
+	for _, fn := range fns {
+		if f, err := os.Open(fn); err == nil {
+			n, err := f.Read(buf)
+			_ = f.Close()
+			if err == nil {
+				ret = append(ret, sshKnownHostnames(buf[:n])...)
+			}
+		}
+	}
+	return
+}
+
 var SshCmd = &cobra.Command{
 	Use:                "ssh",
 	Short:              "Invoke ssh command with signed certificate on ssh-agent",
@@ -12,7 +88,7 @@ var SshCmd = &cobra.Command{
 		ignoreFlagsAfter("ssh")
 		return runExecCommand(RootCmd.Flags().Args()[1:])
 	},
-	ValidArgsFunction: noCompletion,
+	ValidArgsFunction: sshHostnames,
 }
 
 func init() {