From b2e74e4830f1fe9a78fe66cef1b554a2c562d18f Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sat, 12 Mar 2022 15:14:16 +0530 Subject: [PATCH] Switch to using the hostname from the CLI for settings This removes the need to wait for data from the remote machine before sending data to it. --- kittens/ssh/main.py | 49 ++++++++++++------------------ kittens/ssh/options/definition.py | 9 ++---- kitty_tests/ssh.py | 10 +++--- shell-integration/ssh/bootstrap.py | 4 +-- shell-integration/ssh/bootstrap.sh | 13 ++------ 5 files changed, 30 insertions(+), 55 deletions(-) diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index 1e7948b4f..97019434f 100644 --- a/kittens/ssh/main.py +++ b/kittens/ssh/main.py @@ -28,7 +28,6 @@ from kitty.constants import ( runtime_dir, shell_integration_dir, ssh_control_master_template, terminfo_dir ) -from kitty.fast_data_types import get_options from kitty.shm import SharedMemory from kitty.utils import SSHConnectionData @@ -99,14 +98,15 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: st from kitty.shell_integration import get_effective_ksi_env_var if ssh_opts.shell_integration == 'inherited': - ksi = get_effective_ksi_env_var() + from kitty.cli import create_default_opts + ksi = get_effective_ksi_env_var(create_default_opts()) else: from kitty.options.types import Options from kitty.options.utils import shell_integration ksi = get_effective_ksi_env_var(Options({'shell_integration': shell_integration(ssh_opts.shell_integration)})) env = { - 'TERM': get_options().term, + 'TERM': os.environ['TERM'], 'COLORTERM': 'truecolor', } for q in ('KITTY_WINDOW_ID', 'WINDOWID'): @@ -148,12 +148,9 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]: try: msg = standard_b64decode(msg).decode('utf-8') md = dict(x.split('=', 1) for x in msg.split(':')) - hostname = md['hostname'] pw = md['pw'] pwfilename = md['pwfile'] - username = md['user'] rq_id = md['id'] - compression = md['compression'] except Exception: traceback.print_exc() yield fmt_prefix('!invalid ssh data request message') @@ -171,24 +168,15 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]: raise ValueError('Incorrect password') if rq_id != request_id: raise ValueError('Incorrect request id') - cli_hostname = env_data['cli_hostname'] - cli_uname = env_data['cli_uname'] except Exception as e: traceback.print_exc() 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, cli_hostname, cli_uname) - resolved_ssh_opts.copy = {k: CopyInstruction(*v) for k, v in resolved_ssh_opts.copy.items()} - try: - data = make_tarfile(resolved_ssh_opts, env_data['env'], compression) - except Exception: - traceback.print_exc() - yield fmt_prefix('!error while gathering ssh data') - else: - encoded_data = standard_b64encode(data) - yield fmt_prefix(len(encoded_data)) - yield encoded_data + ssh_opts = SSHOptions(env_data['opts']) + ssh_opts.copy = {k: CopyInstruction(*v) for k, v in ssh_opts.copy.items()} + encoded_data = env_data['tarfile'].encode('ascii') + yield fmt_prefix(len(encoded_data)) + yield encoded_data def safe_remove(x: str) -> None: @@ -217,8 +205,7 @@ def prepare_exec_cmd(remote_args: Sequence[str], is_python: bool) -> str: def bootstrap_script( - script_type: str = 'sh', remote_args: Sequence[str] = (), - ssh_opts_dict: Dict[str, Dict[str, Any]] = {}, + ssh_opts: SSHOptions, script_type: str = 'sh', remote_args: Sequence[str] = (), test_script: str = '', request_id: Optional[str] = None, cli_hostname: str = '', cli_uname: str = '', request_data: str = '1', ) -> Tuple[str, Dict[str, str], SharedMemory]: @@ -228,7 +215,8 @@ def bootstrap_script( with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f: ans = f.read() pw = secrets.token_hex() - data = {'pw': pw, 'env': dict(os.environ), 'opts': ssh_opts_dict, 'cli_hostname': cli_hostname, 'cli_uname': cli_uname} + tfd = standard_b64encode(make_tarfile(ssh_opts, dict(os.environ), 'gz' if script_type == 'sh' else 'bz2')).decode('ascii') + data = {'pw': pw, 'opts': ssh_opts._asdict(), 'hostname': cli_hostname, 'uname': cli_uname, 'tarfile': tfd} db = json.dumps(data) with SharedMemory(size=len(db) + SharedMemory.num_bytes_for_size, mode=stat.S_IREAD, prefix=f'kssh-{os.getpid()}-') as shm: shm.write_data_with_size(db) @@ -431,14 +419,15 @@ def wrap_bootstrap_script(sh_script: str, interpreter: str) -> List[str]: def get_remote_command( - remote_args: List[str], hostname: str = 'localhost', cli_hostname: str = '', cli_uname: str = '', - interpreter: str = 'sh', - ssh_opts_dict: Dict[str, Dict[str, Any]] = {} + remote_args: List[str], + ssh_opts: SSHOptions, + hostname: str = 'localhost', cli_hostname: str = '', cli_uname: str = '', ) -> Tuple[List[str], Dict[str, str]]: + interpreter = ssh_opts.interpreter q = os.path.basename(interpreter).lower() is_python = 'python' in q sh_script, replacements, shm = bootstrap_script( - script_type='py' if is_python else 'sh', remote_args=remote_args, ssh_opts_dict=ssh_opts_dict, + ssh_opts, script_type='py' if is_python else 'sh', remote_args=remote_args, cli_hostname=cli_hostname, cli_uname=cli_uname) return wrap_bootstrap_script(sh_script, interpreter), replacements @@ -528,10 +517,10 @@ def main(args: List[str]) -> NoReturn: if overrides: overrides.insert(0, f'hostname {uname}@{hostname_for_match}') so = init_config(overrides) - sod = {k: v._asdict() for k, v in so.items()} host_opts = options_for_host(hostname_for_match, uname, so) - use_control_master = 'KITTY_PID' in os.environ and host_opts.share_connections - rcmd, replacements = get_remote_command(remote_args, hostname, hostname_for_match, uname, host_opts.interpreter, sod) + running_in_kitty = 'KITTY_PID' in os.environ + use_control_master = running_in_kitty and host_opts.share_connections + rcmd, replacements = get_remote_command(remote_args, host_opts, hostname, hostname_for_match, uname) cmd += rcmd if use_control_master: cmd[insertion_point:insertion_point] = connection_sharing_args(host_opts, int(os.environ['KITTY_PID'])) diff --git a/kittens/ssh/options/definition.py b/kittens/ssh/options/definition.py index 3373d7c2d..67320f7ba 100644 --- a/kittens/ssh/options/definition.py +++ b/kittens/ssh/options/definition.py @@ -30,8 +30,7 @@ hosts can be used. Multiple hostnames can also be specified separated by spaces. The hostname can include an optional username in the form :code:`user@host`. When not specified options apply to all hosts, until the first hostname specification is found. Note that matching of hostname is done against -both the hostname used by the remote computer, and the name you pass -to SSH to connect to it. If either matches, it is considered a match. +the name you specify on the command line to connect to the remote computer. If you wish to include the same basic configuration for many different hosts, you can do so with the :ref:`include ` directive. ''') @@ -89,9 +88,7 @@ these are automatically cleaned up by kitty when it quits. opt('interpreter', 'sh', long_text=''' The interpreter to use on the remote host. Must be either a POSIX complaint shell or a python executable. If the default sh is not available for broken, using -an alternate interpreter can be useful. Note that as the interpreter is used for -bootstrapping, hostname specific values are matched again the hostname from the -command line args only. +an alternate interpreter can be useful. ''') opt('remote_dir', '.local/share/kitty-ssh-kitten', option_type='relative_dir', long_text=''' @@ -103,7 +100,7 @@ that resolve to a location outside the HOME are not allowed. opt('shell_integration', 'inherited', long_text=''' Control the shell integration on the remote host. See :ref:`shell_integration` for details on how this setting works. The special value :code:`inherited` means -use the setting from kitty.conf. This setting is useful for overriding +use the setting from :file:`kitty.conf`. This setting is useful for overriding integration on a per-host basis.''') opt('login_shell', '', long_text=''' diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py index 814cc9d9f..86adc325b 100644 --- a/kitty_tests/ssh.py +++ b/kitty_tests/ssh.py @@ -12,11 +12,11 @@ from kittens.ssh.config import load_config, options_for_host from kittens.ssh.main import ( bootstrap_script, get_connection_data, wrap_bootstrap_script ) +from kittens.ssh.options.types import Options as SSHOptions from kittens.ssh.options.utils import DELETE_ENV_VAR from kittens.transfer.utils import set_paths from kitty.constants import is_macos, runtime_dir from kitty.fast_data_types import CURSOR_BEAM -from kitty.options.utils import shell_integration from kitty.utils import SSHConnectionData from . import BaseTest @@ -240,9 +240,9 @@ copy --exclude */w.* d1 test_script = f'print("UNTAR_DONE", flush=True); {test_script}' else: test_script = f'echo "UNTAR_DONE"; {test_script}' + ssh_opts['shell_integration'] = SHELL_INTEGRATION_VALUE or 'disabled' script, replacements, shm = bootstrap_script( - script_type='py' if 'python' in sh else 'sh', request_id="testing", - test_script=test_script, ssh_opts_dict={'*': ssh_opts}, + SSHOptions(ssh_opts), script_type='py' if 'python' in sh else 'sh', request_id="testing", test_script=test_script ) try: env = basic_shell_env(home_dir) @@ -250,9 +250,9 @@ copy --exclude */w.* d1 os.makedirs(os.path.join(home_dir, '.local', 'share', 'fish', 'generated_completions'), exist_ok=True) # prevent newuser-install from running open(os.path.join(home_dir, '.zshrc'), 'w').close() - options = {'shell_integration': shell_integration(SHELL_INTEGRATION_VALUE or 'disabled')} cmd = wrap_bootstrap_script(script, sh) - pty = self.create_pty([launcher, '-c', ' '.join(cmd)], cwd=home_dir, env=env, options=options) + pty = self.create_pty([launcher, '-c', ' '.join(cmd)], cwd=home_dir, env=env) + del cmd if pre_data: pty.write_buf = pre_data.encode('utf-8') del script diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py index b7a58c409..dc23f15e8 100644 --- a/shell-integration/ssh/bootstrap.py +++ b/shell-integration/ssh/bootstrap.py @@ -4,7 +4,6 @@ import atexit import base64 -import getpass import io import os import pwd @@ -59,9 +58,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', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:compression=bz2:pw=DATA_PASSWORD'.format(hostname, getpass.getuser()))) + 'ssh', 'id=REQUEST_ID:pwfile=PASSWORD_FILENAME:pw=DATA_PASSWORD')) def debug(msg): diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index 031d76005..2ace11432 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/shell-integration/ssh/bootstrap.sh @@ -5,7 +5,6 @@ saved_tty_settings="" tdir="" shell_integration_dir="" -compression="gz" cleanup_on_bootstrap_exit() { [ -n "$saved_tty_settings" ] && command stty "$saved_tty_settings" 2> /dev/null < /dev/tty @@ -165,11 +164,6 @@ fi debug() { dcs_to_kitty "print" "debug: $1"; } echo_via_kitty() { dcs_to_kitty "echo" "$1"; } -hostname="$HOSTNAME" -[ -z "$hostname" ] && hostname="$(command hostname 2> /dev/null)" -[ -z "$hostname" ] && hostname="$(command hostnamectl hostname 2> /dev/null)" -[ -z "$hostname" ] && hostname="$(command uname -m 2> /dev/null)" -[ -z "$hostname" ] && hostname="_" # ensure $HOME is set [ -z "$HOME" ] && HOME=~ # ensure $USER is set @@ -183,8 +177,7 @@ request_data="REQUEST_DATA" [ "$request_data" = "1" ] && init_tty trap "cleanup_on_bootstrap_exit" EXIT if [ "$tty_ok" = "y" -a "$request_data" = "1" ]; then - command -v "bzip2" > /dev/null 2> /dev/null && compression="bz2" - dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":compression="$compression":pw="DATA_PASSWORD"" + dcs_to_kitty "ssh" "id="REQUEST_ID":pwfile="PASSWORD_FILENAME":pw="DATA_PASSWORD"" fi record_separator=$(printf "\036") @@ -225,9 +218,7 @@ untar_and_read_env() { tdir=$(command mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX") [ $? = 0 ] || die "Creating temp directory failed" - cflag="j" - [ "$compression" = "gz" ] && cflag="z" - read_n_bytes_from_tty "$1" | base64_decode | command tar "x${cflag}pf" - -C "$tdir" + read_n_bytes_from_tty "$1" | base64_decode | command tar "xpzf" "-" "-C" "$tdir" data_file="$tdir/data.sh" [ -f "$data_file" ] && . "$data_file" [ -z "$KITTY_SSH_KITTEN_DATA_DIR" ] && die "Failed to read SSH data from tty"