diff --git a/pkgs/clan-cli/clan_cli/ssh/upload.py b/pkgs/clan-cli/clan_cli/ssh/upload.py index 955b4882c..7024ba29d 100644 --- a/pkgs/clan-cli/clan_cli/ssh/upload.py +++ b/pkgs/clan-cli/clan_cli/ssh/upload.py @@ -2,13 +2,84 @@ import tarfile from pathlib import Path from shlex import quote from tempfile import TemporaryDirectory +from typing import IO from clan_cli.cmd import Log, RunOpts -from clan_cli.cmd import run as run_local from clan_cli.errors import ClanError from clan_cli.ssh.host import Host +def unpack_archive_as_root( + host: Host, f: IO[bytes], local_src: Path, remote_dest: Path, dir_mode: int = 0o700 +) -> None: + if local_src.is_dir(): + cmd = 'rm -rf "$0" && mkdir -m "$1" -p "$0" && tar -C "$0" -xzf -' + elif local_src.is_file(): + cmd = 'rm -f "$0" && tar -C "$(dirname "$0")" -xzf -' + else: + msg = f"Unsupported source file type: {local_src}" + raise ClanError(msg) + + host.run( + [ + "sudo", + "-p", + f"Enter sudo password for {quote(host.host)}: ", + "--", + "bash", + "-c", + cmd, + str(remote_dest), + f"{dir_mode:o}", + ], + RunOpts( + input=f, + log=Log.BOTH, + ), + ) + + +def unpack_archive_as_user( + host: Host, f: IO[bytes], local_src: Path, remote_dest: Path, dir_mode: int = 0o700 +) -> None: + archive = host.run( + ["bash", "-c", "f=$(mktemp); echo $f; cat > $f"], + RunOpts( + input=f, + log=Log.BOTH, + ), + ).stdout.strip() + + if local_src.is_dir(): + cmd = 'trap "rm -f $0" EXIT; rm -rf "$1" && mkdir -m "$2" -p "$1" && tar -C "$1" -xzf "$0"' + elif local_src.is_file(): + cmd = 'trap "rm -f $0" EXIT; rm -f "$1" && tar -C "$(dirname "$1")" -xzf "$0"' + else: + msg = f"Unsupported source type: {local_src}" + raise ClanError(msg) + + # We also need some sort of locks in case we have multiple prompts + host.run( + [ + "sudo", + "-p", + f"Enter sudo password for {host.host}:\n", + "--", + "bash", + "-c", + cmd, + archive, + str(remote_dest), + f"{dir_mode:o}", + ], + tty=True, + opts=RunOpts( + log=Log.BOTH, + prefix="", + ), + ) + + def upload( host: Host, local_src: Path, @@ -89,33 +160,22 @@ def upload( with local_src.open("rb") as f: tar.addfile(tarinfo, f) - sudo = "" - if host.user != "root": - sudo = "sudo -- " - - cmd = None - if local_src.is_dir(): - cmd = 'rm -rf "$0" && mkdir -m "$1" -p "$0" && tar -C "$0" -xzf -' - elif local_src.is_file(): - cmd = 'rm -f "$0" && tar -C "$(dirname "$0")" -xzf -' - else: - msg = f"Unsupported source type: {local_src}" - raise ClanError(msg) - # TODO accept `input` to be an IO object instead of bytes so that we don't have to read the tarfile into memory. with tar_path.open("rb") as f: - run_local( - [ - *host.ssh_cmd(), - "--", - f"{sudo}bash -c {quote(cmd)}", - str(remote_dest), - f"{dir_mode:o}", - ], - RunOpts( - input=f.read(), - log=Log.BOTH, - prefix=host.command_prefix, - needs_user_terminal=True, - ), - ) + if host.user == "root": + unpack_archive_as_root( + host, + f, + local_src, + remote_dest, + dir_mode=dir_mode, + ) + else: + # For sudo we need to split the upload into two steps + unpack_archive_as_user( + host, + f, + local_src, + remote_dest, + dir_mode=dir_mode, + ) diff --git a/pkgs/clan-cli/clan_cli/tests/sshd.py b/pkgs/clan-cli/clan_cli/tests/sshd.py index caf6414f5..68b045835 100644 --- a/pkgs/clan-cli/clan_cli/tests/sshd.py +++ b/pkgs/clan-cli/clan_cli/tests/sshd.py @@ -80,6 +80,16 @@ exec {bash} -l "${{@}}" fake_sudo.write_text( f"""#!{bash} +# skip over every sudo option +for arg in "${{@}}"; do + if [[ "$arg" == "-p" ]]; then + shift + shift + continue + fi + break +done + exec "${{@}}" """ )