ssh kitten: Passthrough to ssh if run outside of kitty

Also, ensure that the ssh data request is only served if it is received
over the tty of the correct kitty window.

kittens/ssh/main.py

@@ -120,7 +120,7 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
     return buf.getvalue()
 
 
-def get_ssh_data(msg: str) -> Iterator[bytes]:
+def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
     record_sep = b'\036'
 
     def fmt_prefix(msg: Any) -> bytes:
@@ -134,6 +134,7 @@ def get_ssh_data(msg: str) -> Iterator[bytes]:
         pw = md['pw']
         pwfilename = md['pwfile']
         username = md['user']
+        rq_id = md['id']
     except Exception:
         traceback.print_exc()
         yield fmt_prefix('!invalid ssh data request message')
@@ -144,9 +145,11 @@ def get_ssh_data(msg: str) -> Iterator[bytes]:
                 env_data = json.load(f)
                 if pw != env_data['pw']:
                     raise ValueError('Incorrect password')
-        except Exception:
+                if rq_id != request_id:
+                    raise ValueError('Incorrect request id')
+        except Exception as e:
             traceback.print_exc()
-            yield fmt_prefix('!incorrect ssh data password')
+            yield fmt_prefix(f'!{e}')
         else:
             ssh_opts = {k: SSHOptions(v) for k, v in env_data['opts'].items()}
             resolved_ssh_opts = options_for_host(hostname, username, ssh_opts)
@@ -187,7 +190,13 @@ def prepare_exec_cmd(remote_args: Sequence[str], is_python: bool) -> str:
     return f"""exec "$login_shell" -c '{args}'"""
 
 
-def bootstrap_script(script_type: str = 'sh', remote_args: Sequence[str] = (), ssh_opts_dict: Dict[str, Dict[str, Any]] = {}, test_script: str = '') -> str:
+def bootstrap_script(
+    script_type: str = 'sh', remote_args: Sequence[str] = (),
+    ssh_opts_dict: Dict[str, Dict[str, Any]] = {},
+    test_script: str = '', request_id: Optional[str] = None
+) -> str:
+    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 ''
     with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f:
         ans = f.read()
@@ -196,7 +205,10 @@ def bootstrap_script(script_type: str = 'sh', remote_args: Sequence[str] = (), s
         data = {'pw': pw, 'env': dict(os.environ), 'opts': ssh_opts_dict}
         tf.write(json.dumps(data).encode('utf-8'))
     atexit.register(safe_remove, tf.name)
-    replacements = {'DATA_PASSWORD': pw, 'PASSWORD_FILENAME': os.path.basename(tf.name), 'EXEC_CMD': exec_cmd, 'TEST_SCRIPT': test_script}
+    replacements = {
+        'DATA_PASSWORD': pw, 'PASSWORD_FILENAME': os.path.basename(tf.name), 'EXEC_CMD': exec_cmd, 'TEST_SCRIPT': test_script,
+        'REQUEST_ID': request_id
+    }
     return prepare_script(ans, replacements)
 
 
@@ -385,6 +397,8 @@ def main(args: List[str]) -> NoReturn:
         ssh_args, server_args, passthrough, found_extra_args = parse_ssh_args(args, extra_args=('--kitten',))
     except InvalidSSHArgs as e:
         e.system_exit()
+    if not os.environ.get('KITTY_WINDOW_ID'):
+        passthrough = True
     cmd = ['ssh'] + ssh_args
     if passthrough:
         cmd += server_args

kitty/window.py

@@ -882,7 +882,7 @@ class Window:
 
     def handle_remote_ssh(self, msg: str) -> None:
         from kittens.ssh.main import get_ssh_data
-        for line in get_ssh_data(msg):
+        for line in get_ssh_data(msg, f'{os.getpid()}-{self.id}'):
             self.write_to_child(line)
 
     def handle_remote_print(self, msg: str) -> None:

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, "testing"):
                 self.pty.write_to_child(line)
 
     def handle_remote_echo(self, msg):

kitty_tests/ssh.py

@@ -224,7 +224,7 @@ copy --exclude */w.* d1
         else:
             test_script = f'echo "UNTAR_DONE"; {test_script}'
         script = bootstrap_script(
-            script_type='py' if 'python' in sh else 'sh',
+            script_type='py' if 'python' in sh else 'sh', request_id="testing",
             test_script=test_script, ssh_opts_dict={'*': ssh_opts},
         )
         env = basic_shell_env(home_dir)

shell-integration/ssh/bootstrap.py

@@ -59,7 +59,8 @@ def dcs_to_kitty(type, payload):
 
 def send_data_request():
     hostname = os.environ.get('HOSTNAME') or os.uname().nodename
-    write_all(tty_fd, dcs_to_kitty('ssh', 'hostname={}:pwfile=PASSWORD_FILENAME:user={}:pw=DATA_PASSWORD'.format(hostname, getpass.getuser())))
+    write_all(tty_fd, dcs_to_kitty(
+        'ssh', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:pw=DATA_PASSWORD'.format(hostname, getpass.getuser())))
 
 
 def debug(msg):

shell-integration/ssh/bootstrap.sh

@@ -93,12 +93,10 @@ if [ -z "$HOME" ]; then HOME=~; fi
 if [ -z "$USER" ]; then USER=$(command whoami 2> /dev/null); fi
 
 # ask for the SSH data
-data_password="DATA_PASSWORD"
-password_filename="PASSWORD_FILENAME"
 leading_data=""
 
 init_tty && trap 'cleanup_on_bootstrap_exit' EXIT
-[ "$tty_ok" = "y" ] && dcs_to_kitty "ssh" "hostname="$hostname":pwfile="$password_filename":user="$USER":pw="$data_password""
+[ "$tty_ok" = "y" ] && dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":pw="DATA_PASSWORD""
 record_separator=$(printf "\036")
 
 mv_files_and_dirs() {