diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index 7d309f7f6..7f7f6a08a 100644 --- a/kittens/ssh/main.py +++ b/kittens/ssh/main.py @@ -4,13 +4,14 @@ import atexit import io import os +import re import shlex import subprocess import sys import tarfile import tempfile from contextlib import suppress -from typing import Iterator, List, NoReturn, Optional, Set, Tuple +from typing import Dict, Iterator, List, NoReturn, Optional, Set, Tuple from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir from kitty.short_uuid import uuid4 @@ -18,8 +19,10 @@ from kitty.utils import SSHConnectionData from .completion import complete, ssh_options +DEFAULT_SHELL_INTEGRATION_DEST = '.local/share/kitty-ssh-kitten/shell-integration' -def make_tarfile(hostname: str = '') -> bytes: + +def make_tarfile(hostname: str = '', shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> bytes: def filter_files(tarinfo: tarfile.TarInfo) -> Optional[tarfile.TarInfo]: if tarinfo.name.endswith('ssh/bootstrap.sh') or tarinfo.name.endswith('ssh/bootstrap.py'): @@ -30,12 +33,12 @@ def make_tarfile(hostname: str = '') -> bytes: buf = io.BytesIO() with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf: - tf.add(shell_integration_dir, arcname='.local/share/kitty-ssh-kitten/shell-integration', filter=filter_files) + tf.add(shell_integration_dir, arcname=shell_integration_dest, filter=filter_files) tf.add(terminfo_dir, arcname='.terminfo', filter=filter_files) return buf.getvalue() -def get_ssh_data(msg: str) -> Iterator[bytes]: +def get_ssh_data(msg: str, shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> Iterator[bytes]: yield b"KITTY_SSH_DATA_START\n" try: hostname, pwfilename, pw = msg.split(':', 2) @@ -50,7 +53,7 @@ def get_ssh_data(msg: str) -> Iterator[bytes]: yield b' incorrect ssh data password\n' else: try: - data = make_tarfile(hostname) + data = make_tarfile(hostname, shell_integration_dest) except Exception: yield b' error while gathering ssh data\n' else: @@ -68,21 +71,28 @@ def safe_remove(x: str) -> None: os.remove(x) -def prepare_script(ans: str, EXEC_CMD: str = '') -> str: - ans = ans.replace('EXEC_CMD', EXEC_CMD, 1) +def prepare_script(ans: str, replacements: Dict[str, str]) -> str: pw = uuid4() 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 + replacements['DATA_PASSWORD'] = pw + 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) + replacements['SHELL_INTEGRATION_VALUE'] = replacements.get('SHELL_INTEGRATION_VALUE', 'enabled') + + def sub(m: 're.Match[str]') -> str: + return replacements[m.group()] + + return re.sub('|'.join(fr'\b{k}\b' for k in replacements), sub, ans) -def bootstrap_script(EXEC_CMD: str = '', script_type: str = 'sh') -> str: +def bootstrap_script(script_type: str = 'sh', **replacements: str) -> 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) + return prepare_script(ans, replacements) SHELL_SCRIPT = '''\ diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py index d8d95789a..c6ce2d5bb 100644 --- a/kitty_tests/ssh.py +++ b/kitty_tests/ssh.py @@ -45,10 +45,11 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77) q = shutil.which(sh) if q: with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir: - script = bootstrap_script('echo TEST_DONE; return 0') + script = bootstrap_script(EXEC_CMD='echo UNTAR_DONE; exit 0') env = basic_shell_env(tdir) pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=tdir, env=env) self.check_bootstrap(tdir, pty) def check_bootstrap(self, home_dir, pty): - pty.wait_till(lambda: 'TEST_DONE' in pty.screen_contents()) + pty.wait_till(lambda: 'UNTAR_DONE' in pty.screen_contents()) + self.assertTrue(os.path.exists(os.path.join(home_dir, '.terminfo/kitty.terminfo'))) diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index dfd5d85d2..131934307 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/shell-integration/ssh/bootstrap.sh @@ -83,3 +83,153 @@ fi # If a command was passed to SSH execute it here EXEC_CMD + +shell_integration_dir="$HOME/SHELL_INTEGRATION_DIR" + +login_shell_is_ok() { + if [ -z "$login_shell" ] || [ ! -x "$login_shell" ]; then return 1; fi + case "$login_shell" in + *sh) return 0; + esac + return 1; +} + +detect_python() { + python=$(command -v python3) + if [ -z "$python" ]; then python=$(command -v python2); fi + if [ -z "$python" ]; then python=python; fi + if [ -z "$python" || ! -x "$python" ]; then return 1; fi + return 0; +} + +using_getent() { + cmd=$(command -v getent) + if [ -n "$cmd" ]; then + output=$($cmd passwd $USER 2>/dev/null) + if [ $? = 0 ]; then + login_shell=$(echo $output | cut -d: -f7); + if login_shell_is_ok; then return 0; fi + fi + fi + return 1; +} + +using_id() { + cmd=$(command -v id) + if [ -n "$cmd" ]; then + output=$($cmd -P $USER 2>/dev/null) + if [ $? = 0 ]; then + login_shell=$(echo $output | cut -d: -f7); + if login_shell_is_ok; then return 0; fi + fi + fi + return 1; +} + +using_passwd() { + cmd=$(command -v grep) + if [ -n "$cmd" ]; then + output=$($cmd "^$USER:" /etc/passwd 2>/dev/null) + if [ $? = 0 ]; then + login_shell=$(echo $output | cut -d: -f7); + if login_shell_is_ok; then return 0; fi + fi + fi + return 1; +} + +using_python() { + if detect_python; then + output=$($python -c "import pwd, os; print(pwd.getpwuid(os.geteuid()).pw_shell)") + if [ $? = 0 ]; then + login_shell=$output; + if login_shell_is_ok; then return 0; fi + fi + fi + return 1 +} + +execute_with_python() { + if detect_python; then + exec $python -c "import os; os.execl('$login_shell', '-' '$shell_name')" + fi + return 1; +} + +LOGIN_SHELL="OVERRIDE_LOGIN_SHELL" +if [ -n "$LOGIN_SHELL" ]; then + login_shell="$LOGIN_SHELL" +else + using_getent || using_id || using_python || using_passwd || die "Could not detect login shell"; +fi +shell_name=$(basename $login_shell) + +export KITTY_SHELL_INTEGRATION="SHELL_INTEGRATION_VALUE" + +exec_bash_with_integration() { + export ENV="$shell_integration_dir/bash/kitty.bash" + export KITTY_BASH_INJECT="1" + exec "$login_shell" "--posix" +} + +exec_zsh_with_integration() { + zdotdir="$ZDOTDIR" + if [ -z "$zdotdir" ]; then + zdotdir=~; + unset KITTY_ORIG_ZDOTDIR + else + export KITTY_ORIG_ZDOTDIR="$zdotdir" + fi + # dont prevent zsh-new-user from running + if [ -e "$zdotdir/.zshrc" || -e "$zdotdir/.zshenv" || -e "$zdotdir/.zprofile" || -e "$zdotdir/.zlogin" ]; then + export ZDOTDIR="$shell_integration_dir/zsh" + exec "$login_shell" "-l" + fi +} + +exec_fish_with_integration() { + if [ -z "$XDG_DATA_DIRS" ]; then + export XDG_DATA_DIRS="$shell_integration_dir" + else + export XDG_DATA_DIRS="$shell_integration_dir:$XDG_DATA_DIRS" + fi + export KITTY_FISH_XDG_DATA_DIR="$shell_integration_dir" + exec "$login_shell" "-l" +} + +exec_with_shell_integration() { + case "$login_shell" in + *"zsh") + exec_zsh_with_integration + ;; + *"bash") + exec_bash_with_integration + ;; + *"fish") + exec_fish_with_integration + ;; + esac +} + +case "$KITTY_SHELL_INTEGRATION" in + "") + unset KITTY_SHELL_INTEGRATION + ;; + *"no-rc"*) + ;; + *) + exec_with_shell_integration + unset KITTY_SHELL_INTEGRATION + ;; +esac + +# We need to pass the first argument to the executed program with a leading - +# to make sure the shell executes as a login shell. Note that not all shells +# support exec -a so we use the below to try to detect such shells +shell_name=$(basename $login_shell) +if [ -z "$PIPESTATUS" ]; then + # the dash shell does not support exec -a and also does not define PIPESTATUS + execute_with_python + exec $login_shell "-l" +fi +exec -a "-$shell_name" $login_shell