Work on enabling shell integration over ssh

kittens/ssh/main.py

@@ -4,13 +4,14 @@
 import atexit
 import io
 import os
+import re
 import shlex
 import subprocess
 import sys
 import tarfile
 import tempfile
 from contextlib import suppress
-from typing import Iterator, List, NoReturn, Optional, Set, Tuple
+from typing import Dict, Iterator, List, NoReturn, Optional, Set, Tuple
 
 from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
 from kitty.short_uuid import uuid4
@@ -18,8 +19,10 @@ from kitty.utils import SSHConnectionData
 
 from .completion import complete, ssh_options
 
+DEFAULT_SHELL_INTEGRATION_DEST = '.local/share/kitty-ssh-kitten/shell-integration'
 
-def make_tarfile(hostname: str = '') -> bytes:
+
+def make_tarfile(hostname: str = '', shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> bytes:
 
     def filter_files(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]:
         if tarinfo.name.endswith('ssh/bootstrap.sh') or tarinfo.name.endswith('ssh/bootstrap.py'):
@@ -30,12 +33,12 @@ def make_tarfile(hostname: str = '') -> bytes:
 
     buf = io.BytesIO()
     with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
-        tf.add(shell_integration_dir, arcname='.local/share/kitty-ssh-kitten/shell-integration', filter=filter_files)
+        tf.add(shell_integration_dir, arcname=shell_integration_dest, filter=filter_files)
         tf.add(terminfo_dir, arcname='.terminfo', filter=filter_files)
     return buf.getvalue()
 
 
-def get_ssh_data(msg: str) -> Iterator[bytes]:
+def get_ssh_data(msg: str, shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> Iterator[bytes]:
     yield b"KITTY_SSH_DATA_START\n"
     try:
         hostname, pwfilename, pw = msg.split(':', 2)
@@ -50,7 +53,7 @@ def get_ssh_data(msg: str) -> Iterator[bytes]:
         yield b' incorrect ssh data password\n'
     else:
         try:
-            data = make_tarfile(hostname)
+            data = make_tarfile(hostname, shell_integration_dest)
         except Exception:
             yield b' error while gathering ssh data\n'
         else:
@@ -68,21 +71,28 @@ def safe_remove(x: str) -> None:
         os.remove(x)
 
 
-def prepare_script(ans: str, EXEC_CMD: str = '') -> str:
+def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
-    ans = ans.replace('EXEC_CMD', EXEC_CMD, 1)
     pw = uuid4()
     with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', dir=cache_dir(), delete=False) as tf:
         tf.write(pw.encode('utf-8'))
     atexit.register(safe_remove, tf.name)
-    ans = ans.replace('DATA_PASSWORD', pw, 1)
+    replacements['DATA_PASSWORD'] = pw
-    ans = ans.replace("PASSWORD_FILENAME", os.path.basename(tf.name))
+    replacements['PASSWORD_FILENAME'] = os.path.basename(tf.name)
-    return ans
+    for k in ('EXEC_CMD', 'OVERRIDE_LOGIN_SHELL'):
+        replacements[k] = replacements.get(k, '')
+    replacements['SHELL_INTEGRATION_DIR'] = replacements.get('SHELL_INTEGRATION_DIR', DEFAULT_SHELL_INTEGRATION_DEST)
+    replacements['SHELL_INTEGRATION_VALUE'] = replacements.get('SHELL_INTEGRATION_VALUE', 'enabled')
+
+    def sub(m: 're.Match[str]') -> str:
+        return replacements[m.group()]
+
+    return re.sub('|'.join(fr'\b{k}\b' for k in replacements), sub, ans)
 
 
-def bootstrap_script(EXEC_CMD: str = '', script_type: str = 'sh') -> str:
+def bootstrap_script(script_type: str = 'sh', **replacements: str) -> str:
     with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f:
         ans = f.read()
-    return prepare_script(ans, EXEC_CMD)
+    return prepare_script(ans, replacements)
 
 
 SHELL_SCRIPT = '''\

kitty_tests/ssh.py

@@ -45,10 +45,11 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
             q = shutil.which(sh)
             if q:
                 with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
-                    script = bootstrap_script('echo TEST_DONE; return 0')
+                    script = bootstrap_script(EXEC_CMD='echo UNTAR_DONE; exit 0')
                     env = basic_shell_env(tdir)
                     pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=tdir, env=env)
                     self.check_bootstrap(tdir, pty)
 
     def check_bootstrap(self, home_dir, pty):
-        pty.wait_till(lambda: 'TEST_DONE' in pty.screen_contents())
+        pty.wait_till(lambda: 'UNTAR_DONE' in pty.screen_contents())
+        self.assertTrue(os.path.exists(os.path.join(home_dir, '.terminfo/kitty.terminfo')))

shell-integration/ssh/bootstrap.sh

@@ -83,3 +83,153 @@ fi
 
 # If a command was passed to SSH execute it here
 EXEC_CMD
+
+shell_integration_dir="$HOME/SHELL_INTEGRATION_DIR"
+
+login_shell_is_ok() {
+    if [ -z "$login_shell" ] || [ ! -x "$login_shell" ]; then return 1; fi
+    case "$login_shell" in
+        *sh) return 0;
+    esac
+    return 1;
+}
+
+detect_python() {
+    python=$(command -v python3)
+    if [ -z "$python" ]; then python=$(command -v python2); fi
+    if [ -z "$python" ]; then python=python; fi
+    if [ -z "$python" || ! -x "$python" ]; then return 1; fi
+    return 0;
+}
+
+using_getent() {
+    cmd=$(command -v getent)
+    if [ -n "$cmd" ]; then 
+        output=$($cmd passwd $USER 2>/dev/null)
+        if [ $? = 0 ]; then 
+            login_shell=$(echo $output | cut -d: -f7);
+            if login_shell_is_ok; then return 0; fi
+        fi
+    fi
+    return 1;
+}
+
+using_id() {
+    cmd=$(command -v id)
+    if [ -n "$cmd" ]; then 
+        output=$($cmd -P $USER 2>/dev/null)
+        if [ $? = 0 ]; then 
+            login_shell=$(echo $output | cut -d: -f7);
+            if login_shell_is_ok; then return 0; fi
+        fi
+    fi
+    return 1;
+}
+
+using_passwd() {
+    cmd=$(command -v grep)
+    if [ -n "$cmd" ]; then 
+        output=$($cmd "^$USER:" /etc/passwd 2>/dev/null)
+        if [ $? = 0 ]; then 
+            login_shell=$(echo $output | cut -d: -f7);
+            if login_shell_is_ok; then return 0; fi
+        fi
+    fi
+    return 1;
+}
+
+using_python() {
+    if detect_python; then
+        output=$($python -c "import pwd, os; print(pwd.getpwuid(os.geteuid()).pw_shell)")
+        if [ $? = 0 ]; then 
+            login_shell=$output; 
+            if login_shell_is_ok; then return 0; fi
+        fi
+    fi
+    return 1
+}
+
+execute_with_python() {
+    if detect_python; then
+        exec $python -c "import os; os.execl('$login_shell', '-' '$shell_name')"
+    fi
+    return 1;
+}
+
+LOGIN_SHELL="OVERRIDE_LOGIN_SHELL"
+if [ -n "$LOGIN_SHELL" ]; then
+    login_shell="$LOGIN_SHELL"
+else
+    using_getent || using_id || using_python || using_passwd || die "Could not detect login shell";
+fi
+shell_name=$(basename $login_shell)
+
+export KITTY_SHELL_INTEGRATION="SHELL_INTEGRATION_VALUE"
+
+exec_bash_with_integration() {
+    export ENV="$shell_integration_dir/bash/kitty.bash"
+    export KITTY_BASH_INJECT="1"
+    exec "$login_shell" "--posix"
+}
+
+exec_zsh_with_integration() {
+    zdotdir="$ZDOTDIR"
+    if [ -z "$zdotdir" ]; then 
+        zdotdir=~; 
+        unset KITTY_ORIG_ZDOTDIR
+    else
+        export KITTY_ORIG_ZDOTDIR="$zdotdir"
+    fi
+    # dont prevent zsh-new-user from running
+    if [ -e "$zdotdir/.zshrc" || -e "$zdotdir/.zshenv" || -e "$zdotdir/.zprofile" || -e "$zdotdir/.zlogin" ]; then
+        export ZDOTDIR="$shell_integration_dir/zsh"
+        exec "$login_shell" "-l"
+    fi
+}
+
+exec_fish_with_integration() {
+    if [ -z "$XDG_DATA_DIRS" ]; then
+        export XDG_DATA_DIRS="$shell_integration_dir"
+    else
+        export XDG_DATA_DIRS="$shell_integration_dir:$XDG_DATA_DIRS"
+    fi
+    export KITTY_FISH_XDG_DATA_DIR="$shell_integration_dir"
+    exec "$login_shell" "-l"
+}
+
+exec_with_shell_integration() {
+    case "$login_shell" in
+        *"zsh")
+            exec_zsh_with_integration
+            ;;
+        *"bash")
+            exec_bash_with_integration
+            ;;
+        *"fish")
+            exec_fish_with_integration
+            ;;
+    esac
+}
+
+case "$KITTY_SHELL_INTEGRATION" in
+    "")
+        unset KITTY_SHELL_INTEGRATION
+        ;;
+    *"no-rc"*)
+        ;;
+    *)
+        exec_with_shell_integration
+        unset KITTY_SHELL_INTEGRATION
+        ;;
+esac
+
+# We need to pass the first argument to the executed program with a leading -
+# to make sure the shell executes as a login shell. Note that not all shells
+# support exec -a so we use the below to try to detect such shells
+shell_name=$(basename $login_shell)
+if [ -z "$PIPESTATUS" ]; then
+    # the dash shell does not support exec -a and also does not define PIPESTATUS
+    execute_with_python
+    exec $login_shell "-l"
+fi
+exec -a "-$shell_name" $login_shell