From 68df13d3fe99b98eee909c7a08ed6f9ac676ea3f Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 7 Mar 2022 11:07:26 +0530
Subject: [PATCH] Fix recursive definitions in env vars not expanded

---
 kittens/ssh/main.py                | 12 +++++++++++-
 kitty_tests/ssh.py                 |  6 ++++--
 shell-integration/ssh/bootstrap.py |  8 ++++++--
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index 24e0f9a73..03525aa9a 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -33,11 +33,21 @@ from .options.types import Options as SSHOptions
 from .options.utils import DELETE_ENV_VAR
 
 
+# See https://www.gnu.org/software/bash/manual/html_node/Double-Quotes.html
+quote_pat = re.compile('([\\`"\n])')
+
+
+def quote_env_val(x: str) -> str:
+    x = quote_pat.sub(r'\\\1', x)
+    x = x.replace('$(', r'\$(')  # prevent execution with $()
+    return f'"{x}"'
+
+
 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)}')
+        lines.append(f'export {shlex.quote(k)}={quote_env_val(val)}')
 
     for k in sorted(env):
         v = env[k]
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index 5d40c24cb..aaf158a8d 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -139,16 +139,18 @@ copy --exclude */w.* d1
                 self.ae(len(glob.glob(f'{remote_home}/{tname}/*/xterm-kitty')), 2)
 
     def test_ssh_env_vars(self):
+        tset = '$A-$(echo no)-`echo no2` "something"'
         for sh in self.all_possible_sh:
             with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
                 pty = self.check_bootstrap(
                     sh, tdir, test_script='env; exit 0', SHELL_INTEGRATION_VALUE='',
                     ssh_opts={'env': {
-                        'TSET': 'set-works',
+                        'A': 'AAA',
+                        'TSET': tset,
                         'COLORTERM': DELETE_ENV_VAR,
                     }}
                 )
-                pty.wait_till(lambda: 'TSET=set-works' in pty.screen_contents())
+                pty.wait_till(lambda: 'TSET={}'.format(tset.replace('$A', 'AAA')) in pty.screen_contents())
                 self.assertNotIn('COLORTERM', pty.screen_contents())
 
     def test_ssh_leading_data(self):
diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py
index f170a2cec..df70acad5 100644
--- a/shell-integration/ssh/bootstrap.py
+++ b/shell-integration/ssh/bootstrap.py
@@ -8,8 +8,8 @@ import getpass
 import io
 import os
 import pwd
+import re
 import select
-import shlex
 import shutil
 import subprocess
 import sys
@@ -72,6 +72,10 @@ def debug(msg):
         write_all(tty_fd, data)
 
 
+def unquote_env_val(x):
+    return re.sub('\\\\([\\$`"\n])', r'\1', x[1:-1])
+
+
 def apply_env_vars(raw):
     global login_shell
 
@@ -81,7 +85,7 @@ def apply_env_vars(raw):
             key, val = parts[0], ''
         else:
             key, val = parts
-            val = shlex.split(val)[0]
+            val = os.path.expandvars(unquote_env_val(val))
         os.environ[key] = val
 
     for line in raw.splitlines():