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,30 +476,18 @@ 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 main(args: List[str]) -> NoReturn:
+def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple[str, ...], echo_on: bool) -> NoReturn:
-    args = args[1:]
-    if args and args[0] == 'use-python':
-        args = args[1:]  # backwards compat from when we had a python implementation
-    try:
-        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'):
-        passthrough = True
     cmd = ['ssh'] + ssh_args
-    if passthrough:
-        cmd += server_args
-    else:
     hostname, remote_args = server_args[0], server_args[1:]
     if not remote_args:
         cmd.append('-t')
@@ -527,9 +514,8 @@ def main(args: List[str]) -> NoReturn:
         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 = host_opts.share_connections
-        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, echo_on)
-        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']))
@@ -541,13 +527,30 @@ def main(args: List[str]) -> NoReturn:
     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?')
 
 
+def main(args: List[str]) -> NoReturn:
+    args = args[1:]
+    if args and args[0] == 'use-python':
+        args = args[1:]  # backwards compat from when we had a python implementation
+    try:
+        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') or not os.environ.get('KITTY_PID'):
+        raise SystemExit('The SSH kitten is meant to run inside a kitty window')
+    if passthrough:
+        raise SystemExit('The SSH kitten is meant for interactive use via SSH only')
+    if not sys.stdin.isatty():
+        raise SystemExit('The SSH kitten is meant for interactive use only, STDIN must be a terminal')
+    with restore_terminal_state() as echo_on, no_echo(sys.stdin.fileno()):
+        run_ssh(ssh_args, server_args, found_extra_args, echo_on)
+
+
 if __name__ == '__main__':
     main(sys.argv)
 elif __name__ == '__completer__':

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,25 +214,10 @@ def exec_with_shell_integration():
 
 
 def main():
-    global tty_fd, original_termios_state, login_shell
+    global tty_fd, login_shell
-    try:
     tty_fd = os.open(os.ctermid(), os.O_RDWR | os.O_CLOEXEC)
-    except OSError:
-        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()

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
     # 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" > /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