Only use bz2 compression if bzip2 is available

This commit is contained in:
Kovid Goyal 2022-03-07 20:44:26 +05:30
parent 855718b179
commit 5e457da30b
No known key found for this signature in database
GPG Key ID: 06BC317B515ACE7C
3 changed files with 11 additions and 6 deletions

View File

@ -62,7 +62,7 @@ def serialize_env(env: Dict[str, str], base_env: Dict[str, str]) -> bytes:
return '\n'.join(lines).encode('utf-8') return '\n'.join(lines).encode('utf-8')
def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes: def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: str = 'gz') -> bytes:
def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo: def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo:
tarinfo.uname = tarinfo.gname = '' tarinfo.uname = tarinfo.gname = ''
@ -117,7 +117,7 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
env['KITTY_LOGIN_CWD'] = ssh_opts.cwd env['KITTY_LOGIN_CWD'] = ssh_opts.cwd
env_script = serialize_env(env, base_env) env_script = serialize_env(env, base_env)
buf = io.BytesIO() buf = io.BytesIO()
with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf: with tarfile.open(mode=f'w:{compression}', fileobj=buf, encoding='utf-8') as tf:
rd = ssh_opts.remote_dir.rstrip('/') rd = ssh_opts.remote_dir.rstrip('/')
for ci in ssh_opts.copy.values(): for ci in ssh_opts.copy.values():
tf.add(ci.local_path, arcname=ci.arcname, filter=filter_from_globs(*ci.exclude_patterns)) tf.add(ci.local_path, arcname=ci.arcname, filter=filter_from_globs(*ci.exclude_patterns))
@ -147,6 +147,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
pwfilename = md['pwfile'] pwfilename = md['pwfile']
username = md['user'] username = md['user']
rq_id = md['id'] rq_id = md['id']
compression = md['compression']
except Exception: except Exception:
traceback.print_exc() traceback.print_exc()
yield fmt_prefix('!invalid ssh data request message') yield fmt_prefix('!invalid ssh data request message')
@ -169,7 +170,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
resolved_ssh_opts = options_for_host(hostname, username, ssh_opts, cli_hostname, cli_uname) resolved_ssh_opts = options_for_host(hostname, username, ssh_opts, cli_hostname, cli_uname)
resolved_ssh_opts.copy = {k: CopyInstruction(*v) for k, v in resolved_ssh_opts.copy.items()} resolved_ssh_opts.copy = {k: CopyInstruction(*v) for k, v in resolved_ssh_opts.copy.items()}
try: try:
data = make_tarfile(resolved_ssh_opts, env_data['env']) data = make_tarfile(resolved_ssh_opts, env_data['env'], compression)
except Exception: except Exception:
traceback.print_exc() traceback.print_exc()
yield fmt_prefix('!error while gathering ssh data') yield fmt_prefix('!error while gathering ssh data')

View File

@ -60,7 +60,7 @@ def dcs_to_kitty(type, payload):
def send_data_request(): def send_data_request():
hostname = os.environ.get('HOSTNAME') or os.uname().nodename hostname = os.environ.get('HOSTNAME') or os.uname().nodename
write_all(tty_fd, dcs_to_kitty( write_all(tty_fd, dcs_to_kitty(
'ssh', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:pw=DATA_PASSWORD'.format(hostname, getpass.getuser()))) 'ssh', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:compression=bz2:pw=DATA_PASSWORD'.format(hostname, getpass.getuser())))
def debug(msg): def debug(msg):

View File

@ -120,7 +120,9 @@ login_cwd=""
init_tty && trap "cleanup_on_bootstrap_exit" EXIT init_tty && trap "cleanup_on_bootstrap_exit" EXIT
if [ "$tty_ok" = "y" ]; then if [ "$tty_ok" = "y" ]; then
dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":pw="DATA_PASSWORD"" compression="gz"
command -v "bzip2" > /dev/null 2> /dev/null && compression="bz2"
dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":compression="$compression":pw="DATA_PASSWORD""
record_separator=$(printf "\036") record_separator=$(printf "\036")
fi fi
@ -157,7 +159,9 @@ untar_and_read_env() {
tdir=$(command mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX") tdir=$(command mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX")
[ $? = 0 ] || die "Creating temp directory failed" [ $? = 0 ] || die "Creating temp directory failed"
read_n_bytes_from_tty "$1" | base64_decode | command tar xpjf - -C "$tdir" cflag="j"
[ "$compression" = "gz" ] && cflag="z"
read_n_bytes_from_tty "$1" | base64_decode | command tar "x${cflag}pf" - -C "$tdir"
data_file="$tdir/data.sh" data_file="$tdir/data.sh"
[ -f "$data_file" ] && . "$data_file" [ -f "$data_file" ] && . "$data_file"
[ -z "$KITTY_SSH_KITTEN_DATA_DIR" ] && die "Failed to read SSH data from tty" [ -z "$KITTY_SSH_KITTEN_DATA_DIR" ] && die "Failed to read SSH data from tty"