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

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


def load_env_file(path: Path) -> None:
    if not path.exists():
        return

    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()

        if not line or line.startswith("#") or "=" not in line:
            continue

        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip()

        if not key or key in os.environ:
            continue

        if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}:
            value = value[1:-1]

        os.environ[key] = value


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 parse_timeout_seconds(value: str) -> float:
    value = value.strip().lower()

    if value.endswith("ms"):
        return float(value[:-2]) / 1000
    if value.endswith("s"):
        return float(value[:-1])
    if value.endswith("m"):
        return float(value[:-1]) * 60
    if value.endswith("h"):
        return float(value[:-1]) * 3600

    return float(value)


def build_remote_script(container: str, workdir: str, command: str) -> str:
    q_container = shlex.quote(container)
    q_workdir = shlex.quote(workdir)
    q_command = shlex.quote(command)

    return f"""
set -Eeuo pipefail

if ! command -v docker >/dev/null 2>&1; then
  echo 'error: docker not found on remote host' >&2
  exit 127
fi

if ! docker inspect {q_container} >/dev/null 2>&1; then
  echo 'error: docker container not found: {container}' >&2
  exit 127
fi

if [ "$(docker inspect -f '{{{{.State.Running}}}}' {q_container})" != "true" ]; then
  echo 'error: docker container is not running: {container}' >&2
  exit 127
fi

docker exec -i \\
  -w {q_workdir} \\
  {q_container} \\
  bash -lc {q_command}
""".strip()


def build_ssh_command(
    *,
    ssh_target: str,
    ssh_port: str,
    ssh_opts: str,
) -> list[str]:
    ssh_cmd = [
        "ssh",
        "-p",
        ssh_port,
        "-o",
        "BatchMode=yes",
        "-o",
        "StrictHostKeyChecking=accept-new",
    ]

    if ssh_opts:
        ssh_cmd.extend(shlex.split(ssh_opts))

    ssh_cmd.extend([ssh_target, "bash", "-s", "--"])

    return ssh_cmd


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()