zkxw/测试环境部署/deploy_test_env.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试环境日常发版脚本（裸机：Nginx 10001 + backend 127.0.0.1:10002）

在本机仓库根目录或任意目录执行均可（脚本会定位仓库根）：

  python3 测试环境部署/deploy_test_env.py              # 前端打包 + 后端同步 + 重启
  python3 测试环境部署/deploy_test_env.py --frontend-only
  python3 测试环境部署/deploy_test_env.py --backend-only
  python3 测试环境部署/deploy_test_env.py --verify
  python3 测试环境部署/deploy_test_env.py --upload-env  # 额外上传 backend/.env → 远端 .env.test
  python3 测试环境部署/deploy_test_env.py --pip         # 远端重新 pip install -r requirements.txt

SSH 密码默认读取仓库根「测试环境机器的账号密码」（root密码：...），
也可用环境变量 ZKXW_TEST_SSH_PASS，或 --password / --password-file。
勿把密码写进本脚本或提交到 git。
"""

from __future__ import annotations

import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from typing import List, Optional, Sequence

# ---------------------------------------------------------------------------
# 默认拓扑（与 docs/superpowers/specs/2026-08-06-test-env-bare-metal-deploy-design.md 一致）
# ---------------------------------------------------------------------------
DEFAULT_HOST = "101.35.245.122"
DEFAULT_USER = "root"
REMOTE_ROOT = "/root/zkxw"
REMOTE_BACKEND = f"{REMOTE_ROOT}/backend"
REMOTE_DIST = f"{REMOTE_ROOT}/web-admin/dist"
BACKEND_HOST = "127.0.0.1"
BACKEND_PORT = "10002"
PUBLIC_HTTP_PORT = "10001"
DEFAULT_HEALTH_TIMEOUT_SEC = 60
DEFAULT_HEALTH_INTERVAL_SEC = 2

BACKEND_EXCLUDES = (
    ".venv",
    "venv",
    "__pycache__",
    "*.pyc",
    "logs",
    "graphify-out",
    ".env",
    ".env.*",
    ".pytest_cache",
    ".pycache",
    ".DS_Store",
)


def repo_root() -> Path:
    here = Path(__file__).resolve().parent
    # 测试环境部署/ → 仓库根
    return here.parent


def die(msg: str, code: int = 1) -> None:
    print(f"[ERROR] {msg}", file=sys.stderr)
    raise SystemExit(code)


def info(msg: str) -> None:
    print(f"[INFO] {msg}")


def run(
    cmd: Sequence[str],
    *,
    cwd: Optional[Path] = None,
    env: Optional[dict] = None,
    check: bool = True,
) -> subprocess.CompletedProcess:
    print(f"[CMD] {' '.join(cmd)}")
    return subprocess.run(
        list(cmd),
        cwd=str(cwd) if cwd else None,
        env=env,
        check=check,
    )


def load_password(args: argparse.Namespace, root: Path) -> Optional[str]:
    if args.password:
        return args.password
    env_pass = os.environ.get("ZKXW_TEST_SSH_PASS") or os.environ.get("SSHPASS")
    if env_pass:
        return env_pass
    path = Path(args.password_file) if args.password_file else root / "测试环境机器的账号密码"
    if not path.is_file():
        return None
    text = path.read_text(encoding="utf-8", errors="replace")
    m = re.search(r"^root密码[：:]\s*(.+)\s*$", text, re.M)
    if m:
        return m.group(1).strip()
    # 兜底：第二行非空内容
    lines = [ln.strip() for ln in text.splitlines() if ln.strip() and not ln.strip().startswith("#")]
    if len(lines) >= 2 and "密码" in lines[1]:
        return re.sub(r"^.*?[：:]\s*", "", lines[1]).strip()
    return None


def require_bin(name: str) -> str:
    path = shutil.which(name)
    if not path:
        die(f"未找到命令 `{name}`，请先安装并加入 PATH")
    return path


def ssh_argv(password: Optional[str]) -> List[str]:
    """ssh 可执行前缀（含可选 sshpass）。"""
    ssh = require_bin("ssh")
    opts = ["-o", "StrictHostKeyChecking=accept-new", "-o", "RequestTTY=no"]
    if password:
        require_bin("sshpass")
        return ["sshpass", "-e", ssh, *opts]
    return [ssh, "-o", "BatchMode=yes", *opts]


def scp_argv(password: Optional[str]) -> List[str]:
    scp = require_bin("scp")
    opts = ["-o", "StrictHostKeyChecking=accept-new", "-o", "RequestTTY=no"]
    if password:
        require_bin("sshpass")
        return ["sshpass", "-e", scp, *opts]
    return [scp, "-o", "BatchMode=yes", *opts]


def with_sshpass_env(password: Optional[str], base_env: Optional[dict] = None) -> dict:
    env = dict(base_env or os.environ)
    if password:
        env["SSHPASS"] = password
    return env


def remote_target(args: argparse.Namespace) -> str:
    return f"{args.user}@{args.host}"


def ssh_run(args: argparse.Namespace, password: Optional[str], remote_cmd: str) -> None:
    cmd = [*ssh_argv(password), remote_target(args), remote_cmd]
    run(cmd, env=with_sshpass_env(password))


def rsync_to(
    args: argparse.Namespace,
    password: Optional[str],
    local: str,
    remote_path: str,
    *,
    delete: bool = False,
    excludes: Sequence[str] = (),
) -> None:
    require_bin("rsync")
    # rsync -e 需要单个 shell 字符串
    ssh_e = " ".join(ssh_argv(password))
    cmd: List[str] = ["rsync", "-avz"]
    if delete:
        cmd.append("--delete")
    for ex in excludes:
        cmd.extend(["--exclude", ex])
    cmd.extend(["-e", ssh_e, local, f"{remote_target(args)}:{remote_path}"])
    run(cmd, env=with_sshpass_env(password))


def build_frontend(root: Path) -> None:
    web = root / "web-admin"
    if not (web / "package.json").is_file():
        die(f"未找到前端目录: {web}")
    require_bin("pnpm")
    env = dict(os.environ)
    # Vite 大项目默认堆易 OOM
    env.setdefault("NODE_OPTIONS", "--max-old-space-size=8192")
    run(["pnpm", "build:test"], cwd=web, env=env)
    index = web / "dist" / "index.html"
    if not index.is_file():
        die(f"打包未生成 {index}")


def sync_frontend(args: argparse.Namespace, password: Optional[str], root: Path) -> None:
    dist = root / "web-admin" / "dist"
    if not (dist / "index.html").is_file():
        die("本地无 web-admin/dist/index.html，请先打包或去掉 --skip-build")
    rsync_to(
        args,
        password,
        f"{dist}/",
        f"{REMOTE_DIST}/",
        delete=True,
    )
    # Nginx 读 /root 下静态文件需要可穿越权限
    ssh_run(
        args,
        password,
        f"chmod 755 /root && chmod -R a+rX {REMOTE_ROOT}/web-admin",
    )


def sync_backend(args: argparse.Namespace, password: Optional[str], root: Path) -> None:
    backend = root / "backend"
    if not (backend / "start-test.sh").is_file():
        die(f"未找到后端: {backend}/start-test.sh")
    rsync_to(
        args,
        password,
        f"{backend}/",
        f"{REMOTE_BACKEND}/",
        excludes=BACKEND_EXCLUDES,
    )


def upload_env(args: argparse.Namespace, password: Optional[str], root: Path) -> None:
    src = root / "backend" / ".env"
    if not src.is_file():
        die(f"本机无 {src}，无法上传 .env.test")
    require_bin("scp")
    with tempfile.NamedTemporaryFile(prefix="zkxw-env-", suffix=".test", delete=False) as tmp:
        tmp_path = Path(tmp.name)
    try:
        tmp_path.write_bytes(src.read_bytes())
        cmd = [
            *scp_argv(password),
            str(tmp_path),
            f"{remote_target(args)}:{REMOTE_BACKEND}/.env.test",
        ]
        run(cmd, env=with_sshpass_env(password))
        info(f"已上传 {src} → {REMOTE_BACKEND}/.env.test")
    finally:
        tmp_path.unlink(missing_ok=True)


def remote_pip(args: argparse.Namespace, password: Optional[str]) -> None:
    ssh_run(
        args,
        password,
        f"cd {REMOTE_BACKEND} && ./.venv/bin/pip install -r requirements.txt",
    )


def restart_backend(args: argparse.Namespace, password: Optional[str]) -> None:
    ssh_run(
        args,
        password,
        f"cd {REMOTE_BACKEND} && chmod +x start-test.sh && "
        f"HOST={BACKEND_HOST} PORT={BACKEND_PORT} ./start-test.sh restart",
    )


def _health_ok_via_ssh(args: argparse.Namespace, password: Optional[str]) -> bool:
    """本机回环健康检查（不依赖公网 / Nginx）。"""
    remote = (
        f"curl -sf --connect-timeout 2 --max-time 5 "
        f"http://{BACKEND_HOST}:{BACKEND_PORT}/health"
    )
    cmd = [*ssh_argv(password), remote_target(args), remote]
    cp = subprocess.run(
        cmd,
        env=with_sshpass_env(password),
        capture_output=True,
        text=True,
    )
    return cp.returncode == 0 and "healthy" in (cp.stdout or "")


def _health_ok_via_nginx(args: argparse.Namespace) -> bool:
    """经 Nginx 公开入口检查（与验收路径一致）。"""
    url = f"http://{args.host}:{PUBLIC_HTTP_PORT}/health"
    cp = subprocess.run(
        ["curl", "-sf", "--connect-timeout", "3", "--max-time", "8", url],
        capture_output=True,
        text=True,
    )
    return cp.returncode == 0 and "healthy" in (cp.stdout or "")


def wait_for_health(
    args: argparse.Namespace,
    password: Optional[str],
    *,
    prefer_ssh: bool,
) -> None:
    """
    重启后轮询健康检查，成功后再验收。
    prefer_ssh=True：优先 SSH 打 127.0.0.1:10002（刚重启时更准）。
    否则优先 Nginx :10001/health。
    """
    require_bin("curl")
    timeout = max(1, int(args.health_timeout))
    interval = max(0.5, float(args.health_interval))
    deadline = time.monotonic() + timeout
    attempt = 0
    info(f"等待后端健康（最多 {timeout}s，间隔 {interval}s）…")
    while True:
        attempt += 1
        ssh_ok = _health_ok_via_ssh(args, password)
        nginx_ok = _health_ok_via_nginx(args)
        ready = ssh_ok if prefer_ssh else (nginx_ok or ssh_ok)
        detail = f"ssh={'ok' if ssh_ok else 'fail'} nginx={'ok' if nginx_ok else 'fail'}"

        if ready:
            info(f"健康检查通过（第 {attempt} 次，{detail}）")
            # 本机已通但 Nginx 仍 502 时，再短等公开入口
            if prefer_ssh and not nginx_ok:
                info("本机 health 已通，等待 Nginx /health …")
                nginx_deadline = time.monotonic() + min(15.0, float(timeout))
                while time.monotonic() < nginx_deadline:
                    if _health_ok_via_nginx(args):
                        info("Nginx /health 已就绪")
                        break
                    time.sleep(interval)
                else:
                    info("Nginx /health 仍未就绪，将继续验收（可能失败）")
            return

        if time.monotonic() >= deadline:
            die(
                f"等待健康检查超时（{timeout}s，共 {attempt} 次，{detail}）。"
                f"请 SSH 查看: cd {REMOTE_BACKEND} && ./start-test.sh logs",
                code=3,
            )
        print(f"[WAIT] 未就绪（第 {attempt} 次，{detail}），{interval}s 后重试…")
        time.sleep(interval)


def verify(args: argparse.Namespace) -> None:
    require_bin("curl")
    base = f"http://{args.host}:{PUBLIC_HTTP_PORT}"
    checks = [
        ("首页", f"{base}/", True),
        ("health", f"{base}/health", False),
        ("login-rsa", f"{base}/api/web/auth/login-rsa-public-key", True),
    ]
    ok = True
    for name, url, code_only in checks:
        if code_only:
            cp = subprocess.run(
                ["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}", "--connect-timeout", "10", url],
                capture_output=True,
                text=True,
            )
            code = (cp.stdout or "").strip()
            print(f"[VERIFY] {name}: HTTP {code}  ({url})")
            if code != "200":
                ok = False
        else:
            cp = subprocess.run(
                ["curl", "-s", "--connect-timeout", "10", url],
                capture_output=True,
                text=True,
            )
            body = (cp.stdout or "").strip()
            print(f"[VERIFY] {name}: {body[:120]}  ({url})")
            if "healthy" not in body:
                ok = False

    # 10002 应对外不可达
    cp = subprocess.run(
        [
            "curl",
            "-s",
            "-o",
            "/dev/null",
            "-w",
            "%{http_code}",
            "--connect-timeout",
            "5",
            f"http://{args.host}:{BACKEND_PORT}/health",
        ],
        capture_output=True,
        text=True,
    )
    code = (cp.stdout or "").strip()
    if code in ("", "000") or cp.returncode != 0:
        print(f"[VERIFY] 外网 10002: 不可达（符合预期） code={code or 'n/a'}")
    else:
        print(f"[VERIFY] 外网 10002: 意外可达 HTTP {code}（应仅本机监听）")
        ok = False

    if not ok:
        die("验收未全部通过", code=2)
    info("验收通过")


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="ZKXW 测试环境发版（前端 build:test + rsync + backend restart）",
    )
    p.add_argument("--host", default=DEFAULT_HOST, help=f"测试机 IP（默认 {DEFAULT_HOST}）")
    p.add_argument("--user", default=DEFAULT_USER, help=f"SSH 用户（默认 {DEFAULT_USER}）")
    p.add_argument("--password", default=None, help="SSH 密码（不推荐；优先用密码文件或环境变量）")
    p.add_argument(
        "--password-file",
        default=None,
        help="密码文件路径（默认：仓库根/测试环境机器的账号密码）",
    )
    g = p.add_mutually_exclusive_group()
    g.add_argument("--frontend-only", action="store_true", help="仅前端打包并上传")
    g.add_argument("--backend-only", action="store_true", help="仅同步后端并重启")
    g.add_argument("--verify", action="store_true", help="仅跑验收 curl")
    p.add_argument("--skip-build", action="store_true", help="前端跳过 pnpm build:test，直接上传已有 dist")
    p.add_argument("--skip-restart", action="store_true", help="同步后端后不重启")
    p.add_argument("--upload-env", action="store_true", help="上传本机 backend/.env 为远端 .env.test")
    p.add_argument("--pip", action="store_true", help="远端执行 pip install -r requirements.txt")
    p.add_argument("--no-verify", action="store_true", help="发版后不自动验收")
    p.add_argument(
        "--health-timeout",
        type=int,
        default=DEFAULT_HEALTH_TIMEOUT_SEC,
        help=f"重启后等待健康检查的最长时间（秒，默认 {DEFAULT_HEALTH_TIMEOUT_SEC}）",
    )
    p.add_argument(
        "--health-interval",
        type=float,
        default=DEFAULT_HEALTH_INTERVAL_SEC,
        help=f"健康检查轮询间隔（秒，默认 {DEFAULT_HEALTH_INTERVAL_SEC}）",
    )
    return p.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None) -> None:
    args = parse_args(argv)
    root = repo_root()
    os.chdir(root)
    info(f"仓库根目录: {root}")

    if args.verify:
        verify(args)
        return

    password = load_password(args, root)
    if not password and not os.environ.get("SSH_AUTH_SOCK"):
        info("未找到 SSH 密码；将尝试密钥登录（BatchMode）")
    elif password:
        info("已加载 SSH 密码（来自参数/环境变量/密码文件）")

    do_fe = not args.backend_only
    do_be = not args.frontend_only
    did_restart = False

    if do_fe:
        if not args.skip_build:
            info("开始前端打包 pnpm build:test …")
            build_frontend(root)
        else:
            info("跳过打包，使用现有 web-admin/dist")
        info("上传前端 dist …")
        sync_frontend(args, password, root)

    if do_be:
        info("同步后端源码 …")
        sync_backend(args, password, root)
        if args.upload_env:
            upload_env(args, password, root)
        if args.pip:
            info("远端 pip install …")
            remote_pip(args, password)
        if not args.skip_restart:
            info("重启后端 …")
            restart_backend(args, password)
            did_restart = True
        else:
            info("已跳过后端重启")

    if not args.no_verify:
        if did_restart:
            wait_for_health(args, password, prefer_ssh=True)
        else:
            # 仅前端发版：仍确认当前后端/Nginx 可用再验收
            wait_for_health(args, password, prefer_ssh=False)
        info("发版后验收 …")
        verify(args)

    info("完成")


if __name__ == "__main__":
    try:
        main()
    except subprocess.CalledProcessError as e:
        die(f"命令失败（exit {e.returncode}）: {e.cmd}")
    except KeyboardInterrupt:
        die("已中断", code=130)