Nexus/scripts/prod_probe.py

#!/usr/bin/env python3
"""Production probe — public + authenticated API checks (no secrets printed).

Modes:
  external — curl public BASE (may fail on some ISP routes)
  origin   — SSH to host, curl http://127.0.0.1:8600
  auto     — both; origin auth checks are required, external public is best-effort
"""
from __future__ import annotations

import argparse
import json
import os
import shlex
import subprocess
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


@dataclass
class Config:
    base_external: str
    ssh_host: str
    ssh_key: str | None
    origin_url: str
    admin_user: str
    probe_password: str | None
    probe_client_ip: str | None
    timeout: int


def _load_secrets_env() -> None:
    secrets = os.path.join(ROOT, "deploy", "nexus-1panel.secrets.sh")
    if not os.path.isfile(secrets):
        return
    # Parse KEY=VAL / export KEY=VAL (no shell source — avoids arbitrary code)
    with open(secrets, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if line.startswith("export "):
                line = line[7:].strip()
            if "=" not in line:
                continue
            key, _, val = line.partition("=")
            key = key.strip()
            val = val.strip().strip("'\"")
            if key and key not in os.environ:
                os.environ[key] = val


def build_config() -> Config:
    _load_secrets_env()
    pw = os.environ.get("NEXUS_TEST_ADMIN_PASSWORD", "").strip() or None
    probe_ip = os.environ.get("NEXUS_PROBE_CLIENT_IP", "").strip() or None
    return Config(
        base_external=os.environ.get("NEXUS_PROBE_BASE", "https://api.synaglobal.vip").rstrip("/"),
        ssh_host=os.environ.get("NEXUS_SSH", "nexus"),
        ssh_key=os.environ.get("NEXUS_SSH_KEY") or None,
        origin_url=os.environ.get("NEXUS_PROBE_ORIGIN", "http://127.0.0.1:8600").rstrip("/"),
        admin_user=os.environ.get("NEXUS_TEST_ADMIN_USER", "admin"),
        probe_password=pw,
        probe_client_ip=probe_ip,
        timeout=int(os.environ.get("NEXUS_PROBE_TIMEOUT", "20")),
    )


def ssh_base(cfg: Config) -> list[str]:
    cmd = ["ssh", "-o", "BatchMode=yes", "-o", f"ConnectTimeout={cfg.timeout}"]
    if cfg.ssh_key:
        cmd.extend(["-i", cfg.ssh_key])
    cmd.append(cfg.ssh_host)
    return cmd


def fetch_env_from_origin(cfg: Config, key: str) -> str | None:
    remote = (
        "sudo bash -c "
        + shlex.quote(
            "for f in /opt/nexus/docker/.env.prod /opt/nexus/.env "
            "/www/wwwroot/api.synaglobal.vip/.env; do "
            'if [[ -f "$f" ]]; then '
            f'grep -E "^{key}=" "$f" 2>/dev/null | head -1 | cut -d= -f2- | tr -d \'"\'; '
            "exit 0; fi; done"
        )
    )
    try:
        r = subprocess.run(
            ssh_base(cfg) + [remote],
            capture_output=True,
            text=True,
            timeout=cfg.timeout + 5,
            check=False,
        )
    except (subprocess.TimeoutExpired, OSError):
        return None
    out = (r.stdout or "").strip()
    return out if out else None


def fetch_password_from_origin(cfg: Config) -> str | None:
    if cfg.probe_password:
        return cfg.probe_password
    return fetch_env_from_origin(cfg, "NEXUS_TEST_ADMIN_PASSWORD")


def fetch_probe_client_ip(cfg: Config) -> str | None:
    if cfg.probe_client_ip:
        return cfg.probe_client_ip
    from_env = fetch_env_from_origin(cfg, "NEXUS_PROBE_CLIENT_IP")
    if from_env:
        return from_env
    # First IP from login_manual_ips (login allowlist) for origin probe via X-Real-IP
    remote = (
        "sudo docker exec nexus-prod-nexus-1 python3 -c "
        + shlex.quote(
            "import asyncio\n"
            "from server.infrastructure.database.session import AsyncSessionLocal\n"
            "from sqlalchemy import text\n"
            "async def main():\n"
            "    async with AsyncSessionLocal() as s:\n"
            "        for k in ('login_manual_ips','login_allowed_ips','login_subscription_ips'):\n"
            "            r = await s.execute(text('SELECT value FROM settings WHERE `key` = :k'), {'k': k})\n"
            "            row = r.first()\n"
            "            if row and row[0]:\n"
            "                for part in str(row[0]).split(','):\n"
            "                    ip = part.strip().split('/')[0].strip()\n"
            "                    if ip:\n"
            "                        print(ip); return\n"
            "asyncio.run(main())\n"
        )
    )
    try:
        r = subprocess.run(
            ssh_base(cfg) + [remote],
            capture_output=True,
            text=True,
            timeout=cfg.timeout + 15,
            check=False,
        )
    except (subprocess.TimeoutExpired, OSError):
        return None
    out = (r.stdout or "").strip().splitlines()
    return out[0].strip() if out else None


def http_request(
    url: str,
    *,
    method: str = "GET",
    headers: dict[str, str] | None = None,
    body: dict[str, Any] | None = None,
    timeout: int = 20,
) -> tuple[int, str]:
    data = None
    req_headers = dict(headers or {})
    if body is not None:
        data = json.dumps(body).encode("utf-8")
        req_headers.setdefault("Content-Type", "application/json")
    req = urllib.request.Request(url, data=data, headers=req_headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as e:
        payload = e.read().decode("utf-8", errors="replace")
        return e.code, payload


def origin_curl(
    cfg: Config,
    path: str,
    *,
    method: str = "GET",
    token: str | None = None,
    body: dict[str, Any] | None = None,
    client_ip: str | None = None,
) -> tuple[int, str]:
    url = f"{cfg.origin_url}{path}"
    parts = [
        "curl",
        "-sS",
        "--max-time",
        str(cfg.timeout),
        "-w",
        "\\n__HTTP_CODE__%{http_code}",
        "-X",
        method,
    ]
    if client_ip:
        parts.extend(
            ["-H", f"X-Real-IP: {client_ip}", "-H", f"X-Forwarded-For: {client_ip}"]
        )
    if token:
        parts.extend(["-H", f"Authorization: Bearer {token}"])
    if body is not None:
        parts.extend(["-H", "Content-Type: application/json", "-d", json.dumps(body)])
    parts.append(url)
    remote = " ".join(shlex.quote(p) for p in parts)
    r = subprocess.run(
        ssh_base(cfg) + [remote],
        capture_output=True,
        text=True,
        timeout=cfg.timeout + 10,
        check=False,
    )
    if r.returncode != 0 and not r.stdout:
        raise RuntimeError(f"origin curl failed: {(r.stderr or '').strip() or r.returncode}")
    out = r.stdout or ""
    if "\n__HTTP_CODE__" not in out:
        raise RuntimeError(f"origin curl bad output for {path}")
    text, _, code_s = out.rpartition("\n__HTTP_CODE__")
    return int(code_s), text


def check_public_external(cfg: Config) -> tuple[bool, str]:
    try:
        code, body = http_request(f"{cfg.base_external}/health", timeout=cfg.timeout)
        if code == 200 and body.strip() == "ok":
            return True, "✓ external /health → ok"
        return False, f"FAIL external /health HTTP {code}: {body[:80]}"
    except Exception as e:
        return False, f"WARN external /health unreachable ({e}) — use origin path or VPN"


def check_public_external_spa(cfg: Config) -> tuple[bool, str]:
    try:
        code, _ = http_request(f"{cfg.base_external}/app/", timeout=cfg.timeout)
        if code == 200:
            return True, "✓ external /app/ → HTTP 200"
        return False, f"FAIL external /app/ HTTP {code}"
    except Exception as e:
        return False, f"WARN external /app/ unreachable ({e})"


def check_public_origin(cfg: Config) -> tuple[bool, str]:
    try:
        code, body = origin_curl(cfg, "/health")
        if code == 200 and body.strip() == "ok":
            return True, "✓ origin /health → ok"
        return False, f"FAIL origin /health HTTP {code}: {body[:80]}"
    except Exception as e:
        return False, f"FAIL origin /health ({e})"


def login_origin(cfg: Config, password: str, client_ip: str | None) -> str:
    code, body = origin_curl(
        cfg,
        "/api/auth/login",
        method="POST",
        body={"username": cfg.admin_user, "password": password},
        client_ip=client_ip,
    )
    if code != 200:
        raise RuntimeError(f"login HTTP {code}: {body[:200]}")
    data = json.loads(body)
    token = data.get("access_token") or ""
    if not token:
        raise RuntimeError("login missing access_token")
    return token


def assert_json(code: int, body: str, path: str) -> dict[str, Any]:
    if code != 200:
        raise RuntimeError(f"{path} HTTP {code}: {body[:200]}")
    return json.loads(body)


def run_auth_checks(cfg: Config, token: str, client_ip: str | None) -> list[str]:
    lines: list[str] = []
    ip = client_ip
    code, body = origin_curl(cfg, "/health/detail", token=token, client_ip=ip)
    detail = assert_json(code, body, "/health/detail")
    checks = detail.get("checks") or {}
    for key in ("mysql", "redis"):
        if checks.get(key) != "ok":
            raise RuntimeError(f"/health/detail checks.{key}={checks.get(key)!r}")
    lines.append("✓ /health/detail mysql+redis ok")

    code, body = origin_curl(cfg, "/api/alert-history/stats", token=token, client_ip=ip)
    stats = assert_json(code, body, "/api/alert-history/stats")
    for key in ("today", "active", "recovered", "top_server"):
        if key not in stats:
            raise RuntimeError(f"alert stats missing {key}")
    lines.append("✓ /api/alert-history/stats contract OK")

    code, body = origin_curl(
        cfg, "/api/alert-history/?page=1&per_page=50", token=token, client_ip=ip
    )
    alerts = assert_json(code, body, "/api/alert-history/")
    if "items" not in alerts or "total" not in alerts:
        raise RuntimeError("alert-history missing items/total")
    if alerts.get("per_page") not in (50, None) and len(alerts["items"]) > 50:
        raise RuntimeError(f"alert per_page mismatch: {alerts.get('per_page')}")
    lines.append(
        f"✓ alert-history per_page=50 items={len(alerts['items'])} total={alerts['total']}"
    )

    for path in (
        "/api/assets/command-logs?page=1&per_page=5",
        "/api/assets/ssh-sessions?page=1&per_page=5",
    ):
        code, body = origin_curl(cfg, path, token=token, client_ip=ip)
        data = assert_json(code, body, path)
        if "items" not in data or "total" not in data:
            raise RuntimeError(f"{path} missing pagination")
        lines.append(f"  {path}: items={len(data['items'])} total={data['total']}")
    lines.append("✓ paginated command/session APIs OK")

    code, body = origin_curl(cfg, "/api/schedules/", token=token, client_ip=ip)
    schedules = assert_json(code, body, "/api/schedules/")
    if not isinstance(schedules, list):
        raise RuntimeError("/api/schedules/ not a list")
    lines.append(f"✓ /api/schedules/ count={len(schedules)}")

    code, body = origin_curl(cfg, "/api/servers/?per_page=1", token=token, client_ip=ip)
    servers = assert_json(code, body, "/api/servers/")
    items = servers.get("items") or []
    if not items:
        raise RuntimeError("no servers for schedule probe")
    sid = items[0]["id"]

    probe_name = "probe-once-health-check"
    create_body = {
        "name": probe_name,
        "schedule_type": "push",
        "run_mode": "once",
        "cron_expr": "0 10 * * *",
        "fire_at": "2026-12-01T10:00:00.000Z",
        "source_path": "/tmp/nexus_probe",
        "server_ids": f"[{sid}]",
        "enabled": False,
    }
    code, body = origin_curl(
        cfg,
        "/api/schedules/",
        method="POST",
        token=token,
        body=create_body,
        client_ip=ip,
    )
    if code != 201:
        raise RuntimeError(f"once schedule create HTTP {code}: {body[:300]}")
    created = json.loads(body)
    rid = created.get("id")
    if not rid:
        raise RuntimeError("schedule create missing id")
    lines.append(f"✓ once push schedule create HTTP 201 id={rid}")

    d_code, _ = origin_curl(
        cfg, f"/api/schedules/{rid}", method="DELETE", token=token, client_ip=ip
    )
    if d_code not in (200, 204):
        raise RuntimeError(f"schedule delete HTTP {d_code}")
    lines.append(f"✓ probe schedule deleted id={rid}")

    # Operator timezone: invalid date filters must 422 (no silent ignore)
    for bad_path in (
        "/api/alert-history/?date_from=not-a-date",
        "/api/audit/?date_from=bad",
    ):
        code, body = origin_curl(cfg, bad_path, token=token, client_ip=ip)
        if code != 422:
            raise RuntimeError(f"{bad_path} expected HTTP 422, got {code}: {body[:120]}")
    lines.append("✓ alert/audit invalid date_from → HTTP 422")

    # Cron Beijing semantics: 0 2 * * * → next_run UTC naive 18:00 (Beijing 02:00 next day)
    cron_body = {
        "name": "probe-cron-beijing-0200",
        "schedule_type": "push",
        "run_mode": "cron",
        "cron_expr": "0 2 * * *",
        "fire_at": None,
        "source_path": "/tmp/nexus_probe",
        "server_ids": f"[{sid}]",
        "enabled": True,
    }
    code, body = origin_curl(
        cfg, "/api/schedules/", method="POST", token=token, body=cron_body, client_ip=ip
    )
    if code != 201:
        raise RuntimeError(f"cron schedule create HTTP {code}: {body[:200]}")
    cron_sched = json.loads(body)
    cron_id = cron_sched.get("id")
    next_run = cron_sched.get("next_run") or ""
    if not cron_id or "18:00:00" not in next_run:
        raise RuntimeError(f"cron next_run Beijing mismatch: {next_run!r}")
    lines.append(f"✓ cron 0 2 * * * next_run={next_run} (UTC naive → Beijing 02:00)")
    d_code, _ = origin_curl(
        cfg, f"/api/schedules/{cron_id}", method="DELETE", token=token, client_ip=ip
    )
    if d_code not in (200, 204):
        raise RuntimeError(f"cron schedule delete HTTP {d_code}")

    # SPA entry: index.html must reference a loadable main bundle
    code, html = origin_curl(cfg, "/app/", client_ip=ip)
    if code != 200:
        raise RuntimeError(f"/app/ HTTP {code}")
    import re

    m = re.search(r'src="(/app/assets/index-[^"]+\.js)"', html)
    if not m:
        raise RuntimeError("/app/ index.html missing main bundle script")
    bundle_path = m.group(1)
    b_code, _ = origin_curl(cfg, bundle_path, client_ip=ip)
    if b_code != 200:
        raise RuntimeError(f"{bundle_path} HTTP {b_code}")
    lines.append(f"✓ /app/ main bundle OK ({bundle_path})")

    return lines


def main() -> int:
    parser = argparse.ArgumentParser(description="Nexus production probe")
    parser.add_argument(
        "--mode",
        choices=("auto", "external", "origin"),
        default="auto",
        help="auto: external best-effort + origin required",
    )
    args = parser.parse_args()
    cfg = build_config()

    print("=== Public (external) ===")
    ext_public_ok = True
    if args.mode in ("auto", "external"):
        ok, msg = check_public_external(cfg)
        print(msg)
        if not ok:
            ext_public_ok = False
        ok2, msg2 = check_public_external_spa(cfg)
        print(msg2)
        if not ok2:
            ext_public_ok = False
        try:
            code, _ = http_request(f"{cfg.base_external}/app/install.html", timeout=cfg.timeout)
            print(f"  /app/install.html → HTTP {code} (404 expected when locked)")
        except Exception:
            print("  /app/install.html → skip (external unreachable)")
    else:
        print("  (skipped)")

    print("=== Public (origin via SSH) ===")
    if args.mode in ("auto", "origin"):
        ok, msg = check_public_origin(cfg)
        print(msg)
        if not ok:
            print("=== FAILED: origin unreachable ===")
            return 1
    else:
        print("  (skipped)")

    print("=== Authenticated (origin via SSH) ===")
    pw = fetch_password_from_origin(cfg)
    if not pw:
        print("SKIP authenticated checks (NEXUS_TEST_ADMIN_PASSWORD not on server)")
        print("  → run: NEXUS_TEST_ADMIN_PASSWORD='…' bash scripts/prod_probe_install.sh")
        if args.mode == "auto" and not ext_public_ok:
            print("=== PASSED (origin only; external path blocked) ===")
            return 0
        if args.mode == "origin":
            return 0
        return 0

    client_ip = fetch_probe_client_ip(cfg)
    if not client_ip:
        print("SKIP authenticated checks (no NEXUS_PROBE_CLIENT_IP / login allowlist IP)")
        print("  → set NEXUS_PROBE_CLIENT_IP in .env.prod to a whitelisted login IP")
        return 0

    try:
        token = login_origin(cfg, pw, client_ip)
        print(f"✓ login OK (probe client IP {client_ip})")
        for line in run_auth_checks(cfg, token, client_ip):
            print(line)
    except Exception as e:
        print(f"FAIL authenticated: {e}")
        return 1

    print("=== All production probe checks passed ===")
    return 0


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