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,19 +70,31 @@ def make_tarfile(ssh_opts: SSHOptions) -> bytes:
             return None
         return normalize_tarinfo(tarinfo)
 
+    from kitty.shell_integration import get_effective_ksi_env_var
+    if ssh_opts.shell_integration == 'inherit':
+        ksi = get_effective_ksi_env_var()
+    else:
+        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('/')
-        from kitty.shell_integration import get_effective_ksi_env_var
-        if ssh_opts.shell_integration == 'inherit':
-            ksi = get_effective_ksi_env_var()
-        else:
-            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)}))
+        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:
-        tf.write(pw.encode('utf-8'))
+    with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=cache_dir(), delete=False) as tf:
+        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)