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 filter(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]:
     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 fmt_prefix(msg: Any) -> 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 fmt_prefix(msg: Any) -> 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 @@ def handle_remote_echo(self, msg: str) -> None:
 
     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 @@ def handle_remote_print(self, msg):
     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 @@ def check_bootstrap(self, sh, home_dir, login_shell='', SHELL_INTEGRATION_VALUE=
         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() {