diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index 73c38197a..438c8e88e 100644 --- a/kittens/ssh/main.py +++ b/kittens/ssh/main.py @@ -19,14 +19,14 @@ from typing import ( 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 make_tarfile(hostname: str = '', shell_integration_dest: str = DEFAULT_SHELL 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 get_ssh_data(msg: str, shell_integration_dest: str = DEFAULT_SHELL_INTEGRATI 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()] diff --git a/kittens/ssh/options/definition.py b/kittens/ssh/options/definition.py index 49627eaa8..08d44bab0 100644 --- a/kittens/ssh/options/definition.py +++ b/kittens/ssh/options/definition.py @@ -27,10 +27,13 @@ to SSH to connect to it. ''' ) -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 @@ The value of MYVAR2 will be :code:`a//b`. Using :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() # }}} diff --git a/kittens/ssh/options/parse.py b/kittens/ssh/options/parse.py index ffb003d28..aba3c82d2 100644 --- a/kittens/ssh/options/parse.py +++ b/kittens/ssh/options/parse.py @@ -14,6 +14,9 @@ class Parser: 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 { diff --git a/kittens/ssh/options/types.py b/kittens/ssh/options/types.py index 7f05263c7..6b92a65a7 100644 --- a/kittens/ssh/options/types.py +++ b/kittens/ssh/options/types.py @@ -4,11 +4,12 @@ import typing 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, ...] = () diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index 3f4065a64..4a35bd325 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/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"