From 2404eba11f1e0fc99b8f7c8c566184783d177ff4 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Mon, 7 Mar 2022 06:48:18 +0530 Subject: [PATCH] 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 | 24 +++++++++++++++++++----- kitty/window.py | 2 +- kitty_tests/__init__.py | 2 +- kitty_tests/ssh.py | 2 +- shell-integration/ssh/bootstrap.py | 3 ++- shell-integration/ssh/bootstrap.sh | 4 +--- 6 files changed, 25 insertions(+), 12 deletions(-) diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index 9378611d9..42813599f 100644 --- a/kittens/ssh/main.py +++ b/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 diff --git a/kitty/window.py b/kitty/window.py index 5540530f5..2c0c1ecfb 100644 --- a/kitty/window.py +++ b/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: diff --git a/kitty_tests/__init__.py b/kitty_tests/__init__.py index f344754cd..4e882a40b 100644 --- a/kitty_tests/__init__.py +++ b/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): diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py index 113a84a4d..5d40c24cb 100644 --- a/kitty_tests/ssh.py +++ b/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) diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py index 2c7f0a981..f170a2cec 100644 --- a/shell-integration/ssh/bootstrap.py +++ b/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): diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index fc43909d7..44a688c5c 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/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() {