Avoid needing to initialize tty state in bootstrap scripts

kittens/ssh/main.py

@@ -31,7 +31,7 @@ from kitty.constants import (
 from kitty.options.types import Options
 from kitty.shm import SharedMemory
 from kitty.types import run_once
-from kitty.utils import SSHConnectionData
+from kitty.utils import SSHConnectionData, no_echo
 
 from .completion import complete, ssh_options
 from .config import init_config, options_for_host
@@ -216,7 +216,7 @@ def prepare_exec_cmd(remote_args: Sequence[str], is_python: bool) -> str:
 def bootstrap_script(
     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',
+    request_data: str = '1', echo_on: bool = True
 ) -> Tuple[str, Dict[str, str], SharedMemory]:
     if request_id is None:
         request_id = os.environ['KITTY_PID'] + '-' + os.environ['KITTY_WINDOW_ID']
@@ -233,7 +233,7 @@ def bootstrap_script(
         atexit.register(shm.unlink)
     replacements = {
         'DATA_PASSWORD': pw, 'PASSWORD_FILENAME': shm.name, 'EXEC_CMD': exec_cmd, 'TEST_SCRIPT': test_script,
-        'REQUEST_ID': request_id, 'REQUEST_DATA': request_data,
+        'REQUEST_ID': request_id, 'REQUEST_DATA': request_data, 'ECHO_ON': '1' if echo_on else '0',
     }
     return prepare_script(ans, replacements), replacements, shm
 
@@ -428,16 +428,15 @@ def wrap_bootstrap_script(sh_script: str, interpreter: str) -> List[str]:
 
 
 def get_remote_command(
-    remote_args: List[str],
+    remote_args: List[str], ssh_opts: SSHOptions,
-    ssh_opts: SSHOptions,
+    hostname: str = 'localhost', cli_hostname: str = '', cli_uname: str = '', echo_on: bool = True,
-    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(
         ssh_opts, script_type='py' if is_python else 'sh', remote_args=remote_args,
-        cli_hostname=cli_hostname, cli_uname=cli_uname)
+        cli_hostname=cli_hostname, cli_uname=cli_uname, echo_on=echo_on)
     return wrap_bootstrap_script(sh_script, interpreter), replacements
 
 
@@ -477,16 +476,63 @@ def connection_sharing_args(opts: SSHOptions, kitty_pid: int) -> List[str]:
 
 
 @contextmanager
-def restore_terminal_state() -> Iterator[None]:
+def restore_terminal_state() -> Iterator[bool]:
     import termios
     with open(os.ctermid()) as f:
         val = termios.tcgetattr(f.fileno())
         try:
-            yield
+            yield bool(val[3] & termios.ECHO)
         finally:
             termios.tcsetattr(f.fileno(), termios.TCSAFLUSH, val)
 
 
+def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple[str, ...], echo_on: bool) -> NoReturn:
+    cmd = ['ssh'] + ssh_args
+    hostname, remote_args = server_args[0], server_args[1:]
+    if not remote_args:
+        cmd.append('-t')
+    insertion_point = len(cmd)
+    cmd.append('--')
+    cmd.append(hostname)
+    uname = getuser()
+    if hostname.startswith('ssh://'):
+        from urllib.parse import urlparse
+        purl = urlparse(hostname)
+        hostname_for_match = purl.hostname or hostname
+        uname = purl.username or uname
+    elif '@' in hostname and hostname[0] != '@':
+        uname, hostname_for_match = hostname.split('@', 1)
+    else:
+        hostname_for_match = hostname
+    hostname_for_match = hostname.split('@', 1)[-1].split(':', 1)[0]
+    overrides = []
+    pat = re.compile(r'^([a-zA-Z0-9_]+)[ \t]*=')
+    for i, a in enumerate(found_extra_args):
+        if i % 2 == 1:
+            overrides.append(pat.sub(r'\1 ', a.lstrip()))
+    if overrides:
+        overrides.insert(0, f'hostname {uname}@{hostname_for_match}')
+    so = init_config(overrides)
+    host_opts = options_for_host(hostname_for_match, uname, so)
+    use_control_master = host_opts.share_connections
+    rcmd, replacements = get_remote_command(remote_args, host_opts, hostname, hostname_for_match, uname, echo_on)
+    cmd += rcmd
+    if use_control_master:
+        cmd[insertion_point:insertion_point] = connection_sharing_args(host_opts, int(os.environ['KITTY_PID']))
+    # We force use of askpass so that OpenSSH does not use the tty leaving
+    # it free for us to use
+    os.environ['SSH_ASKPASS_REQUIRE'] = 'force'
+    if not os.environ.get('SSH_ASKPASS'):
+        os.environ['SSH_ASKPASS'] = os.path.join(shell_integration_dir, 'ssh', 'askpass.py')
+    import subprocess
+    with suppress(FileNotFoundError):
+        try:
+            raise SystemExit(subprocess.run(cmd).returncode)
+        except KeyboardInterrupt:
+            raise SystemExit(1)
+    raise SystemExit('Could not find the ssh executable, is it in your PATH?')
+
+
 def main(args: List[str]) -> NoReturn:
     args = args[1:]
     if args and args[0] == 'use-python':
@@ -495,57 +541,14 @@ def main(args: List[str]) -> NoReturn:
         ssh_args, server_args, passthrough, found_extra_args = parse_ssh_args(args, extra_args=('--kitten',))
     except InvalidSSHArgs as e:
         e.system_exit()
-    if not os.environ.get('KITTY_WINDOW_ID'):
+    if not os.environ.get('KITTY_WINDOW_ID') or not os.environ.get('KITTY_PID'):
-        passthrough = True
+        raise SystemExit('The SSH kitten is meant to run inside a kitty window')
-    cmd = ['ssh'] + ssh_args
     if passthrough:
-        cmd += server_args
+        raise SystemExit('The SSH kitten is meant for interactive use via SSH only')
-    else:
+    if not sys.stdin.isatty():
-        hostname, remote_args = server_args[0], server_args[1:]
+        raise SystemExit('The SSH kitten is meant for interactive use only, STDIN must be a terminal')
-        if not remote_args:
+    with restore_terminal_state() as echo_on, no_echo(sys.stdin.fileno()):
-            cmd.append('-t')
+        run_ssh(ssh_args, server_args, found_extra_args, echo_on)
-        insertion_point = len(cmd)
-        cmd.append('--')
-        cmd.append(hostname)
-        uname = getuser()
-        if hostname.startswith('ssh://'):
-            from urllib.parse import urlparse
-            purl = urlparse(hostname)
-            hostname_for_match = purl.hostname or hostname
-            uname = purl.username or uname
-        elif '@' in hostname and hostname[0] != '@':
-            uname, hostname_for_match = hostname.split('@', 1)
-        else:
-            hostname_for_match = hostname
-        hostname_for_match = hostname.split('@', 1)[-1].split(':', 1)[0]
-        overrides = []
-        pat = re.compile(r'^([a-zA-Z0-9_]+)[ \t]*=')
-        for i, a in enumerate(found_extra_args):
-            if i % 2 == 1:
-                overrides.append(pat.sub(r'\1 ', a.lstrip()))
-        if overrides:
-            overrides.insert(0, f'hostname {uname}@{hostname_for_match}')
-        so = init_config(overrides)
-        host_opts = options_for_host(hostname_for_match, uname, so)
-        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']))
-        # We force use of askpass so that OpenSSH does not use the tty leaving
-        # it free for us to use
-        os.environ['SSH_ASKPASS_REQUIRE'] = 'force'
-        if not os.environ.get('SSH_ASKPASS'):
-            os.environ['SSH_ASKPASS'] = os.path.join(shell_integration_dir, 'ssh', 'askpass.py')
-    import subprocess
-    with suppress(FileNotFoundError):
-        try:
-            with restore_terminal_state():
-                raise SystemExit(subprocess.run(cmd).returncode)
-        except KeyboardInterrupt:
-            raise SystemExit(1)
-    raise SystemExit('Could not find the ssh executable, is it in your PATH?')
 
 
 if __name__ == '__main__':

kitty_tests/__init__.py

@@ -194,6 +194,15 @@ class PTY:
         self.screen = Screen(self.callbacks, rows, columns, scrollback, cell_width, cell_height, 0, self.callbacks)
         self.received_bytes = b''
 
+    def turn_off_echo(self):
+        s = termios.tcgetattr(self.master_fd)
+        s[3] &= ~termios.ECHO
+        termios.tcsetattr(self.master_fd, termios.TCSANOW, s)
+
+    def is_echo_on(self):
+        s = termios.tcgetattr(self.master_fd)
+        return True if s[3] & termios.ECHO else False
+
     def __del__(self):
         if not self.is_child:
             fd = self.master_fd

kitty_tests/ssh.py

@@ -156,6 +156,7 @@ copy --exclude */w.* d1
                 pty.wait_till(lambda: 'TSET={}'.format(tset.replace('$A', 'AAA')) in pty.screen_contents())
                 self.assertNotIn('COLORTERM', pty.screen_contents())
                 pty.wait_till(lambda: '/cwd' in pty.screen_contents())
+                self.assertTrue(pty.is_echo_on())
 
     def test_ssh_bootstrap_with_different_launchers(self):
         for launcher in self.all_possible_sh:
@@ -252,6 +253,7 @@ copy --exclude */w.* d1
             open(os.path.join(home_dir, '.zshrc'), 'w').close()
             cmd = wrap_bootstrap_script(script, sh)
             pty = self.create_pty([launcher, '-c', ' '.join(cmd)], cwd=home_dir, env=env)
+            pty.turn_off_echo()
             del cmd
             if pre_data:
                 pty.write_buf = pre_data.encode('utf-8')

shell-integration/ssh/bootstrap.py

@@ -16,7 +16,7 @@ import tempfile
 import termios
 
 tty_fd = -1
-original_termios_state = None
+echo_on = int('ECHO_ON')
 data_dir = shell_integration_dir = ''
 request_data = int('REQUEST_DATA')
 leading_data = b''
@@ -25,11 +25,12 @@ login_shell = pwd.getpwuid(os.geteuid()).pw_shell or 'sh'
 
 
 def cleanup():
-    global tty_fd, original_termios_state
+    global tty_fd
     if tty_fd > -1:
-        if original_termios_state is not None:
+        if echo_on:
-            termios.tcsetattr(tty_fd, termios.TCSANOW, original_termios_state)
+            s = termios.tcgetattr(tty_fd)
-            original_termios_state = None
+            s[3] |= termios.ECHO
+            termios.tcsetattr(tty_fd, termios.TCSANOW, s)
         os.close(tty_fd)
         tty_fd = -1
 
@@ -150,7 +151,6 @@ def get_data():
     data = []
     with open(tty_fd, 'rb', closefd=False) as f:
         data = b''.join(iter_base64_data(f))
-    cleanup()
     if leading_data:
         # 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
@@ -214,26 +214,11 @@ def exec_with_shell_integration():
 
 
 def main():
-    global tty_fd, original_termios_state, login_shell
+    global tty_fd, login_shell
+    tty_fd = os.open(os.ctermid(), os.O_RDWR | os.O_CLOEXEC)
     try:
-        tty_fd = os.open(os.ctermid(), os.O_RDWR | os.O_CLOEXEC)
+        send_data_request()
-    except OSError:
+        get_data()
-        pass
-    else:
-        if request_data:
-            try:
-                original_termios_state = termios.tcgetattr(tty_fd)
-            except OSError:
-                pass
-            else:
-                new_state = termios.tcgetattr(tty_fd)
-                new_state[3] &= ~termios.ECHO
-                termios.tcsetattr(tty_fd, termios.TCSANOW, new_state)
-    try:
-        if original_termios_state is not None:
-            send_data_request()
-        if tty_fd > -1:
-            get_data()
     finally:
         cleanup()
     cwd = os.environ.pop('KITTY_LOGIN_CWD', '')

shell-integration/ssh/bootstrap.sh

@@ -2,14 +2,14 @@
 # Copyright (C) 2022 Kovid Goyal <kovid at kovidgoyal.net>
 # Distributed under terms of the GPLv3 license.
 
-saved_tty_settings=""
 tdir=""
 shell_integration_dir=""
+echo_on="ECHO_ON"
 
 cleanup_on_bootstrap_exit() {
-    [ -n "$saved_tty_settings" ] && command stty "$saved_tty_settings" 2> /dev/null < /dev/tty
+    [ "$echo_on" = "1" ] && command stty "echo" 2> /dev/null < /dev/tty
+    echo_on="0"
     [ -n "$tdir" ] && command rm -rf "$tdir"
-    saved_tty_settings=""
     tdir=""
 }
 
@@ -58,18 +58,6 @@ else
     die "base64 executable not present on remote host, ssh kitten cannot function."
 fi
 
-init_tty() {
-    saved_tty_settings=$(command stty -g 2> /dev/null < /dev/tty)
-    tty_ok="n"
-    [ -n "$saved_tty_settings" ] && tty_ok="y"
-
-    if [ "$tty_ok" = "y" ]; then
-        command stty -echo 2> /dev/null < /dev/tty || die "stty failed to set raw mode"
-        return 0
-    fi
-    return 1
-}
-
 dcs_to_kitty() { printf "\033P@kitty-$1|%s\033\134" "$(printf "%s" "$2" | base64_encode)" > /dev/tty; }
 debug() { dcs_to_kitty "print" "debug: $1"; }
 echo_via_kitty() { dcs_to_kitty "echo" "$1"; }
@@ -84,11 +72,8 @@ leading_data=""
 login_cwd=""
 
 request_data="REQUEST_DATA"
-[ "$request_data" = "1" ] && init_tty
 trap "cleanup_on_bootstrap_exit" EXIT
-if [ "$tty_ok" = "y" -a "$request_data" = "1" ]; then
+dcs_to_kitty "ssh" "id="REQUEST_ID":pwfile="PASSWORD_FILENAME":pw="DATA_PASSWORD""
-    dcs_to_kitty "ssh" "id="REQUEST_ID":pwfile="PASSWORD_FILENAME":pw="DATA_PASSWORD""
-fi
 record_separator=$(printf "\036")
 
 mv_files_and_dirs() {
@@ -177,18 +162,16 @@ get_data() {
     untar_and_read_env
 }
 
-if [ "$tty_ok" = "y" ]; then
+# ask for the SSH data
-    # ask for the SSH data
+get_data
-    get_data
+cleanup_on_bootstrap_exit
-    cleanup_on_bootstrap_exit
+if [ -n "$leading_data" ]; then
-    if [ -n "$leading_data" ]; then
+    # clear current line as it might have things echoed on it from leading_data
-        # 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
-        # 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
-        # have been sent before the script had a chance to run
+    printf "\r\033[K" > /dev/tty
-        printf "\r\033[K" > /dev/tty
-    fi
-    [ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data"
 fi
+[ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data"
 
 login_shell_is_ok() {
     if [ -n "$login_shell" -a -x "$login_shell" ]; then return 0; fi
@@ -282,14 +265,6 @@ shell_name=$(command basename $login_shell)
 # If a command was passed to SSH execute it here
 EXEC_CMD
 
-if [ "$tty_ok" = "n" ]; then
-    if [ -z "$(command -v stty)" ]; then
-        printf "%s\n" "stty missing ssh kitten cannot function" > /dev/stderr
-    else
-        printf "%s\n" "stty failed ssh kitten cannot function" > /dev/stderr
-    fi
-fi
-
 exec_zsh_with_integration() {
     zdotdir="$ZDOTDIR"
     if [ -z "$zdotdir" ]; then