xiaobai-review/tools/verify_datahub_console.py

"""数据中枢控制台端到端自测：主站与中枢真的对话一遍，不是 mock。

跑法：python tools/verify_datahub_console.py
覆盖邀请码一次性注册、管理员门禁、凭证掩码、模型池与会员桥接读写、服务令牌校验。
两个服务都起在临时端口 + 临时数据目录，跑完自动清理，不碰任何现网数据。
"""
from __future__ import annotations

import http.cookies
import json
import os
import shutil
import socket
import sqlite3
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
# 从 tools/ 下运行，仓库根不在 sys.path 上；主站包按仓库根导入。
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
HUB_TOKEN = "smoke-hub-admin-token-0123456789abcdef"
PASSWORD = "SmokePass123"
FAILURES: list[str] = []


def free_port() -> int:
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def check(label: str, ok: bool, detail: str = "") -> None:
    print(f"  {'PASS' if ok else 'FAIL'}  {label}{(' — ' + detail) if detail else ''}")
    if not ok:
        FAILURES.append(label)


def request(url: str, payload=None, method="GET", headers=None, cookie="") -> tuple[int, dict, str]:
    data = json.dumps(payload).encode("utf-8") if payload is not None else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Content-Type", "application/json; charset=utf-8")
    for key, value in (headers or {}).items():
        req.add_header(key, value)
    if cookie:
        req.add_header("Cookie", cookie)
    try:
        with urllib.request.urlopen(req, timeout=15) as response:
            raw = response.read().decode("utf-8")
            set_cookie = response.headers.get("Set-Cookie") or ""
            return response.status, _parse(raw), set_cookie
    except urllib.error.HTTPError as exc:
        return exc.code, _parse(exc.read().decode("utf-8")), exc.headers.get("Set-Cookie") or ""


def _parse(raw: str) -> dict:
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        return {"_raw": raw[:200]}
    return parsed if isinstance(parsed, dict) else {"_list": parsed}


def session_cookie(header: str) -> str:
    jar = http.cookies.SimpleCookie()
    jar.load(header)
    morsel = jar.get("xiaobai_session")
    return f"xiaobai_session={morsel.value}" if morsel else ""


def wait_for(url: str, seconds: float = 20.0) -> bool:
    deadline = time.time() + seconds
    while time.time() < deadline:
        try:
            urllib.request.urlopen(url, timeout=2)
            return True
        except urllib.error.HTTPError:
            return True
        except OSError:
            time.sleep(0.25)
    return False


REVIEW_DB = ROOT / "data" / "review.db"
ENV_FILE = ROOT / ".env"


class RepoSandbox:
    """主站的库路径和 .env 都写死在仓库里，跑之前挪开、跑完原样放回。"""

    def __enter__(self) -> "RepoSandbox":
        self.stash = Path(tempfile.mkdtemp(prefix="hel560-stash-"))
        for path in (REVIEW_DB, ENV_FILE):
            if path.exists():
                shutil.copy2(path, self.stash / path.name)
        if REVIEW_DB.exists():
            REVIEW_DB.unlink()  # 自测需要一个空库来验证"首个账号免邀请码"
        return self

    def __exit__(self, *exc_info: object) -> None:
        for path in (REVIEW_DB, ENV_FILE):
            saved = self.stash / path.name
            if saved.exists():
                shutil.copy2(saved, path)
            elif path.exists():
                path.unlink()
        for extra in REVIEW_DB.parent.glob("review.db-*"):
            extra.unlink()
        shutil.rmtree(self.stash, ignore_errors=True)


def start_review(workdir: Path, port: int) -> None:
    os.environ["HUB_ADMIN_TOKEN"] = HUB_TOKEN
    from backend.application import RequestHandler, SERVICE  # noqa: F401
    from http.server import ThreadingHTTPServer

    server = ThreadingHTTPServer(("127.0.0.1", port), RequestHandler)
    threading.Thread(target=server.serve_forever, daemon=True).start()


def start_hub(workdir: Path, port: int, review_port: int) -> None:
    sys.path.insert(0, str(ROOT / "xiaobai-datahub"))
    from cryptography.fernet import Fernet

    os.environ.update(
        {
            "DATAHUB_DB_PATH": str(workdir / "hub.db"),
            "DATAHUB_BACKUP_DIR": str(workdir / "backups"),
            "DATAHUB_ENCRYPTION_KEY": Fernet.generate_key().decode(),
            "DATAHUB_TOKEN": "smoke-datahub-token",
            "HUB_ADMIN_TOKEN": HUB_TOKEN,
            "REVIEW_BASE_URL": f"http://127.0.0.1:{review_port}",
            "REVIEW_PUBLIC_URL": f"http://127.0.0.1:{review_port}",
            "DATAHUB_SCHEDULER_ENABLED": "0",
        }
    )
    from datahub.httpapp import make_handler
    from datahub.hub import Hub
    from datahub.settings import load_settings
    from http.server import ThreadingHTTPServer

    hub = Hub(load_settings(os.environ))
    server = ThreadingHTTPServer(("127.0.0.1", port), make_handler(hub))
    threading.Thread(target=server.serve_forever, daemon=True).start()


def main() -> int:
    if any(argument in {"-h", "--help"} for argument in sys.argv[1:]):
        print(__doc__.strip())
        return 0
    workdir = Path(tempfile.mkdtemp(prefix="hel560-smoke-"))
    review_port, hub_port = free_port(), free_port()
    review = f"http://127.0.0.1:{review_port}"
    hub = f"http://127.0.0.1:{hub_port}"
    with RepoSandbox():
        start_review(workdir, review_port)
        if not wait_for(f"{review}/api/session"):
            print("主站没起来")
            return 1
        start_hub(workdir, hub_port, review_port)
        if not wait_for(f"{hub}/livez"):
            print("数据中枢没起来")
            return 1

        print("\n[1] 首个账号免邀请码，之后注册强制邀请码")
        status, body, cookie_header = request(f"{review}/api/auth/register", {"username": "boss", "password": PASSWORD}, "POST")
        check("首个账号可直接注册（自动成为管理员）", status == 201, f"{status} {body.get('error', '')}")
        admin_cookie = session_cookie(cookie_header)
        status, body, _ = request(f"{review}/api/auth/register", {"username": "nobody", "password": PASSWORD}, "POST")
        check("第二个账号没邀请码被拒", status >= 400 and "邀请码" in str(body.get("error", "")), f"{status} {body.get('error', '')}")

        print("\n[2] 未登录 / 非管理员进不了数据中枢")
        status, body, _ = request(f"{hub}/admin/api/session")
        check("未登录访问控制台返回 401 并给出主站登录地址", status == 401 and "/login/" in str(body.get("login_url", "")), f"{status} {body}")

        print("\n[3] 主站管理员会话直接进控制台（跨端口共享 cookie）")
        status, session, _ = request(f"{hub}/admin/api/session", cookie=admin_cookie)
        check("带主站会话访问控制台返回 200", status == 200, f"{status} {session}")
        check("控制台回显主站用户名", session.get("username") == "boss", str(session.get("username")))
        csrf = str(session.get("csrf") or "")
        check("下发了 CSRF 令牌", len(csrf) >= 32, csrf[:12])
        write_headers = {"X-CSRF-Token": csrf}

        print("\n[4] 写接口必须带 CSRF")
        status, body, _ = request(f"{hub}/admin/api/invites/create", {"count": 1}, "POST", cookie=admin_cookie)
        check("缺 CSRF 的写请求被拒", status == 401, f"{status} {body}")

        print("\n[5] 控制台生成邀请码 → 注册消耗一次 → 二次使用失败")
        status, created, _ = request(f"{hub}/admin/api/invites/create", {"count": 2}, "POST", write_headers, admin_cookie)
        check("控制台生成邀请码成功", status == 200 and len(created.get("created") or []) == 2, f"{status} {created.get('error', '')}")
        codes = [item["code"] for item in created.get("created") or []]
        check("列表只给掩码，不回明文", all("•" in row["code_masked"] for row in created.get("codes") or []))
        status, body, member_cookie_header = request(
            f"{review}/api/auth/register", {"username": "xiaochen", "password": PASSWORD, "invite_code": codes[0]}, "POST"
        )
        check("凭邀请码注册成功", status == 201, f"{status} {body.get('error', '')}")
        member_cookie = session_cookie(member_cookie_header)
        status, body, _ = request(
            f"{review}/api/auth/register", {"username": "again", "password": PASSWORD, "invite_code": codes[0]}, "POST"
        )
        check("同一邀请码第二次注册被拒", status >= 400, f"{status} {body.get('error', '')}")

        print("\n[6] 作废后的邀请码不能注册")
        handles = {row["code_masked"][:7]: row["code_id"] for row in created.get("codes") or []}
        target = handles.get(codes[1][:7])
        status, body, _ = request(f"{hub}/admin/api/invites/revoke", {"code_id": target}, "POST", write_headers, admin_cookie)
        check("控制台作废未使用的邀请码", status == 200, f"{status} {body.get('error', '')}")
        status, body, _ = request(
            f"{review}/api/auth/register", {"username": "revoked", "password": PASSWORD, "invite_code": codes[1]}, "POST"
        )
        check("已作废邀请码无法注册", status >= 400, f"{status} {body.get('error', '')}")

        print("\n[7] 普通会员账号进不了控制台")
        status, body, _ = request(f"{hub}/admin/api/session", cookie=member_cookie)
        check("非管理员访问控制台返回 403", status == 403, f"{status} {body}")
        status, body, _ = request(f"{hub}/admin/api/members", cookie=member_cookie)
        check("非管理员读会员接口同样 403", status == 403, f"{status} {body}")

        print("\n[8] 数据源凭证在线写入 + 掩码回显")
        status, body, _ = request(
            f"{hub}/admin/api/credentials/tushare", {"tushare_token": "tok-abcdefgh1234"}, "POST", write_headers, admin_cookie
        )
        check("Tushare Token 保存成功", status == 200, f"{status} {body.get('error', '')}")
        status, sources, _ = request(f"{hub}/admin/api/sources", cookie=admin_cookie)
        tushare = next((row for row in sources.get("items") or [] if row.get("provider") == "tushare"), {})
        credential = tushare.get("credential") or {}
        check("数据源卡回显掩码而非明文", credential.get("configured") and "1234" in str(credential.get("last4")), str(credential))
        check("接口不回传明文 Token", "tok-abcdefgh1234" not in json.dumps(sources, ensure_ascii=False))

        print("\n[9] 模型池 / 会员 / 邀请码三页都能从控制台读到")
        for label, path in (("模型池", "/admin/api/models"), ("会员", "/admin/api/members"), ("邀请码", "/admin/api/invites")):
            status, body, _ = request(f"{hub}{path}", cookie=admin_cookie)
            check(f"{label}接口可读", status == 200, f"{status} {body.get('error', '')}")

        print("\n[10] 控制台改模型池 → 主站落库")
        models = [{"id": "smoke-main", "name": "冒烟主模型", "model": "gpt-4o", "base_url": "https://api.openai.com/v1", "api_key": "sk-smoke-key-9911"}]
        status, body, _ = request(
            f"{hub}/admin/api/models/save", {"models": models, "primary_model_id": "smoke-main"}, "POST", write_headers, admin_cookie
        )
        check("控制台保存模型池成功", status == 200, f"{status} {body.get('error', '')}")
        groups = body.get("groups") or []
        check("按供应商归组返回", len(groups) == 1 and groups[0]["base_url"] == "https://api.openai.com/v1", str(groups)[:120])
        check("供应商显示密钥后四位而非明文", groups and groups[0].get("key_last4") == "9911", str(groups[0].get("key_last4") if groups else ""))
        check("模型接口不回传明文密钥", "sk-smoke-key-9911" not in json.dumps(body, ensure_ascii=False))
        status, status_body, _ = request(
            f"{review}/api/admin/settings", None, "GET", {"X-Hub-Admin-Token": HUB_TOKEN}, admin_cookie
        )
        status, mainsite, _ = request(f"{review}/api/hub-admin/status", {}, "POST", {"X-Hub-Admin-Token": HUB_TOKEN})
        pool = (mainsite.get("llm") or {}).get("models") or []
        check("主站确实存下了这个模型", any(m["id"] == "smoke-main" for m in pool), str([m.get("id") for m in pool]))

        print("\n[11] 会员额度与会员开通经控制台落到主站")
        status, body, _ = request(f"{hub}/admin/api/members/quota", {"member_daily_limit": 88}, "POST", write_headers, admin_cookie)
        check("保存会员每日额度成功", status == 200 and (body.get("membership") or {}).get("member_daily_limit") == 88, f"{status} {body.get('membership')}")
        member_id = next((u["id"] for u in body.get("users") or [] if u["username"] == "xiaochen"), 0)
        status, body, _ = request(
            f"{hub}/admin/api/members/save", {"user_id": member_id, "status": "active", "duration": "3_months"}, "POST", write_headers, admin_cookie
        )
        row = next((u for u in body.get("users") or [] if u["id"] == member_id), {})
        check("开通 3 个月会员生效", status == 200 and row.get("membership_status") == "active" and row.get("membership_expires_at"), f"{status} {row.get('membership_status')} {row.get('membership_expires_at')}")

        print("\n[11b] 主站行情任务（后台刷新开关 / 手动刷新 / 历史回补）经桥接可用")
        status, body, _ = request(f"{hub}/admin/api/market/status", cookie=admin_cookie)
        check("读取主站行情状态", status == 200 and body.get("background_refresh_enabled") is True, f"{status} {body.get('error', '')}")
        status, body, _ = request(
            f"{hub}/admin/api/market/settings", {"background_refresh_enabled": False}, "POST", write_headers, admin_cookie
        )
        check("暂停后台刷新生效", status == 200 and body.get("background_refresh_enabled") is False, f"{status} {body.get('background_refresh_enabled')}")
        status, mainsite_setting, _ = request(f"{review}/api/admin/settings", cookie=admin_cookie)
        check("主站侧确认开关已落库", (mainsite_setting.get("data") or {}).get("background_refresh_enabled") is False, str((mainsite_setting.get("data") or {}).get("background_refresh_enabled")))
        status, body, _ = request(
            f"{hub}/admin/api/market/settings", {"background_refresh_enabled": True}, "POST", write_headers, admin_cookie
        )
        check("重新启用后台刷新", status == 200 and body.get("background_refresh_enabled") is True, f"{status}")

        status, body, _ = request(f"{hub}/admin/api/market/refresh", {}, "POST", write_headers, admin_cookie)
        refresh_key = str(body.get("job_key") or "")
        check("提交手动后台刷新任务", status == 200 and body.get("started") and refresh_key.startswith("manual:"), f"{status} {body}")
        status, body, _ = request(f"{hub}/admin/api/market/backfill", {"start_date": "", "end_date": ""}, "POST", write_headers, admin_cookie)
        check("回补缺日期被拒绝", status == 400, f"{status} {body.get('error', '')}")
        status, body, _ = request(
            f"{hub}/admin/api/market/backfill", {"start_date": "2026-09-01", "end_date": "2026-09-08"}, "POST", write_headers, admin_cookie
        )
        backfill_key = str(body.get("job_key") or "")
        check("提交历史回补任务", status == 200 and body.get("started") and backfill_key.startswith("backfill:"), f"{status} {body}")
        terminal = []
        for _ in range(60):
            status, body, _ = request(f"{hub}/admin/api/market/status", cookie=admin_cookie)
            jobs = {str(item.get("idempotency_key")): str(item.get("status")) for item in body.get("jobs") or []}
            if jobs.get(refresh_key) in {"success", "failed"} and jobs.get(backfill_key) in {"success", "failed"}:
                terminal = [jobs[refresh_key], jobs[backfill_key]]
                break
            time.sleep(0.5)
        check("两个任务都进入终态（沙箱无行情源，失败也是正确终态）", len(terminal) == 2 and all(t in {"success", "failed"} for t in terminal), str(terminal))

        print("\n[12] 桥接令牌是唯一信任边界")
        status, body, _ = request(f"{review}/api/hub-admin/status", {}, "POST", {"X-Hub-Admin-Token": "wrong-token"})
        check("桥接端点拒绝错误令牌", status == 401, f"{status} {body}")
        status, body, _ = request(f"{review}/api/hub-admin/status", {}, "POST")
        check("桥接端点拒绝无令牌", status == 401, f"{status} {body}")
        status, body, _ = request(f"{review}/api/hub-admin/invites", {}, "POST", {"X-Hub-Admin-Token": HUB_TOKEN}, admin_cookie)
        check("带正确令牌可读邀请码", status == 200, f"{status} {body.get('error', '')}")

        print("\n[13] 退出登录会真的销毁主站会话")
        status, body, _ = request(f"{hub}/admin/api/logout", {}, "POST", write_headers, admin_cookie)
        check("控制台退出返回主站登录地址", status == 200 and "/login/" in str(body.get("login_url", "")), f"{status} {body}")
        status, body, _ = request(f"{review}/api/session", cookie=admin_cookie)
        check("主站会话已失效", not (body.get("authenticated") or body.get("user")), str(body)[:120])

        print("\n[14] 并发使用同一邀请码只成功一次")
        os.environ["HUB_ADMIN_TOKEN"] = HUB_TOKEN
        with sqlite3.connect(REVIEW_DB) as connection:
            rows = connection.execute("SELECT status, COUNT(*) FROM invite_codes GROUP BY status").fetchall()
        counts = dict(rows)
        check("邀请码状态落库正确（1 已用 / 1 已作废）", counts.get("used") == 1 and counts.get("revoked") == 1, str(counts))
        shutil.rmtree(workdir, ignore_errors=True)

    print("\n" + "=" * 60)
    if FAILURES:
        print(f"FAILED {len(FAILURES)} 项：")
        for item in FAILURES:
            print("  - " + item)
        return 1
    print("数据中枢控制台端到端自测全部通过")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())