more work on ssh data transmission

kittens/ssh/main.py

@@ -1,25 +1,28 @@
 #!/usr/bin/env python3
 # License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
 
+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

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);
+#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();
-#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
 
+            } 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]);
             }

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];

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*);

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()

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()

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