Merge branch 'ssh-kitten' of https://github.com/page-down/kitty

kittens/ssh/main.py

@@ -226,7 +226,7 @@ def safe_remove(x: str) -> None:
 
 
 def prepare_script(ans: str, replacements: Dict[str, str], script_type: str) -> str:
-    for k in ('EXEC_CMD',):
+    for k in ('EXEC_CMD', 'EXPORT_HOME_CMD'):
         replacements[k] = replacements.get(k, '')
 
     def sub(m: 're.Match[str]') -> str:
@@ -242,7 +242,19 @@ def prepare_exec_cmd(remote_args: Sequence[str], is_python: bool) -> str:
     if is_python:
         return standard_b64encode(' '.join(remote_args).encode('utf-8')).decode('ascii')
     args = ' '.join(c.replace("'", """'"'"'""") for c in remote_args)
-    return f"""exec "$login_shell" -c '{args}'"""
+    return f"""unset KITTY_SHELL_INTEGRATION; exec "$login_shell" -c '{args}'"""
+
+
+def prepare_export_home_cmd(ssh_opts: SSHOptions, is_python: bool) -> str:
+    home = ssh_opts.env.get('HOME')
+    if home == '_kitty_copy_env_var_':
+        home = os.environ.get('HOME')
+    if home:
+        if is_python:
+            return standard_b64encode(home.encode('utf-8')).decode('ascii')
+        else:
+            return f'export HOME={quote_env_val(home)}; cd "$HOME"'
+    return ''
 
 
 def bootstrap_script(
@@ -252,7 +264,9 @@ def bootstrap_script(
 ) -> Tuple[str, Dict[str, str], SharedMemory]:
     if request_id is None:
         request_id = os.environ['KITTY_PID'] + '-' + os.environ['KITTY_WINDOW_ID']
-    exec_cmd = prepare_exec_cmd(remote_args, script_type == 'py') if remote_args else ''
+    is_python = script_type == 'py'
+    export_home_cmd = prepare_export_home_cmd(ssh_opts, is_python) if 'HOME' in ssh_opts.env else ''
+    exec_cmd = prepare_exec_cmd(remote_args, is_python) if remote_args else ''
     with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f:
         ans = f.read()
     pw = secrets.token_hex()
@@ -265,7 +279,9 @@ def bootstrap_script(
         atexit.register(shm.unlink)
     sensitive_data = {'REQUEST_ID': request_id, 'DATA_PASSWORD': pw, 'PASSWORD_FILENAME': shm.name}
     replacements = {
-        'EXEC_CMD': exec_cmd, 'TEST_SCRIPT': test_script, 'REQUEST_DATA': '1' if request_data else '0', 'ECHO_ON': '1' if echo_on else '0',
+        'EXPORT_HOME_CMD': export_home_cmd,
+        'EXEC_CMD': exec_cmd, 'TEST_SCRIPT': test_script,
+        'REQUEST_DATA': '1' if request_data else '0', 'ECHO_ON': '1' if echo_on else '0',
     }
     sd = replacements.copy()
     if request_data:

shell-integration/ssh/bootstrap-utils.sh

@@ -178,12 +178,7 @@ prepare_for_exec() {
     fi
     [ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data"
 
-    if [ -n "$KITTY_LOGIN_SHELL" ]; then
-        login_shell="$KITTY_LOGIN_SHELL"
-        unset KITTY_LOGIN_SHELL
-    else
-        using_getent || using_id || using_python || using_perl || using_passwd || using_shell_env || login_shell="sh"
-    fi
+    [ -n "$login_shell" ] || using_getent || using_id || using_python || using_perl || using_passwd || using_shell_env || login_shell="sh"
     shell_name=$(command basename $login_shell)
     [ -n "$login_cwd" ] && cd "$login_cwd"
 }

shell-integration/ssh/bootstrap.py

@@ -20,8 +20,13 @@
 data_dir = shell_integration_dir = ''
 request_data = int('REQUEST_DATA')
 leading_data = b''
-HOME = os.path.expanduser('~')
 login_shell = pwd.getpwuid(os.geteuid()).pw_shell or os.environ.get('SHELL') or 'sh'
+export_home_cmd = b'EXPORT_HOME_CMD'
+if export_home_cmd:
+    HOME = base64.standard_b64decode(export_home_cmd).decode('utf-8')
+    os.chdir(HOME)
+else:
+    HOME = os.path.expanduser('~')
 
 
 def set_echo(fd, on=False):
@@ -219,7 +224,7 @@ def exec_zsh_with_integration():
         os.environ['KITTY_ORIG_ZDOTDIR'] = zdotdir
     # dont prevent zsh-newuser-install from running
     for q in ('.zshrc', '.zshenv', '.zprofile', '.zlogin'):
-        if os.path.exists(os.path.join(HOME, q)):
+        if os.path.exists(os.path.join(zdotdir, q)):
             os.environ['ZDOTDIR'] = shell_integration_dir + '/zsh'
             os.execlp(login_shell, os.path.basename(login_shell), '-l')
     os.environ.pop('KITTY_ORIG_ZDOTDIR', None)  # ensure this is not propagated
@@ -271,6 +276,7 @@ def main():
     ksi = frozenset(filter(None, os.environ.get('KITTY_SHELL_INTEGRATION', '').split()))
     exec_cmd = b'EXEC_CMD'
     if exec_cmd:
+        os.environ.pop('KITTY_SHELL_INTEGRATION', None)
         cmd = base64.standard_b64decode(exec_cmd).decode('utf-8')
         os.execlp(login_shell, os.path.basename(login_shell), '-c', cmd)
     TEST_SCRIPT  # noqa

shell-integration/ssh/bootstrap.sh

@@ -62,6 +62,8 @@ dcs_to_kitty() { printf "\033P@kitty-$1|%s\033\134" "$(printf "%s" "$2" | base64
 debug() { dcs_to_kitty "print" "debug: $1"; }
 echo_via_kitty() { dcs_to_kitty "echo" "$1"; }
 
+# If $HOME is configured set it here
+EXPORT_HOME_CMD
 # ensure $HOME is set
 [ -z "$HOME" ] && HOME=~
 # ensure $USER is set
@@ -69,6 +71,7 @@ echo_via_kitty() { dcs_to_kitty "echo" "$1"; }
 [ -z "$USER" ] && USER="$(command whoami 2> /dev/null)"
 
 leading_data=""
+login_shell=""
 login_cwd=""
 
 request_data="REQUEST_DATA"
@@ -98,6 +101,8 @@ untar_and_read_env() {
     data_dir="$HOME/$KITTY_SSH_KITTEN_DATA_DIR"
     shell_integration_dir="$data_dir/shell-integration"
     unset KITTY_SSH_KITTEN_DATA_DIR
+    login_shell="$KITTY_LOGIN_SHELL"
+    unset KITTY_LOGIN_SHELL
     login_cwd="$KITTY_LOGIN_CWD"
     unset KITTY_LOGIN_CWD
     compile_terminfo "$tdir/home"