Allow controlling where on the remote computer the ssh kitten installs its data

kittens/ssh/main.py

@@ -19,14 +19,14 @@
 from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
 from kitty.shell_integration import get_effective_ksi_env_var
 from kitty.short_uuid import uuid4
+from kitty.types import run_once
 from kitty.utils import SSHConnectionData
 
 from .completion import complete, ssh_options
-
-DEFAULT_SHELL_INTEGRATION_DEST = '.local/share/kitty-ssh-kitten/shell-integration'
+from .options.types import Options as SSHOptions
 
 
-def make_tarfile(hostname: str = '', shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> bytes:
+def make_tarfile(ssh_opts: SSHOptions) -> bytes:
 
     def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo:
         tarinfo.uname = tarinfo.gname = 'kitty'
@@ -51,17 +51,32 @@ def filter_files(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]:
 
     buf = io.BytesIO()
     with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
-        tf.add(shell_integration_dir, arcname=shell_integration_dest, filter=filter_files)
+        rd = ssh_opts.remote_dir.rstrip('/')
+        ksi = get_effective_ksi_env_var()
+        if ksi:
+            tf.add(shell_integration_dir, arcname=rd + '/shell-integration', filter=filter_files)
+            add_data_as_file(tf, rd + '/settings/ksi_env_var', ksi)
         tf.add(terminfo_dir, arcname='.terminfo', filter=filter_files)
-        add_data_as_file(tf, shell_integration_dest.rstrip('/') + '/settings/ksi_env_var', get_effective_ksi_env_var())
     return buf.getvalue()
 
 
-def get_ssh_data(msg: str, shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> Iterator[bytes]:
+@run_once
+def load_ssh_options() -> Dict[str, SSHOptions]:
+    from .config import init_config
+    return init_config()
+
+
+def get_ssh_data(msg: str, ssh_opts: Optional[Dict[str, SSHOptions]] = None) -> Iterator[bytes]:
+    from .config import options_for_host
+    record_sep = b'\036'
+
+    if ssh_opts is None:
+        ssh_opts = load_ssh_options()
 
     def fmt_prefix(msg: Any) -> bytes:
-        return f'\036{msg}:'.encode('ascii')
+        return str(msg).encode('ascii') + record_sep
 
+    yield record_sep  # to discard leading data
     try:
         msg = standard_b64decode(msg).decode('utf-8')
         md = dict(x.split('=', 1) for x in msg.split(':'))
@@ -70,23 +85,26 @@ def fmt_prefix(msg: Any) -> bytes:
         pwfilename = md['pwfile']
     except Exception:
         yield fmt_prefix('!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 fmt_prefix('!incorrect ssh data password')
     else:
         try:
-            data = make_tarfile(hostname, shell_integration_dest)
+            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 fmt_prefix('!error while gathering ssh data')
+            yield fmt_prefix('!incorrect ssh data password')
         else:
-            from base64 import standard_b64encode
-            encoded_data = standard_b64encode(data)
-            yield fmt_prefix(len(encoded_data))
-            yield encoded_data
+            resolved_ssh_opts = options_for_host(hostname, ssh_opts)
+            try:
+                data = make_tarfile(resolved_ssh_opts)
+            except Exception:
+                yield fmt_prefix('!error while gathering ssh data')
+            else:
+                from base64 import standard_b64encode
+                encoded_data = standard_b64encode(data)
+                yield fmt_prefix(len(encoded_data))
+                yield fmt_prefix(resolved_ssh_opts.remote_dir)
+                yield encoded_data
 
 
 def safe_remove(x: str) -> None:
@@ -103,7 +121,6 @@ def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
     replacements['PASSWORD_FILENAME'] = os.path.basename(tf.name)
     for k in ('EXEC_CMD', 'OVERRIDE_LOGIN_SHELL'):
         replacements[k] = replacements.get(k, '')
-    replacements['SHELL_INTEGRATION_DIR'] = replacements.get('SHELL_INTEGRATION_DIR', DEFAULT_SHELL_INTEGRATION_DEST)
 
     def sub(m: 're.Match[str]') -> str:
         return replacements[m.group()]

kittens/ssh/options/definition.py

@@ -27,10 +27,13 @@
 '''
     )
 
-opt('+env', '',
-    option_type='env',
-    add_to_default=False,
-    long_text='''
+opt('remote_dir', '.local/share/kitty-ssh-kitten', long_text='''
+The location on the remote computer where the files needed for this kitten
+are installed. The location is relative to the HOME directory.
+'''
+    )
+
+opt('+env', '', option_type='env', add_to_default=False, long_text='''
 Specify environment variables to set on the remote host. Note that
 environment variables can refer to each other, so if you use::
 
@@ -41,8 +44,5 @@
 :code:`VAR=` will set it to the empty string and using just :code:`VAR`
 will delete the variable from the child process' environment. The definitions
 are processed alphabetically.
-'''
-    )
-
-
+''')
 egr()  # }}}

kittens/ssh/options/parse.py

@@ -14,6 +14,9 @@ def env(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
     def hostname(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
         hostname(val, ans)
 
+    def remote_dir(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
+        ans['remote_dir'] = str(val)
+
 
 def create_result_dict() -> typing.Dict[str, typing.Any]:
     return {

kittens/ssh/options/types.py

@@ -4,11 +4,12 @@
 
 
 option_names = (  # {{{
- 'env', 'hostname')  # }}}
+ 'env', 'hostname', 'remote_dir')  # }}}
 
 
 class Options:
     hostname: str = '*'
+    remote_dir: str = '.local/share/kitty-ssh-kitten'
     env: typing.Dict[str, str] = {}
     config_paths: typing.Tuple[str, ...] = ()
     config_overrides: typing.Tuple[str, ...] = ()

shell-integration/ssh/bootstrap.sh

@@ -36,12 +36,13 @@ leading_data=""
 
 dsc_to_kitty "ssh" "hostname=$hostname:pwfile=$password_filename:pw=$data_password"
 size=""
+record_separator=$(printf "\036")
 
 untar() {
     command base64 -d | command tar xjf - --no-same-owner -C "$HOME" 
 }
 
-get_data() {
+read_record() {
     # We need a way to read a single byte at a time and to read a specified number of bytes in one invocation.
     # The options are head -c, read -N and dd
     #
@@ -52,45 +53,49 @@ get_data() {
     #
     # POSIX dd works for one byte at a time but for reading X bytes it needs the GNU iflag=count_bytes
     # extension, and is anyway unsafe as it can lead to corrupt output when the read syscall is interrupted.
-    record_started=0
-    record_separator=$(printf "\036")
+    record=""
     while :; do
-        n=$(command dd bs=1 count=1 2> /dev/null) 
-        if [ $record_started = 1 ]; then
-            if [ "$n" = ":" ]; then break; fi
-            size="$size$n"
-        else
-            if [ "$n" = "$record_separator" ]; then 
-                record_started=1; 
-            else
-                leading_data="$leading_data$n"
-            fi
-        fi
+        n=$(command dd bs=1 count=1 2> /dev/null < /dev/tty) 
+        [ "$n" = "$record_separator" ] && break
+        record="$record$n"
     done
+    printf "%s" "$record"
+}
+
+get_data() {
+    leading_data=$(read_record)
+    size=$(read_record)
     case "$size" in
         ("!"*)
             die "$size"
             ;;
     esac
+    data_dir=$(read_record)
+    case "$data_dir" in 
+        ("/"*)
+            ;;
+        (*)
+            data_dir="$HOME/$data_dir"
+            ;;
+    esac
     # using dd with bs=1 is very slow on Linux, so use head 
-    command head -c "$size" | untar
-    # command dd bs=1 "count=$size" 2> /dev/null | untar
+    command head -c "$size" < /dev/tty | untar
     rc="$?";
 }
 
 get_data
 command stty "$saved_tty_settings"
 saved_tty_settings=""
-if [ "$rc" != "0" ]; then die "Failed to extract data transmitted by ssh kitten over the TTY device"; fi
-shell_integration_dir="$HOME/SHELL_INTEGRATION_DIR"
-shell_integration_settings_file="$shell_integration_dir/settings/ksi_env_var"
-if [ ! -f "$shell_integration_settings_file" ]; then die "Extracted data transmitted by ssh kitten is incomplete"; fi
 if [ -n "$leading_data" ]; then
     # clear current line as it might have things echoed on it from leading_data
     # because we only turn off echo in this script whereas the leading bytes could 
     # have been sent before the script had a chance to run
     printf "\r\033[K"  
 fi
+if [ "$rc" != "0" ]; then die "Failed to extract data transmitted by ssh kitten over the TTY device"; fi
+[ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data, no kitty.terminfo found";
+shell_integration_dir="$data_dir/shell-integration"
+shell_integration_settings_file="$data_dir/settings/ksi_env_var"
 
 # export TERMINFO
 tname=".terminfo"
@@ -194,7 +199,11 @@ fi
 shell_name=$(basename $login_shell)
 
 # read the variable and remove all leading and trailing spaces and collapse multiple spaces using xargs
-export KITTY_SHELL_INTEGRATION="$(cat $shell_integration_settings_file | xargs echo)"
+if [  -f "$shell_integration_settings_file" ]; then 
+    export KITTY_SHELL_INTEGRATION="$(cat $shell_integration_settings_file | xargs echo)"
+else
+    unset KITTY_SHELL_INTEGRATION
+fi
 
 exec_bash_with_integration() {
     export ENV="$shell_integration_dir/bash/kitty.bash"