diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index bc946c885..6f8f5288d 100644 --- a/kittens/ssh/main.py +++ b/kittens/ssh/main.py @@ -1,25 +1,28 @@ #!/usr/bin/env python3 # License: GPL v3 Copyright: 2018, Kovid Goyal +import atexit import io import os import shlex import subprocess import sys import tarfile +import tempfile from contextlib import suppress -from typing import List, NoReturn, Optional, Set, Tuple +from typing import Iterator, List, NoReturn, Optional, Set, Tuple -from kitty.constants import shell_integration_dir, terminfo_dir +from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir +from kitty.short_uuid import uuid4_for_escape_code from kitty.utils import SSHConnectionData from .completion import complete, ssh_options -def make_tarfile() -> bytes: +def make_tarfile(hostname: str = '') -> bytes: def filter_files(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]: - if tarinfo.name.endswith('ssh/bootstrap.sh'): + if tarinfo.name.endswith('ssh/bootstrap.sh') or tarinfo.name.endswith('ssh/bootstrap.py'): return None tarinfo.uname = tarinfo.gname = 'kitty' tarinfo.uid = tarinfo.gid = 0 @@ -32,6 +35,55 @@ def make_tarfile() -> bytes: return buf.getvalue() +def get_ssh_data(msg: str) -> Iterator[bytes]: + yield b"KITTY_SSH_DATA_START" + try: + hostname, pwfilename, pw = msg.split(':', 2) + except Exception: + yield b' invalid ssh data request message' + try: + with open(os.path.join(cache_dir(), pwfilename)) as f: + os.unlink(f.name) + if pw != f.read(): + raise ValueError('Incorrect password') + except Exception: + yield b' incorrect ssh data password' + else: + try: + data = make_tarfile(hostname) + except Exception: + yield b' error while gathering ssh data' + else: + from base64 import standard_b64encode + encoded_data = memoryview(standard_b64encode(data)) + while encoded_data: + yield encoded_data[:1024] + encoded_data = encoded_data[1024:] + yield b"KITTY_SSH_DATA_END" + + +def safe_remove(x: str) -> None: + with suppress(OSError): + os.remove(x) + + +def prepare_script(ans: str, EXEC_CMD: str = '') -> str: + ans = ans.replace('EXEC_CMD', EXEC_CMD, 1) + pw = uuid4_for_escape_code() + with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', dir=cache_dir(), delete=False) as tf: + tf.write(pw.encode('utf-8')) + atexit.register(safe_remove, tf.name) + ans = ans.replace('DATA_PASSWORD', pw, 1) + ans = ans.replace("PASSWORD_FILENAME", os.path.basename(tf.name)) + return ans + + +def bootstrap_script(EXEC_CMD: str = '', script_type: str = 'sh') -> str: + with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f: + ans = f.read() + return prepare_script(ans, EXEC_CMD) + + SHELL_SCRIPT = '''\ #!/bin/sh # macOS ships with an ancient version of tic that cannot read from stdin, so we diff --git a/kitty/parser.c b/kitty/parser.c index a5fa03544..f80f0fd1d 100644 --- a/kitty/parser.c +++ b/kitty/parser.c @@ -1071,27 +1071,20 @@ dispatch_dcs(Screen *screen, PyObject DUMP_UNUSED *dump_callback) { Py_DECREF(cmd); } else PyErr_Clear(); #undef CMD_PREFIX -#define PRINT_PREFIX "kitty-print|" - } else if (startswith(screen->parser_buf + 1, screen->parser_buf_pos - 1, PRINT_PREFIX)) { - const size_t pp_size = sizeof(PRINT_PREFIX); - PyObject *msg = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, screen->parser_buf + pp_size, screen->parser_buf_pos - pp_size); - if (msg != NULL) { - REPORT_OSC2(screen_handle_print, (char)screen->parser_buf[0], msg); - screen_handle_print(screen, msg); - Py_DECREF(msg); - } else PyErr_Clear(); -#undef PRINT_PREFIX -#define ECHO_PREFIX "kitty-echo|" - } else if (startswith(screen->parser_buf + 1, screen->parser_buf_pos - 1, ECHO_PREFIX)) { - const size_t pp_size = sizeof(ECHO_PREFIX); - PyObject *msg = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, screen->parser_buf + pp_size, screen->parser_buf_pos - pp_size); - if (msg != NULL) { - REPORT_OSC2(screen_handle_echo, (char)screen->parser_buf[0], msg); - screen_handle_echo(screen, msg); - Py_DECREF(msg); - } else PyErr_Clear(); -#undef ECHO_PREFIX +#define IF_SIMPLE_PREFIX(prefix, func) \ + if (startswith(screen->parser_buf + 1, screen->parser_buf_pos - 1, prefix)) { \ + const size_t pp_size = sizeof(prefix); \ + PyObject *msg = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, screen->parser_buf + pp_size, screen->parser_buf_pos - pp_size); \ + if (msg != NULL) { \ + REPORT_OSC2(func, (char)screen->parser_buf[0], msg); \ + func(screen, msg); \ + Py_DECREF(msg); \ + } else PyErr_Clear(); + } else IF_SIMPLE_PREFIX("kitty-print|", screen_handle_print) + } else IF_SIMPLE_PREFIX("kitty-echo|", screen_handle_echo) + } else IF_SIMPLE_PREFIX("kitty-ssh|", screen_handle_ssh) +#undef IF_SIMPLE_PREFIX } else { REPORT_ERROR("Unrecognized DCS @ code: 0x%x", screen->parser_buf[1]); } diff --git a/kitty/screen.c b/kitty/screen.c index bee895b61..b5fcbfcb8 100644 --- a/kitty/screen.c +++ b/kitty/screen.c @@ -2109,6 +2109,11 @@ screen_handle_echo(Screen *self, PyObject *msg) { CALLBACK("handle_remote_echo", "O", msg); } +void +screen_handle_ssh(Screen *self, PyObject *msg) { + CALLBACK("handle_remote_ssh", "O", msg); +} + void screen_request_capabilities(Screen *self, char c, PyObject *q) { static char buf[128]; diff --git a/kitty/screen.h b/kitty/screen.h index baca3e7f8..736651c68 100644 --- a/kitty/screen.h +++ b/kitty/screen.h @@ -209,6 +209,7 @@ void screen_pop_colors(Screen *, unsigned int); void screen_report_color_stack(Screen *); void screen_handle_print(Screen *, PyObject *cmd); void screen_handle_echo(Screen *, PyObject *cmd); +void screen_handle_ssh(Screen *, PyObject *cmd); void screen_designate_charset(Screen *, uint32_t which, uint32_t as); void screen_use_latin1(Screen *, bool); void set_title(Screen *self, PyObject*); diff --git a/kitty/short_uuid.py b/kitty/short_uuid.py index 6015d4702..3f5ed8ce7 100644 --- a/kitty/short_uuid.py +++ b/kitty/short_uuid.py @@ -55,3 +55,11 @@ _global_instance = ShortUUID() uuid4 = _global_instance.uuid4 uuid5 = _global_instance.uuid5 decode = _global_instance.decode +_escape_code_instance: Optional[ShortUUID] = None + + +def uuid4_for_escape_code() -> str: + global _escape_code_instance + if _escape_code_instance is None: + _escape_code_instance = ShortUUID(escape_code_safe_alphabet) + return _escape_code_instance.uuid4() diff --git a/kitty/window.py b/kitty/window.py index a268979b5..cced08484 100644 --- a/kitty/window.py +++ b/kitty/window.py @@ -321,7 +321,7 @@ def cmd_output(screen: Screen, which: CommandOutput = CommandOutput.last_run, as return ''.join(lines) -def process_remote_print(msg: bytes) -> str: +def process_remote_print(msg: str) -> str: from base64 import standard_b64decode from .cli import green text = standard_b64decode(msg).decode('utf-8', 'replace') @@ -875,12 +875,18 @@ class Window: def handle_remote_cmd(self, cmd: str) -> None: get_boss().handle_remote_cmd(cmd, self) - def handle_remote_echo(self, msg: bytes) -> None: + def handle_remote_echo(self, msg: str) -> None: from base64 import standard_b64decode data = standard_b64decode(msg) self.write_to_child(data) - def handle_remote_print(self, msg: bytes) -> None: + def handle_remote_ssh(self, msg: str) -> None: + from kittens.ssh.main import get_ssh_data + for line in get_ssh_data(msg): + self.write_to_child(line) + self.write_to_child('\n') + + def handle_remote_print(self, msg: str) -> None: text = process_remote_print(msg) print(text, end='', file=sys.stderr) sys.stderr.flush() diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index a4d58029c..b7ed3bc11 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/shell-integration/ssh/bootstrap.sh @@ -22,10 +22,15 @@ if [[ -z "$HOSTNAME" ]]; then else hostname=$(HOSTNAME) fi +# ensure $HOME is set +if [[ -z "$HOME" ]]; then HOME=~; fi +# ensure $USER is set +if [[ -z "$USER" ]]; then USER=$(whoami); fi + # ask for the SSH data data_password="DATA_PASSWORD" password_filename="PASSWORD_FILENAME" -printf "\eP@kitty-ssh|$password_filename:$data_password:$hostname\e\\" +printf "\eP@kitty-ssh|$hostname:$password_filename:$data_password\e\\" while IFS= read -r line; do case "$line" in @@ -52,28 +57,26 @@ saved_tty_settings="" if [[ ! -z "$pending_data" ]]; then printf "\eP@kitty-echo|$(echo -n "$pending_data" | base64)\e\\" fi -command base64 -d < "$encoded_data_file" | command tar xjf - --no-same-owner +command base64 -d < "$encoded_data_file" | command tar xjf - --no-same-owner -C "$HOME" rc=$? command rm -f "$encoded_data_file" encoded_data_file="" -if [ "$rc" != "0" ]; then echo "Failed to extract data transmitted by ssh kitten over the TTY device" > /dev/stderr; exit 1; fi +if [[ "$rc" != "0" ]]; then echo "Failed to extract data transmitted by ssh kitten over the TTY device" > /dev/stderr; exit 1; fi # compile terminfo for this system -if [ -x "$(command -v tic)" ]; then +if [[ -x "$(command -v tic)" ]]; then tname=".terminfo" - if [ -e "/usr/share/misc/terminfo.cdb" ]; then + if [[ -e "/usr/share/misc/terminfo.cdb" ]]; then # NetBSD requires this see https://github.com/kovidgoyal/kitty/issues/4622 tname=".terminfo.cdb" fi tic_out=$(command tic -x -o "$HOME/$tname" ".terminfo/kitty.terminfo" 2>&1) rc=$? - if [ "$rc" != "0" ]; then echo "$tic_out"; exit 1; fi + if [[ "$rc" != "0" ]]; then echo "$tic_out"; exit 1; fi export TERMINFO="$HOME/$tname" fi # If a command was passed to SSH execute it here EXEC_CMD -# ensure $USER is set -if [ -z "$USER" ]; then export USER=$(whoami); fi