Avoid needing to initialize tty state in bootstrap scripts
This commit is contained in:
Kovid Goyal
parent
74f0057ec8
commit
e1504c4775
@ -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],
|
||||
ssh_opts: SSHOptions,
|
||||
hostname: str = 'localhost', cli_hostname: str = '', cli_uname: str = '',
|
||||
remote_args: List[str], ssh_opts: SSHOptions,
|
||||
hostname: str = 'localhost', cli_hostname: str = '', cli_uname: str = '', echo_on: bool = True,
|
||||
) -> 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:
|
||||
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
|
||||
def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple[str, ...], echo_on: bool) -> NoReturn:
|
||||
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 = running_in_kitty and host_opts.share_connections
|
||||
rcmd, replacements = get_remote_command(remote_args, host_opts, hostname, hostname_for_match, uname)
|
||||
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']))
|
||||
@ -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__':
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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')
|
||||
|
||||
@ -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:
|
||||
termios.tcsetattr(tty_fd, termios.TCSANOW, original_termios_state)
|
||||
original_termios_state = None
|
||||
if echo_on:
|
||||
s = termios.tcgetattr(tty_fd)
|
||||
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
|
||||
try:
|
||||
global tty_fd, login_shell
|
||||
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()
|
||||
|
||||
@ -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""
|
||||
fi
|
||||
dcs_to_kitty "ssh" "id="REQUEST_ID":pwfile="PASSWORD_FILENAME":pw="DATA_PASSWORD""
|
||||
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
|
||||
get_data
|
||||
cleanup_on_bootstrap_exit
|
||||
if [ -n "$leading_data" ]; then
|
||||
# ask for the SSH data
|
||||
get_data
|
||||
cleanup_on_bootstrap_exit
|
||||
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
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user