Start work on testing the bootstrap script

This commit is contained in:
Kovid Goyal 2022-02-23 15:34:19 +05:30
parent f37d947dd5
commit e73525d0a2
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C
6 changed files with 64 additions and 55 deletions

View File

@ -37,7 +37,7 @@ def install_deps():
run('sudo apt-get install -y libgl1-mesa-dev libxi-dev libxrandr-dev libxinerama-dev ca-certificates' run('sudo apt-get install -y libgl1-mesa-dev libxi-dev libxrandr-dev libxinerama-dev ca-certificates'
' libxcursor-dev libxcb-xkb-dev libdbus-1-dev libxkbcommon-dev libharfbuzz-dev libx11-xcb-dev zsh' ' libxcursor-dev libxcb-xkb-dev libdbus-1-dev libxkbcommon-dev libharfbuzz-dev libx11-xcb-dev zsh'
' libpng-dev liblcms2-dev libfontconfig-dev libxkbcommon-x11-dev libcanberra-dev librsync-dev uuid-dev' ' libpng-dev liblcms2-dev libfontconfig-dev libxkbcommon-x11-dev libcanberra-dev librsync-dev uuid-dev'
' zsh bash dash') ' zsh bash dash posh')
# for some reason these directories are world writable which causes zsh # for some reason these directories are world writable which causes zsh
# compinit to break # compinit to break
run('sudo chmod -R og-w /usr/share/zsh') run('sudo chmod -R og-w /usr/share/zsh')

View File

@ -13,7 +13,7 @@ from contextlib import suppress
from typing import Iterator, List, NoReturn, Optional, Set, Tuple from typing import Iterator, List, NoReturn, Optional, Set, Tuple
from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
from kitty.short_uuid import uuid4_for_escape_code from kitty.short_uuid import uuid4
from kitty.utils import SSHConnectionData from kitty.utils import SSHConnectionData
from .completion import complete, ssh_options from .completion import complete, ssh_options
@ -36,30 +36,31 @@ def make_tarfile(hostname: str = '') -> bytes:
def get_ssh_data(msg: str) -> Iterator[bytes]: def get_ssh_data(msg: str) -> Iterator[bytes]:
yield b"KITTY_SSH_DATA_START" yield b"KITTY_SSH_DATA_START\n"
try: try:
hostname, pwfilename, pw = msg.split(':', 2) hostname, pwfilename, pw = msg.split(':', 2)
except Exception: except Exception:
yield b' invalid ssh data request message' yield b' invalid ssh data request message\n'
try: try:
with open(os.path.join(cache_dir(), pwfilename)) as f: with open(os.path.join(cache_dir(), pwfilename)) as f:
os.unlink(f.name) os.unlink(f.name)
if pw != f.read(): if pw != f.read():
raise ValueError('Incorrect password') raise ValueError('Incorrect password')
except Exception: except Exception:
yield b' incorrect ssh data password' yield b' incorrect ssh data password\n'
else: else:
try: try:
data = make_tarfile(hostname) data = make_tarfile(hostname)
except Exception: except Exception:
yield b' error while gathering ssh data' yield b' error while gathering ssh data\n'
else: else:
from base64 import standard_b64encode from base64 import standard_b64encode
encoded_data = memoryview(standard_b64encode(data)) encoded_data = memoryview(standard_b64encode(data))
while encoded_data: while encoded_data:
yield encoded_data[:1024] yield encoded_data[:2048]
encoded_data = encoded_data[1024:] yield b'\n'
yield b"KITTY_SSH_DATA_END" encoded_data = encoded_data[2048:]
yield b"KITTY_SSH_DATA_END\n"
def safe_remove(x: str) -> None: def safe_remove(x: str) -> None:
@ -69,7 +70,7 @@ def safe_remove(x: str) -> None:
def prepare_script(ans: str, EXEC_CMD: str = '') -> str: def prepare_script(ans: str, EXEC_CMD: str = '') -> str:
ans = ans.replace('EXEC_CMD', EXEC_CMD, 1) ans = ans.replace('EXEC_CMD', EXEC_CMD, 1)
pw = uuid4_for_escape_code() pw = uuid4()
with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', dir=cache_dir(), delete=False) as tf: with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', dir=cache_dir(), delete=False) as tf:
tf.write(pw.encode('utf-8')) tf.write(pw.encode('utf-8'))
atexit.register(safe_remove, tf.name) atexit.register(safe_remove, tf.name)
@ -338,18 +339,6 @@ def parse_ssh_args(args: List[str]) -> Tuple[List[str], List[str], bool]:
return ssh_args, server_args, passthrough return ssh_args, server_args, passthrough
def quote(x: str) -> str:
# we have to escape unbalanced quotes and other unparsable
# args as they will break the shell script
# But we do not want to quote things like * or 'echo hello'
# See https://github.com/kovidgoyal/kitty/issues/1787
try:
shlex.split(x)
except ValueError:
x = shlex.quote(x)
return x
def get_posix_cmd(terminfo: str, remote_args: List[str]) -> List[str]: def get_posix_cmd(terminfo: str, remote_args: List[str]) -> List[str]:
sh_script = SHELL_SCRIPT.replace('TERMINFO', terminfo, 1) sh_script = SHELL_SCRIPT.replace('TERMINFO', terminfo, 1)
command_to_execute = '' command_to_execute = ''

View File

@ -884,7 +884,6 @@ class Window:
from kittens.ssh.main import get_ssh_data from kittens.ssh.main import get_ssh_data
for line in get_ssh_data(msg): for line in get_ssh_data(msg):
self.write_to_child(line) self.write_to_child(line)
self.write_to_child('\n')
def handle_remote_print(self, msg: str) -> None: def handle_remote_print(self, msg: str) -> None:
text = process_remote_print(msg) text = process_remote_print(msg)

View File

@ -26,8 +26,9 @@ from kitty.window import process_remote_print, process_title_from_child
class Callbacks: class Callbacks:
def __init__(self) -> None: def __init__(self, pty=None) -> None:
self.clear() self.clear()
self.pty = pty
def write(self, data) -> None: def write(self, data) -> None:
self.wtcbuf += data self.wtcbuf += data
@ -93,6 +94,18 @@ class Callbacks:
text = process_remote_print(msg) text = process_remote_print(msg)
print(text, file=sys.__stderr__) print(text, file=sys.__stderr__)
def handle_remote_ssh(self, msg):
from kittens.ssh.main import get_ssh_data
for line in get_ssh_data(msg):
self.pty.write_to_child(line)
self.pty.process_input_from_child(timeout=0)
def handle_remote_echo(self, msg):
from base64 import standard_b64decode
data = standard_b64decode(msg)
self.pty.write_to_child(data)
self.pty.process_input_from_child(timeout=0)
def filled_line_buf(ynum=5, xnum=5, cursor=Cursor()): def filled_line_buf(ynum=5, xnum=5, cursor=Cursor()):
ans = LineBuf(ynum, xnum) ans = LineBuf(ynum, xnum)
@ -177,7 +190,7 @@ class PTY:
self.cell_width = cell_width self.cell_width = cell_width
self.cell_height = cell_height self.cell_height = cell_height
self.set_window_size(rows=rows, columns=columns) self.set_window_size(rows=rows, columns=columns)
self.callbacks = Callbacks() self.callbacks = Callbacks(self)
self.screen = Screen(self.callbacks, rows, columns, scrollback, cell_width, cell_height, 0, self.callbacks) self.screen = Screen(self.callbacks, rows, columns, scrollback, cell_width, cell_height, 0, self.callbacks)
self.received_bytes = b'' self.received_bytes = b''

View File

@ -3,13 +3,15 @@
import os import os
import shlex
import shutil import shutil
import tempfile import tempfile
from kittens.ssh.main import get_connection_data from kittens.ssh.main import bootstrap_script, get_connection_data
from kitty.utils import SSHConnectionData from kitty.utils import SSHConnectionData
from . import BaseTest from . import BaseTest
from .shell_integration import basic_shell_env
class SSHTest(BaseTest): class SSHTest(BaseTest):
@ -39,11 +41,14 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
t('ssh -p 33 main', port=33) t('ssh -p 33 main', port=33)
def test_ssh_launcher_script(self): def test_ssh_launcher_script(self):
for sh in ('sh', 'zsh', 'bash', 'dash'): for sh in ('dash', 'zsh', 'bash', 'posh', 'sh'):
q = shutil.which(sh) q = shutil.which(sh)
if q: if q:
with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir: with self.subTest(sh=sh), tempfile.TemporaryDirectory() as tdir:
self.run_launcher(sh, tdir) script = bootstrap_script('echo TEST_DONE; return 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 run_launcher(self, sh, tdir): def check_bootstrap(self, home_dir, pty):
pass pty.wait_till(lambda: 'TEST_DONE' in pty.screen_contents())

View File

@ -4,35 +4,39 @@
# read the transmitted data from STDIN # read the transmitted data from STDIN
saved_tty_settings=$(command stty -g) saved_tty_settings=$(command stty -g)
command stty -echo command stty raw -echo
encoded_data_file=$(mktemp) encoded_data_file=$(mktemp)
cleanup_on_bootstrap_exit() { cleanup_on_bootstrap_exit() {
[[ ! -z "$encoded_data_file" ]] && command rm -f "$encoded_data_file" [ ! -z "$encoded_data_file" ] && command rm -f "$encoded_data_file"
[[ ! -z "$saved_tty_settings" ]] && command stty "$saved_tty_settings" [ ! -z "$saved_tty_settings" ] && command stty "$saved_tty_settings"
} }
trap 'cleanup_on_bootstrap_exit' EXIT trap 'cleanup_on_bootstrap_exit' EXIT
die() { echo "$*" >/dev/stderr; cleanup_on_bootstrap_exit; exit 1; }
data_started="n" data_started="n"
data_complete="n" data_complete="n"
pending_data="" pending_data=""
if [[ -z "$HOSTNAME" ]]; then if [ -z "$HOSTNAME" ]; then
hostname=$(hostname) hostname=$(hostname)
if [[ -z "$hostname" ]]; then hostname="_"; fi if [ -z "$hostname" ]; then hostname="_"; fi
else else
hostname=$(HOSTNAME) hostname="$HOSTNAME"
fi fi
# ensure $HOME is set # ensure $HOME is set
if [[ -z "$HOME" ]]; then HOME=~; fi if [ -z "$HOME" ]; then HOME=~; fi
# ensure $USER is set # ensure $USER is set
if [[ -z "$USER" ]]; then USER=$(whoami); fi if [ -z "$USER" ]; then USER=$(whoami); fi
# ask for the SSH data # ask for the SSH data
data_password="DATA_PASSWORD" data_password="DATA_PASSWORD"
password_filename="PASSWORD_FILENAME" password_filename="PASSWORD_FILENAME"
printf "\eP@kitty-ssh|$hostname:$password_filename:$data_password\e\\" pending_data=""
data_complete="n"
printf "\033P@kitty-ssh|%s:%s:%s\033\\" "$hostname" "$password_filename" "$data_password"
while IFS= read -r line; do while [ "$data_complete" = "n" ]; do
IFS= read -r line || die "Incomplete ssh data";
case "$line" in case "$line" in
*"KITTY_SSH_DATA_START") *"KITTY_SSH_DATA_START")
prefix=$(command expr "$line" : "\(.*\)KITTY_SSH_DATA_START") prefix=$(command expr "$line" : "\(.*\)KITTY_SSH_DATA_START")
@ -43,40 +47,39 @@ while IFS= read -r line; do
data_complete="y"; data_complete="y";
;; ;;
*) *)
if [[ "$data_started" == "y" ]]; then if [ "$data_started" = "y" ]; then
echo -n "$line" >> "$encoded_data_file" echo -n "$line" >> "$encoded_data_file"
else else
pending_data="$pending_data$line\n" pending_data="$pending_data$line\n"
fi fi
;; ;;
esac esac
if [[ "$data_complete" == "y" ]]; then break; fi
done done
command stty "$saved_tty_settings" command stty "$saved_tty_settings"
saved_tty_settings="" saved_tty_settings=""
if [[ ! -z "$pending_data" ]]; then if [ -n "$pending_data" ]; then
printf "\eP@kitty-echo|$(echo -n "$pending_data" | base64)\e\\" printf "\033P@kitty-echo|%s\033\\" "$(echo -n "$pending_data" | base64)"
fi fi
command base64 -d < "$encoded_data_file" | command tar xjf - --no-same-owner -C "$HOME" command base64 -d < "$encoded_data_file" | command tar xjf - --no-same-owner -C "$HOME"
rc=$? rc=$?
command rm -f "$encoded_data_file" command rm -f "$encoded_data_file"
encoded_data_file="" encoded_data_file=""
if [[ "$rc" != "0" ]]; then echo "Failed to extract data transmitted by ssh kitten over the TTY device" > /dev/stderr; exit 1; fi if [ "$rc" != "0" ]; then die "Failed to extract data transmitted by ssh kitten over the TTY device"; fi
# export TERMINFO
tname=".terminfo"
if [ -e "/usr/share/misc/terminfo.cdb" ]; then
# NetBSD requires this see https://github.com/kovidgoyal/kitty/issues/4622
tname=".terminfo.cdb"
fi
export TERMINFO="$HOME/$tname"
# compile terminfo for this system # compile terminfo for this system
if [[ -x "$(command -v tic)" ]]; then if [ -x "$(command -v tic)" ]; then
tname=".terminfo"
if [[ -e "/usr/share/misc/terminfo.cdb" ]]; then
# NetBSD requires this see https://github.com/kovidgoyal/kitty/issues/4622
tname=".terminfo.cdb"
fi
tic_out=$(command tic -x -o "$HOME/$tname" ".terminfo/kitty.terminfo" 2>&1) tic_out=$(command tic -x -o "$HOME/$tname" ".terminfo/kitty.terminfo" 2>&1)
rc=$? rc=$?
if [[ "$rc" != "0" ]]; then echo "$tic_out"; exit 1; fi if [ "$rc" != "0" ]; then die "$tic_out"; fi
export TERMINFO="$HOME/$tname"
fi fi
# If a command was passed to SSH execute it here # If a command was passed to SSH execute it here
EXEC_CMD EXEC_CMD