planet/scripts/docker_command.py

"""Run Docker CLI work with live evidence, a deadline, and process-group cleanup."""

import argparse
from contextlib import nullcontext
import os
from pathlib import Path
import re
import selectors
import signal
import subprocess
import sys
import time
from typing import TextIO


def safe_line(value: str) -> str:
    value = re.sub(r"\x1b\[[0-9;?]*[a-zA-Z]", "", value)
    return re.sub(r"(https?://)[^\s/@]+:[^\s/@]+@", r"\1***@", value)


def concise_line(value: str) -> str:
    value = safe_line(value).strip()
    if "Bind for " in value:
        return value[value.index("Bind for "):]
    return re.sub(r"sha256:([0-9a-f]{12})[0-9a-f]{52}", r"sha256:\1…", value)


def write_status(path: Path | None, value: str) -> None:
    if path is not None:
        temporary = path.with_suffix(".pending")
        temporary.write_text(value, encoding="utf-8")
        temporary.replace(path)


def stop_group(process: subprocess.Popen[bytes]) -> None:
    # A Compose plugin or credential helper can outlive the Docker CLI parent.
    for sig in (signal.SIGTERM, signal.SIGKILL):
        try:
            os.killpg(process.pid, sig)
        except ProcessLookupError:
            break
        if sig == signal.SIGTERM:
            time.sleep(0.2)
    process.wait()


def monitor(
    process: subprocess.Popen[bytes],
    args: argparse.Namespace,
    log: TextIO | None,
) -> int:
    started = time.monotonic()
    next_notice = started
    latest = args.label
    failure = ""
    pending = b""

    def emit(raw: bytes) -> None:
        nonlocal latest, failure
        line = safe_line(raw.decode("utf-8", errors="replace"))
        if not line.strip():
            return
        latest = concise_line(line)
        if re.search(r"error|failed|denied|timeout|refused|不存在|失败|被占用|未向宿主机", line, re.I):
            failure = latest
        print(line, flush=True)
        if log is not None:
            log.write(line + "\n")
            log.flush()

    with selectors.DefaultSelector() as selector:
        selector.register(process.stdout, selectors.EVENT_READ)
        while selector.get_map() or process.poll() is None:
            now = time.monotonic()
            if now - started >= args.timeout:
                if pending:
                    emit(pending)
                message = (
                    f"PLANET_DOCKER_COMMAND_TIMEOUT: {args.label}；"
                    f"已等待 {now - started:.1f}s；最近输出: {latest}"
                )
                write_status(args.status_file, f"{latest}（等待 {now - started:.0f} 秒后超时）")
                if log is not None:
                    log.write(message + "\n")
                    log.flush()
                return 124
            if now >= next_notice:
                detail = concise_line(pending.decode("utf-8", errors="replace")) or latest
                write_status(args.status_file, f"{detail} · {now - started:.0f}s")
                next_notice = now + args.heartbeat
            for key, _ in selector.select(timeout=min(0.1, args.heartbeat)):
                chunk = os.read(key.fd, 65536)
                if not chunk:
                    selector.unregister(key.fileobj)
                    continue
                parts = re.split(b"[\r\n]", pending + chunk)
                pending = parts.pop()
                for part in parts:
                    emit(part)
                write_status(args.status_file, concise_line(pending.decode("utf-8", errors="replace")) or latest)
        if pending:
            emit(pending)
    code = process.wait()
    write_status(args.status_file, (failure or latest) if code else "")
    return code if code >= 0 else 128 - code


def positive_seconds(value: str) -> float:
    seconds = float(value)
    if not 0 < seconds < float("inf"):
        raise argparse.ArgumentTypeError("时间必须为有限正数")
    return seconds


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--label", required=True)
    parser.add_argument("--timeout", type=positive_seconds, required=True)
    parser.add_argument("--heartbeat", type=positive_seconds, default=5)
    parser.add_argument("--log", type=Path)
    parser.add_argument("--status-file", type=Path)
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    command = args.command[1:] if args.command[:1] == ["--"] else args.command
    if not command:
        parser.error("缺少命令")
    log_context = nullcontext(None)
    if args.log:
        args.log.parent.mkdir(parents=True, exist_ok=True)
        descriptor = os.open(args.log, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        os.fchmod(descriptor, 0o600)
        log_context = os.fdopen(descriptor, "w", encoding="utf-8")
    with log_context as log:
        try:
            process = subprocess.Popen(
                command,
                stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        except OSError as error:
            write_status(args.status_file, safe_line(str(error)))
            return 127
        previous_handler = signal.signal(signal.SIGTERM, lambda *_: sys.exit(143))
        try:
            return monitor(process, args, log)
        except KeyboardInterrupt:
            return 130
        finally:
            stop_group(process)
            signal.signal(signal.SIGTERM, previous_handler)


if __name__ == "__main__":
    sys.exit(main())