Mini-Nav/scripts/remote_docker_run.py

#!/usr/bin/env python3
from __future__ import annotations

import os
import shlex
import subprocess
from pathlib import Path

import typer

from common import (
    build_remote_script,
    build_ssh_command,
    load_env_file,
    parse_timeout_seconds,
)

app = typer.Typer(
    help="Run commands inside a Docker container on a remote host over SSH.",
    no_args_is_help=True,
)


def required_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        typer.secho(
            f"error: missing required env: {name}",
            fg=typer.colors.RED,
            err=True,
        )
        raise typer.Exit(2)
    return value


def build_user_command(command_parts: list[str]) -> str:
    if command_parts and command_parts[0] == "--":
        command_parts = command_parts[1:]

    if not command_parts:
        typer.secho("error: missing command", fg=typer.colors.RED, err=True)
        raise typer.Exit(2)

    # 支持：
    #   run -- "pwd && ls -la"
    #   run -- pwd
    #   run -- uv run pytest -q
    if len(command_parts) == 1:
        return command_parts[0]

    return " ".join(shlex.quote(part) for part in command_parts)


def print_remote_info(
    *,
    ssh_target: str,
    container: str,
    workdir: str,
    command: str,
) -> None:
    typer.echo(f"[remote] target    = {ssh_target}")
    typer.echo(f"[remote] container = {container}")
    typer.echo(f"[remote] workdir   = {workdir}")
    typer.echo(f"[remote] command   = {command}")


def run_ssh_script(
    *,
    ssh_cmd: list[str],
    remote_script: str,
    timeout_raw: str,
) -> int:
    try:
        # 关键点：
        # 1. 手动统一为 LF
        # 2. 以 bytes 写入 stdin，避免 Windows text mode 把 \n 转为 \r\n
        script_bytes = (
            remote_script.replace("\r\n", "\n").replace("\r", "\n").encode("utf-8")
        )

        result = subprocess.run(
            ssh_cmd,
            input=script_bytes,
            check=False,
            timeout=parse_timeout_seconds(timeout_raw) if timeout_raw else None,
        )
        return result.returncode

    except subprocess.TimeoutExpired:
        typer.secho(
            f"error: remote command timed out: {timeout_raw}",
            fg=typer.colors.RED,
            err=True,
        )
        return 124


@app.command(
    context_settings={"allow_extra_args": True, "ignore_unknown_options": True}
)
def run(
    ctx: typer.Context,
    env_file: Path = typer.Option(
        Path(".env"),
        "--env-file",
        help="Path to remote environment file.",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Print generated ssh command and remote script without executing it.",
    ),
) -> None:
    """
    Run an arbitrary command inside the configured remote Docker container.

    Example:
      uv run python scripts/remote_docker_run.py run -- "uv run pytest -q"
    """
    load_env_file(env_file)

    user_command = build_user_command(list(ctx.args))

    ssh_target = required_env("REMOTE_SSH_TARGET")
    container = required_env("REMOTE_DOCKER_CONTAINER")
    workdir = required_env("REMOTE_WORKDIR")

    ssh_port = os.environ.get("REMOTE_SSH_PORT", "22")
    ssh_opts = os.environ.get("REMOTE_SSH_OPTS", "")
    timeout_raw = os.environ.get("REMOTE_TIMEOUT", "")

    remote_script = build_remote_script(
        container=container,
        workdir=workdir,
        command=user_command,
    )

    ssh_cmd = build_ssh_command(
        ssh_target=ssh_target,
        ssh_port=ssh_port,
        ssh_opts=ssh_opts,
    )

    print_remote_info(
        ssh_target=ssh_target,
        container=container,
        workdir=workdir,
        command=user_command,
    )

    if dry_run:
        typer.echo()
        typer.echo("[dry-run] ssh argv:")
        typer.echo(" ".join(shlex.quote(part) for part in ssh_cmd))

        typer.echo()
        typer.echo("[dry-run] remote script sent via stdin:")
        typer.echo(remote_script)

        raise typer.Exit(0)

    exit_code = run_ssh_script(
        ssh_cmd=ssh_cmd,
        remote_script=remote_script,
        timeout_raw=timeout_raw,
    )
    raise typer.Exit(exit_code)


@app.command()
def check(
    env_file: Path = typer.Option(
        Path(".env"),
        "--env-file",
        help="Path to remote environment file.",
    ),
    dry_run: bool = typer.Option(
        False,
        "--dry-run",
        help="Print generated ssh command and remote script without executing it.",
    ),
) -> None:
    """
    Check whether the configured remote Docker container is reachable.
    """
    load_env_file(env_file)

    ssh_target = required_env("REMOTE_SSH_TARGET")
    container = required_env("REMOTE_DOCKER_CONTAINER")
    workdir = required_env("REMOTE_WORKDIR")

    ssh_port = os.environ.get("REMOTE_SSH_PORT", "22")
    ssh_opts = os.environ.get("REMOTE_SSH_OPTS", "")
    timeout_raw = os.environ.get("REMOTE_TIMEOUT", "")

    check_command = "pwd && python --version || true && ls -la"

    remote_script = build_remote_script(
        container=container,
        workdir=workdir,
        command=check_command,
    )

    ssh_cmd = build_ssh_command(
        ssh_target=ssh_target,
        ssh_port=ssh_port,
        ssh_opts=ssh_opts,
    )

    print_remote_info(
        ssh_target=ssh_target,
        container=container,
        workdir=workdir,
        command=check_command,
    )

    if dry_run:
        typer.echo()
        typer.echo("[dry-run] ssh argv:")
        typer.echo(" ".join(shlex.quote(part) for part in ssh_cmd))

        typer.echo()
        typer.echo("[dry-run] remote script sent via stdin:")
        typer.echo(remote_script)

        raise typer.Exit(0)

    exit_code = run_ssh_script(
        ssh_cmd=ssh_cmd,
        remote_script=remote_script,
        timeout_raw=timeout_raw,
    )
    raise typer.Exit(exit_code)


if __name__ == "__main__":
    app()