From ec420b80126784ecfe1d8299bd60cc7f90263f30 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Fri, 25 Jun 2021 10:48:58 +0530
Subject: [PATCH] Start work on completion for ssh kitten

---
 kittens/ssh/completion.py | 65 +++++++++++++++++++++++++++++++++++++++
 kittens/ssh/main.py       | 23 ++++++--------
 2 files changed, 74 insertions(+), 14 deletions(-)
 create mode 100644 kittens/ssh/completion.py

diff --git a/kittens/ssh/completion.py b/kittens/ssh/completion.py
new file mode 100644
index 000000000..31e3db6b7
--- /dev/null
+++ b/kittens/ssh/completion.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+# vim:fileencoding=utf-8
+# License: GPLv3 Copyright: 2021, Kovid Goyal <kovid at kovidgoyal.net>
+
+import os
+import subprocess
+from typing import Dict, Iterator, Tuple
+from kitty.types import run_once
+
+
+def iter_known_hosts() -> Iterator[str]:
+    try:
+        f = open(os.path.expanduser('~/.ssh/config'))
+    except OSError:
+        pass
+    else:
+        for line in f:
+            parts = line.split()
+            if parts and parts[0] == 'Host' and len(parts) > 1:
+                yield parts[1]
+
+    try:
+        f = open(os.path.expanduser('~/.ssh/known_hosts'))
+    except OSError:
+        pass
+    else:
+        for line in f:
+            parts = line.split()
+            if parts:
+                yield parts[0]
+
+
+@run_once
+def known_hosts() -> Tuple[str, ...]:
+    return tuple(iter_known_hosts())
+
+
+@run_once
+def ssh_options() -> Dict[str, str]:
+    stderr = subprocess.Popen(['ssh'], stderr=subprocess.PIPE).stderr
+    assert stderr is not None
+    raw = stderr.read().decode('utf-8')
+    ans: Dict[str, str] = {}
+    pos = 0
+    while True:
+        pos = raw.find('[', pos)
+        if pos < 0:
+            break
+        num = 1
+        epos = pos
+        while num > 0:
+            epos += 1
+            if raw[epos] not in '[]':
+                continue
+            num += 1 if raw[epos] == '[' else -1
+        q = raw[pos+1:epos]
+        pos = epos
+        if len(q) < 2 or q[0] != '-':
+            continue
+        if ' ' in q:
+            opt, desc = q.split(' ', 1)
+            ans[opt[1:]] = desc
+        else:
+            ans.update(dict.fromkeys(q[1:], ''))
+    return ans
diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index 6a53b1d9e..2e2c4d4a7 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -3,12 +3,12 @@
 # License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
 
 import os
-import re
 import shlex
 import subprocess
 import sys
 from contextlib import suppress
 from typing import List, NoReturn, Optional, Set, Tuple
+from .completion import ssh_options
 
 from kitty.utils import SSHConnectionData
 
@@ -129,20 +129,15 @@ os.execlp(shell_path, shell_name)
 
 
 def get_ssh_cli() -> Tuple[Set[str], Set[str]]:
-    other_ssh_args: List[str] = []
-    boolean_ssh_args: List[str] = []
-    stderr = subprocess.Popen(['ssh'], stderr=subprocess.PIPE).stderr
-    assert stderr is not None
-    raw = stderr.read().decode('utf-8')
-    for m in re.finditer(r'\[(.+?)\]', raw):
-        q = m.group(1)
-        if len(q) < 2 or q[0] != '-':
-            continue
-        if ' ' in q:
-            other_ssh_args.append(q[1])
+    other_ssh_args: Set[str] = set()
+    boolean_ssh_args: Set[str] = set()
+    for k, v in ssh_options().items():
+        k = '-' + k
+        if v:
+            other_ssh_args.add(k)
         else:
-            boolean_ssh_args.extend(q[1:])
-    return set('-' + x for x in boolean_ssh_args), set('-' + x for x in other_ssh_args)
+            boolean_ssh_args.add(k)
+    return boolean_ssh_args, other_ssh_args
 
 
 def get_connection_data(args: List[str]) -> Optional[SSHConnectionData]: