Implement setting of env vars

kittens/ssh/main.py

@@ -3,6 +3,7 @@
 
 import atexit
 import io
+import json
 import os
 import re
 import shlex
@@ -17,15 +18,36 @@ from typing import (
 )
 
 from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
+from kitty.fast_data_types import get_options
 from kitty.short_uuid import uuid4
 from kitty.types import run_once
 from kitty.utils import SSHConnectionData
 
 from .completion import complete, ssh_options
 from .options.types import Options as SSHOptions
+from .options.utils import DELETE_ENV_VAR
 
 
-def make_tarfile(ssh_opts: SSHOptions) -> bytes:
+def serialize_env(env: Dict[str, str], base_env: Dict[str, str]) -> bytes:
+    lines = []
+
+    def a(k: str, val: str) -> None:
+        lines.append(f'export {k}={shlex.quote(val)}')
+
+    for k in sorted(env):
+        v = env[k]
+        if v is DELETE_ENV_VAR:
+            lines.append(f'unset {shlex.quote(k)}')
+        elif v == '_kitty_copy_env_var_':
+            q = base_env.get(k)
+            if q is not None:
+                a(k, q)
+        else:
+            a(k, v)
+    return '\n'.join(lines).encode('utf-8')
+
+
+def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
 
     def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo:
         tarinfo.uname = tarinfo.gname = 'kitty'
@@ -48,9 +70,6 @@ def make_tarfile(ssh_opts: SSHOptions) -> bytes:
             return None
         return normalize_tarinfo(tarinfo)
 
-    buf = io.BytesIO()
-    with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
-        rd = ssh_opts.remote_dir.rstrip('/')
     from kitty.shell_integration import get_effective_ksi_env_var
     if ssh_opts.shell_integration == 'inherit':
         ksi = get_effective_ksi_env_var()
@@ -58,9 +77,24 @@ def make_tarfile(ssh_opts: SSHOptions) -> bytes:
         from kitty.options.types import Options
         from kitty.options.utils import shell_integration
         ksi = get_effective_ksi_env_var(Options({'shell_integration': shell_integration(ssh_opts.shell_integration)}))
+
+    env = {
+        'TERM': get_options().term,
+        'COLORTERM': 'truecolor',
+    }
+    for q in ('KITTY_WINDOW_ID', 'WINDOWID'):
+        val = os.environ.get(q)
+        if val is not None:
+            env[q] = val
+    env.update(ssh_opts.env)
+    env['KITTY_SHELL_INTEGRATION'] = ksi or DELETE_ENV_VAR
+    env_script = serialize_env(env, base_env)
+    buf = io.BytesIO()
+    with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
+        rd = ssh_opts.remote_dir.rstrip('/')
+        add_data_as_file(tf, rd + '/settings/env-vars.sh', env_script)
         if ksi:
             tf.add(shell_integration_dir, arcname=rd + '/shell-integration', filter=filter_files)
-            add_data_as_file(tf, rd + '/settings/ksi_env_var', ksi)
         tf.add(terminfo_dir, arcname='.terminfo', filter=filter_files)
     return buf.getvalue()
 
@@ -92,16 +126,19 @@ def get_ssh_data(msg: str, ssh_opts: Optional[Dict[str, SSHOptions]] = None) ->
         yield fmt_prefix('!invalid ssh data request message')
     else:
         try:
-            with open(os.path.join(cache_dir(), pwfilename)) as f:
+            with open(os.path.join(cache_dir(), pwfilename), 'rb') as f:
                 os.unlink(f.name)
-                if pw != f.read():
+                env_data = json.load(f)
+                if pw != env_data['pw']:
                     raise ValueError('Incorrect password')
         except Exception:
+            import traceback
+            traceback.print_exc()
             yield fmt_prefix('!incorrect ssh data password')
         else:
             resolved_ssh_opts = options_for_host(hostname, ssh_opts)
             try:
-                data = make_tarfile(resolved_ssh_opts)
+                data = make_tarfile(resolved_ssh_opts, env_data['env'])
             except Exception:
                 yield fmt_prefix('!error while gathering ssh data')
             else:
@@ -119,8 +156,9 @@ def safe_remove(x: str) -> None:
 
 def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
     pw = uuid4()
-    with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', dir=cache_dir(), delete=False) as tf:
+    with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=cache_dir(), delete=False) as tf:
-        tf.write(pw.encode('utf-8'))
+        data = {'pw': pw, 'env': dict(os.environ)}
+        tf.write(json.dumps(data).encode('utf-8'))
     atexit.register(safe_remove, tf.name)
     replacements['DATA_PASSWORD'] = pw
     replacements['PASSWORD_FILENAME'] = os.path.basename(tf.name)

kittens/ssh/options/definition.py

@@ -28,13 +28,13 @@ to SSH to connect to it.
 
 opt('remote_dir', '.local/share/kitty-ssh-kitten', long_text='''
 The location on the remote computer where the files needed for this kitten
-are installed. The location is relative to the HOME directory.
+are installed. The location is relative to the HOME directory for relative paths.
 ''')
 
 opt('shell_integration', 'inherit', long_text='''
 Control the shell integration on the remote host. See ref:`shell_integration`
 for details on how this setting works. The special value :code:`inherit` means
-use the setting from kitty.conf. This setting is  mainly useful for overriding
+use the setting from kitty.conf. This setting is useful for overriding
 integration on a per-host basis.''')
 
 opt('+env', '', option_type='env', add_to_default=False, long_text='''
@@ -47,6 +47,7 @@ environment variables can refer to each other, so if you use::
 The value of MYVAR2 will be :code:`a/<path to home directory>/b`. Using
 :code:`VAR=` will set it to the empty string and using just :code:`VAR`
 will delete the variable from the child process' environment. The definitions
-are processed alphabetically.
+are processed alphabetically. The special value :code:`_kitty_copy_env_var_`
+will cause the value of the variable to be copied from the local machine.
 ''')
 egr()  # }}}

kitty_tests/__init__.py

@@ -97,7 +97,7 @@ class Callbacks:
     def handle_remote_ssh(self, msg):
         from kittens.ssh.main import get_ssh_data
         if self.pty:
-            for line in get_ssh_data(msg):
+            for line in get_ssh_data(msg, {'*': self.pty.ssh_opts} if self.pty.ssh_opts else None):
                 self.pty.write_to_child(line)
 
     def handle_remote_echo(self, msg):
@@ -159,9 +159,9 @@ class BaseTest(TestCase):
         s = Screen(c, lines, cols, scrollback, cell_width, cell_height, 0, c)
         return s
 
-    def create_pty(self, argv, cols=80, lines=25, scrollback=100, cell_width=10, cell_height=20, options=None, cwd=None, env=None):
+    def create_pty(self, argv, cols=80, lines=25, scrollback=100, cell_width=10, cell_height=20, options=None, cwd=None, env=None, ssh_opts=None):
         self.set_options(options)
-        return PTY(argv, lines, cols, scrollback, cell_width, cell_height, cwd, env)
+        return PTY(argv, lines, cols, scrollback, cell_width, cell_height, cwd, env, ssh_opts)
 
     def assertEqualAttributes(self, c1, c2):
         x1, y1, c1.x, c1.y = c1.x, c1.y, 0, 0
@@ -174,7 +174,12 @@ class BaseTest(TestCase):
 
 class PTY:
 
-    def __init__(self, argv, rows=25, columns=80, scrollback=100, cell_width=10, cell_height=20, cwd=None, env=None):
+    def __init__(self, argv, rows=25, columns=80, scrollback=100, cell_width=10, cell_height=20, cwd=None, env=None, ssh_opts=None):
+        if ssh_opts:
+            from kittens.ssh.options.types import Options as SSHOptions
+            self.ssh_opts = SSHOptions(ssh_opts or {})
+        else:
+            self.ssh_opts = None
         if isinstance(argv, str):
             argv = shlex.split(argv)
         pid, self.master_fd = fork()

kitty_tests/ssh.py

@@ -6,9 +6,11 @@ import os
 import shlex
 import shutil
 import tempfile
+from functools import lru_cache
 
 from kittens.ssh.config import load_config, options_for_host
 from kittens.ssh.main import bootstrap_script, get_connection_data
+from kittens.ssh.options.utils import DELETE_ENV_VAR
 from kitty.constants import is_macos
 from kitty.fast_data_types import CURSOR_BEAM
 from kitty.options.utils import shell_integration
@@ -65,10 +67,25 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         self.ae(for_host('x', 'env a=').env, {'a': ''})
         self.ae(for_host('x', 'env a').env, {'a': '_delete_this_env_var_'})
 
+    @property
+    @lru_cache()
+    def all_possible_sh(self):
+        return tuple(sh for sh in ('dash', 'zsh', 'bash', 'posh', 'sh') if shutil.which(sh))
+
     def test_ssh_bootstrap_script(self):
+        # test setting env vars
+        with tempfile.TemporaryDirectory() as tdir:
+            pty = self.check_bootstrap(
+                'dash', tdir, extra_exec='env; exit 0', SHELL_INTEGRATION_VALUE='',
+                ssh_opts={'env': {
+                    'TSET': 'set-works',
+                    'COLORTERM': DELETE_ENV_VAR,
+                }}
+            )
+            pty.wait_till(lambda: 'TSET=set-works' in pty.screen_contents())
+            self.assertNotIn('COLORTERM', pty.screen_contents())
         # test handling of data in tty before tarfile is sent
-        all_possible_sh = tuple(sh for sh in ('dash', 'zsh', 'bash', 'posh', 'sh') if shutil.which(sh))
+        for sh in self.all_possible_sh:
-        for sh in all_possible_sh:
             with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
                 pty = self.check_bootstrap(
                     sh, tdir, extra_exec='echo "ld:$leading_data"; exit 0',
@@ -90,15 +107,15 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         import pwd
         expected_login_shell = pwd.getpwuid(os.geteuid()).pw_shell
         for m in methods:
-            for sh in all_possible_sh:
+            for sh in self.all_possible_sh:
                 with self.subTest(sh=sh, method=m), tempfile.TemporaryDirectory() as tdir:
                     pty = self.check_bootstrap(sh, tdir, extra_exec=f'{m}; echo "$login_shell"; exit 0', SHELL_INTEGRATION_VALUE='')
                     self.assertIn(expected_login_shell, pty.screen_contents())
 
         # check that shell integration works
         ok_login_shell = ''
-        for sh in all_possible_sh:
+        for sh in self.all_possible_sh:
-            for login_shell in {'fish', 'zsh', 'bash'} & set(all_possible_sh):
+            for login_shell in {'fish', 'zsh', 'bash'} & set(self.all_possible_sh):
                 if login_shell == 'bash' and not bash_ok():
                     continue
                 ok_login_shell = login_shell
@@ -110,7 +127,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
                 with tempfile.TemporaryDirectory() as tdir:
                     self.check_bootstrap('sh', tdir, ok_login_shell, val)
 
-    def check_bootstrap(self, sh, home_dir, login_shell='', SHELL_INTEGRATION_VALUE='enabled', extra_exec='', pre_data=''):
+    def check_bootstrap(self, sh, home_dir, login_shell='', SHELL_INTEGRATION_VALUE='enabled', extra_exec='', pre_data='', ssh_opts=None):
         script = bootstrap_script(
             EXEC_CMD=f'echo "UNTAR_DONE"; {extra_exec}',
             OVERRIDE_LOGIN_SHELL=login_shell,
@@ -121,7 +138,7 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         # prevent newuser-install from running
         open(os.path.join(home_dir, '.zshrc'), 'w').close()
         options = {'shell_integration': shell_integration(SHELL_INTEGRATION_VALUE or 'disabled')}
-        pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=home_dir, env=env, options=options)
+        pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=home_dir, env=env, options=options, ssh_opts=ssh_opts)
         if pre_data:
             pty.write_buf = pre_data.encode('utf-8')
         del script

shell-integration/ssh/bootstrap.sh

@@ -70,14 +70,7 @@ get_data() {
             die "$size"
             ;;
     esac
-    data_dir=$(read_record)
+    data_dir="$HOME/$(read_record)"
-    case "$data_dir" in 
-        ("/"*)
-            ;;
-        (*)
-            data_dir="$HOME/$data_dir"
-            ;;
-    esac
     # using dd with bs=1 is very slow on Linux, so use head 
     command head -c "$size" < /dev/tty | untar
     rc="$?";
@@ -95,7 +88,8 @@ fi
 if [ "$rc" != "0" ]; then die "Failed to extract data transmitted by ssh kitten over the TTY device"; fi
 [ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data, no kitty.terminfo found";
 shell_integration_dir="$data_dir/shell-integration"
-shell_integration_settings_file="$data_dir/settings/ksi_env_var"
+settings_dir="$data_dir/settings"
+env_var_file="$settings_dir/env-vars.sh"
 
 # export TERMINFO
 tname=".terminfo"
@@ -105,6 +99,9 @@ if [ -e "/usr/share/misc/terminfo.cdb" ]; then
 fi
 export TERMINFO="$HOME/$tname"
 
+# setup env vars
+. "$env_var_file"
+
 # compile terminfo for this system
 if [ -x "$(command -v tic)" ]; then
     tic_out=$(command tic -x -o "$HOME/$tname" "$HOME/.terminfo/kitty.terminfo" 2>&1)
@@ -198,13 +195,6 @@ else
 fi
 shell_name=$(basename $login_shell)
 
-# read the variable and remove all leading and trailing spaces and collapse multiple spaces using xargs
-if [  -f "$shell_integration_settings_file" ]; then 
-    export KITTY_SHELL_INTEGRATION="$(cat $shell_integration_settings_file | xargs echo)"
-else
-    unset KITTY_SHELL_INTEGRATION
-fi
-
 exec_bash_with_integration() {
     export ENV="$shell_integration_dir/bash/kitty.bash"
     export KITTY_BASH_INJECT="1"