From 5c02c370d4894774c0c1eb933425fcf3cb8e3f6f Mon Sep 17 00:00:00 2001
From: pagedown <pg.dn@outlook.com>
Date: Thu, 24 Mar 2022 22:00:41 +0800
Subject: [PATCH] ssh kitten: Allow to configure HOME environment variable

Useful if the user does not have a home directory or the directory is
not writable.
---
 kittens/ssh/main.py                | 22 +++++++++++++++++++---
 shell-integration/ssh/bootstrap.py |  7 ++++++-
 shell-integration/ssh/bootstrap.sh |  2 ++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index cf6b1ca26..8b94e7990 100644
--- a/kittens/ssh/main.py
+++ b/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:
@@ -245,6 +245,18 @@ def prepare_exec_cmd(remote_args: Sequence[str], is_python: bool) -> str:
     return f"""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(
     ssh_opts: SSHOptions, script_type: str = 'sh', remote_args: Sequence[str] = (),
     test_script: str = '', request_id: Optional[str] = None, cli_hostname: str = '', cli_uname: str = '',
@@ -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:
diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py
index f47debd62..e4976e08d 100644
--- a/shell-integration/ssh/bootstrap.py
+++ b/shell-integration/ssh/bootstrap.py
@@ -20,8 +20,13 @@ echo_on = int('ECHO_ON')
 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):
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index 137399eff..222c0c46b 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/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