workbuddy-portal/tools/check_live.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli

"""端到端验收：对**运行中的**服务发真实 HTTP 请求，走完整登录/CSRF/API 链路。

与 tests 里用 Flask test_client 的冒烟测试互补——这里验证的是「真的起起来了、
真的能登录、真的能取到数」，适合部署到局域网后随手跑一遍。

用法：
    python tools/check_live.py                        # 默认 http://127.0.0.1:8848
    python tools/check_live.py --base http://192.168.1.50:8848   # 换成你的部署主机
    python tools/check_live.py -u admin -p 你的密码
    python tools/check_live.py --from 2026-09-08 --to 2026-09-14

退出码：0 全通过；1 有失败项（会打印失败清单）。

注意：脚本会读取窗口数据但**不写库**（不触发采集、不改配置），可安全反复运行。
"""
from __future__ import annotations

import argparse
import base64
import http.cookiejar
import json
import os
import re
import sqlite3
import sys
import urllib.error
import urllib.parse
import urllib.request
import zlib
from datetime import datetime

OK = 0
FAIL = 0
FAILS: list[str] = []


def _d(s: str):
    """把 YYYY-MM-DD 解析成本地 datetime（不用 date.fromisoformat 之外的时区处理）。"""
    return datetime.strptime(s, "%Y-%m-%d")


def decode_session(cj) -> dict:
    """从 Flask 会话 cookie 里解出那份**未加密**的载荷。

    Flask 的会话是「签名 + base64，**不加密**」的 —— 也就是说持有 cookie 的人
    就能读到里面的内容。本项目因此把验证码答案放在服务端 captchas 表里，
    会话里只留一个随机 id；本函数存在的意义就是取出那个 id，
    好让自动化验收能跨过验证码这一关（顺便也验证了「答案不在会话里」）。
    """
    for c in cj:
        if not c.name.startswith("workbuddy_portal_sid"):
            continue
        seg = urllib.parse.unquote(c.value).split(".")[0]
        seg += "=" * (-len(seg) % 4)
        try:
            raw = base64.urlsafe_b64decode(seg)
            try:
                raw = zlib.decompress(raw)      # 某些版本的 itsdangerous 会压
            except zlib.error:
                pass
            return json.loads(raw.decode("utf-8"))
        except Exception:                       # noqa: BLE001
            return {}
    return {}


def chk(name: str, cond: bool, extra: str = "") -> None:
    global OK, FAIL
    if cond:
        OK += 1
        print("  [OK]   %s %s" % (name, extra))
    else:
        FAIL += 1
        FAILS.append(name)
        print("  [FAIL] %s %s" % (name, extra))


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    """不要自动跟随 302 —— 检查跳转目标本身是否安全时必须看到原始 Location。"""

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class Live:
    def __init__(self, base: str, timeout: int = 20, db_path: str | None = None):
        self.base = base.rstrip("/")
        self.timeout = timeout
        self.db_path = db_path
        # 关键：显式清空代理，否则本机代理会把 127.0.0.1 也拦成 502
        self.cj = http.cookiejar.CookieJar()
        self.op = urllib.request.build_opener(
            urllib.request.ProxyHandler({}),
            urllib.request.HTTPCookieProcessor(self.cj),
        )
        self.op.addheaders = [("User-Agent", "workbuddy-portal-check/1.1")]
        # 不跟随跳转的 opener：共用同一个 cookie jar，保证是同一会话
        self.op_nr = urllib.request.build_opener(
            urllib.request.ProxyHandler({}),
            urllib.request.HTTPCookieProcessor(self.cj),
            _NoRedirect,
        )
        self.op_nr.addheaders = [("User-Agent", "workbuddy-portal-check/1.1")]

    def get(self, path: str):
        try:
            r = self.op.open(urllib.request.Request(self.base + path), timeout=self.timeout)
            return r.status, r.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as e:
            return e.code, e.read().decode("utf-8", "replace")

    def post(self, path: str, data: dict, csrf: str | None = None, as_json: bool = False):
        if as_json:
            body, ct = json.dumps(data).encode(), "application/json"
        else:
            body, ct = urllib.parse.urlencode(data).encode(), "application/x-www-form-urlencoded"
        req = urllib.request.Request(self.base + path, data=body, method="POST")
        req.add_header("Content-Type", ct)
        if csrf:
            req.add_header("X-CSRF-Token", csrf)
        try:
            r = self.op.open(req, timeout=self.timeout)
            return r.status, r.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as e:
            return e.code, e.read().decode("utf-8", "replace")

    def post_raw(self, path: str, data: dict):
        """表单 POST 且**不跟随**跳转，返回 (status, Location)。"""
        body = urllib.parse.urlencode(data).encode()
        req = urllib.request.Request(self.base + path, data=body, method="POST")
        req.add_header("Content-Type", "application/x-www-form-urlencoded")
        try:
            r = self.op_nr.open(req, timeout=self.timeout)
            return r.status, r.headers.get("Location")
        except urllib.error.HTTPError as e:
            return e.code, e.headers.get("Location")

    def jget(self, path: str) -> dict:
        st, body = self.get(path)
        return json.loads(body) if st == 200 else {}

    def raw(self, path: str):
        """返回 (status, headers, bytes)——验证码/响应头这类要原始字节的场景用。"""
        try:
            r = self.op.open(urllib.request.Request(self.base + path), timeout=self.timeout)
            return r.status, r.headers, r.read()
        except urllib.error.HTTPError as e:
            return e.code, e.headers, e.read()

    def form_csrf(self, path: str) -> str:
        """取某个页面里的 CSRF 隐藏域（该页面必须与当前会话同源）。"""
        _, html = self.get(path)
        m = re.search(r'name="_csrf"\s+value="([^"]+)"', html)
        return m.group(1) if m else ""

    # ---- 验证码辅助（仅验收脚本用）----
    def solve_captcha(self, purpose: str):
        """取一张图 -> 从会话里读 id -> 从本地库里取答案。返回 (答案, 会话载荷)。"""
        self.raw("/captcha.png?purpose=" + purpose)
        sess = decode_session(self.cj)
        cid = sess.get("cap_" + purpose)
        if not cid or not self.db_path or not os.path.exists(self.db_path):
            return None, sess
        try:
            con = sqlite3.connect(self.db_path)
            try:
                row = con.execute("SELECT answer FROM captchas WHERE id=?", (cid,)).fetchone()
            finally:
                con.close()
        except sqlite3.Error:
            return None, sess
        return (row[0] if row else None), sess

    def login(self, user: str, pwd: str, nxt: str = "", follow: bool = True):
        """完整登录（验证码策略为 always 时自动解）。

        follow=False 时返回原始 (status, Location)，用于验证跳转目标是否安全。
        返回 (status, location, need_captcha, session_payload)。
        """
        html = self.get("/login")[1]
        need_cap = 'name="captcha"' in html
        m = re.search(r'name="_csrf"\s+value="([^"]+)"', html)
        data = {"username": user, "password": pwd, "_csrf": m.group(1) if m else ""}
        if nxt:
            data["next"] = nxt
        sess = {}
        if need_cap:
            ans, sess = self.solve_captcha("login")
            if ans is None:
                return None, None, True, sess
            data["captcha"] = ans
        if follow:
            st, _ = self.post("/login", data)
            return st, None, need_cap, sess
        st, loc = self.post_raw("/login", data)
        return st, loc, need_cap, sess


def run(L: Live, user: str, pwd: str, frm: str, to: str) -> None:
    print("== 1. 未登录访问受保护资源 ==")
    st, body = L.get("/")
    chk("GET / 未登录落登录页", st == 200 and "登录" in body, "status=%s" % st)
    for p in ("/api/summary", "/api/bundle", "/api/manifest"):
        st, _ = L.get(p)
        chk("GET %-14s 未登录=401" % p, st == 401, "status=%s" % st)

    print("== 2. 登录（含 CSRF；验证码策略为 always 时自动解） ==")
    st, html = L.get("/login")
    chk("登录页含 CSRF 隐藏域", bool(re.search(r'name="_csrf"\s+value="([^"]+)"', html)))
    st, _, need_cap, sess = L.login(user, pwd)
    chk("登录成功", st in (200, 302), "status=%s" % st)
    if need_cap:
        # 会话里只应有 id，不该有答案本身
        chk("会话里只存验证码 id（不是答案）", bool(sess.get("cap_login")),
            "cap_login=%s" % (sess.get("cap_login") or "无"))
    st, html = L.get("/")
    chk("登录后 GET / 到概览", st == 200 and "概览" in html, "len=%d" % len(html))

    print("== 3. 后台页面均可达 ==")
    for p, kw in [("/", "概览"), ("/tasks", "任务"), ("/config", "配置"),
                  ("/logs", "日志"), ("/records", "记录")]:
        st, html = L.get(p)
        chk("GET %-10s" % p, st == 200 and kw in html, "status=%s len=%d" % (st, len(html)))

    print("== 4. 登录后写操作仍需 CSRF ==")
    st, _ = L.post("/api/collect", {}, csrf=None, as_json=True)
    chk("POST /api/collect 缺 CSRF=400", st == 400, "status=%s" % st)

    print("== 5. 结构与数值 ==")
    mf = L.jget("/api/manifest")
    chk("manifest 含 health/archive/totals/sources",
        all(k in mf for k in ("health", "archive", "totals", "sources")))
    src = (mf.get("sources") or [{}])[0]
    chk("manifest 存档条数与数据源一致",
        mf["totals"]["records"] == src.get("count"),
        "records=%s src=%s" % (mf["totals"]["records"], src.get("count")))

    sm = L.jget("/api/summary?from=%s&to=%s" % (frm, to))
    want_days = (_d(to) - _d(frm)).days + 1
    chk("summary 窗口天数正确", sm.get("window", {}).get("days") == want_days,
        "window=%s 期望 %d 天" % (sm.get("window"), want_days))
    chk("summary 窗口内有记录", (sm.get("records") or 0) > 0, "records=%s" % sm.get("records"))
    chk("summary 环比 prev 存在", bool(sm.get("prev")),
        "prev=%s~%s" % ((sm.get("prev") or {}).get("firstDay"), (sm.get("prev") or {}).get("lastDay")))
    chk("summary avgPerCall 自洽",
        not sm.get("calls") or abs(sm["avgPerCall"] - round(sm["credits"] / sm["calls"], 4)) < 1e-6)

    bd = L.jget("/api/bundle?from=%s&to=%s" % (frm, to))
    chk("bundle 顶层键齐全",
        {"manifest", "daily", "dims", "top", "records", "totals", "window"} <= set(bd),
        "keys=%s" % list(bd.keys()))
    recs, daily = bd.get("records", []), bd.get("daily", [])
    rsum = round(sum(float(x["c"]) for x in recs), 2)
    tsum = round(float(bd.get("totals", {}).get("credits", 0)), 2)
    print("      窗口 %d 条，records 求和 %.2f ；totals.credits %.2f" % (len(recs), rsum, tsum))
    chk("records 求和 == totals.credits", abs(rsum - tsum) < 0.005, "diff=%.4f" % (rsum - tsum))
    chk("records 求和 == summary.credits",
        abs(rsum - float(sm.get("credits", 0))) < 0.005,
        "diff=%.4f" % (rsum - float(sm.get("credits", 0))))
    chk("daily 为全量（多于窗口天数，供日历/日期轴）", len(daily) > want_days,
        "daily=%d 天 > 窗口 %d 天" % (len(daily), want_days))
    chk("daily 全量求和 == 存档总额",
        abs(round(sum(float(x["c"]) for x in daily), 2)
            - round(float(mf["totals"]["credits"]), 2)) < 0.005)
    chk("逐日 h[24] 求和 == 当日 c",
        all(abs(round(sum(d["h"]), 2) - round(d["c"], 2)) < 0.005 for d in daily))
    chk("dims.hour 补齐 24 槽", len(bd.get("dims", {}).get("hour", [])) == 24)
    chk("dims.model 非空", len(bd.get("dims", {}).get("model", [])) > 0)
    top = bd.get("top", [])
    chk("top 榜按积分降序",
        all(top[i]["c"] >= top[i + 1]["c"] for i in range(len(top) - 1)), "n=%d" % len(top))

    print("== 6. 明细分页/筛选/排序 ==")
    rj = L.jget("/api/records?page=1&size=5")
    chk("分页返回 5 条", len(rj.get("items", [])) == 5,
        "total=%s pages=%s" % (rj.get("total"), rj.get("pages")))
    chk("分页 total 与存档一致", rj.get("total") == mf["totals"]["records"])
    chk("分页字段为可读全名", "request_id" in (rj.get("items") or [{}])[0])
    chk("分页页码自洽",
        rj.get("pages") == max(1, (rj.get("total", 0) + rj.get("size", 1) - 1) // rj.get("size", 1)))
    r2 = L.jget("/api/records?page=2&size=5")
    chk("第 2 页与第 1 页不重叠",
        set(x["request_id"] for x in r2.get("items", [])).isdisjoint(
            set(x["request_id"] for x in rj.get("items", []))))
    models = bd.get("dims", {}).get("model", [])
    if models:
        mn = models[0]["name"]
        rf = L.jget("/api/records?page=1&size=5&model=" + urllib.parse.quote(mn))
        chk("按模型筛选生效", all(x["model"] == mn for x in rf.get("items", [])),
            "model=%s total=%s" % (mn, rf.get("total")))
    ro = L.jget("/api/records?page=1&size=10&order=credits_desc")
    chk("按积分降序生效",
        all(ro["items"][i]["credits"] >= ro["items"][i + 1]["credits"]
            for i in range(len(ro.get("items", [])) - 1)))

    print("== 7. 凭据不外泄 ==")
    stj = L.jget("/api/settings")
    # 契约：settings 里 cookie 这个键**必须为空**（db.get_settings 统一置空），
    # 真正的状态只通过 cookie_hint / cookie_broken 这两个派生字段暴露。
    chk("settings 里 cookie 字段为空串",
        "cookie" in stj and not str(stj.get("cookie") or "").strip(),
        "cookie=%r" % stj.get("cookie"))
    chk("settings 用 cookie_hint/cookie_broken 代替明文",
        "cookie_hint" in stj and "cookie_broken" in stj)
    chk("settings 仅回 cookie_hint 掩码",
        bool(stj.get("cookie_hint")) and len(str(stj.get("cookie_hint"))) < 200,
        "hint=%s" % stj.get("cookie_hint"))
    chk("settings 回传实例级键清单", isinstance(stj.get("_globalKeys"), list)
        and bool(stj.get("_globalKeys")), "%s" % stj.get("_globalKeys"))
    chk("settings 标明能否改实例级配置", stj.get("_canEditGlobal") is True)
    chk("配置页 HTML 不含 cookie 明文", "eyJ" not in L.get("/config")[1])
    # 密文形态：v1.<b64salt>.<b64nonce>.<b64ct>.<b64tag>，恰好用正则判定，
    # 免得把版本号 "v1.2.0" 当成泄漏（这两者前缀撞车）
    cipher_re = re.compile(r"v1\.[A-Za-z0-9+/=]{8,}\.[A-Za-z0-9+/=]{8,}\.")
    for p in ("/config", "/profile", "/"):
        chk("%-9s HTML 里没有 Cookie 密文" % p, not cipher_re.search(L.get(p)[1]))

    print("== 8. 错误处理 ==")
    for p in ("/api/nope", "/nope"):
        st, _ = L.get(p)
        chk("GET %-12s =404" % p, st == 404, "status=%s" % st)
    for p in ("/api/summary?from=abc&to=def", "/api/daily?from=2026-13-99"):
        st, _ = L.get(p)
        chk("GET %-32s 非法日期=400" % p, st == 400, "status=%s" % st)

    print("== 9. 新增能力：用户管理 / 审计 / 流式导出 ==")
    st, html = L.get("/users")
    chk("GET /users 管理员可达", st == 200 and "用户管理" in html, "status=%s" % st)
    chk("用户管理页不回传口令散列", "pbkdf2:" not in html)
    au = L.jget("/api/audit?size=5")
    chk("GET /api/audit 结构完整",
        all(k in au for k in ("total", "page", "size", "pages", "actions", "items")),
        "keys=%s" % list(au.keys()))
    chk("审计条目带 actor/action/at",
        not au.get("items") or {"actor", "action", "at"} <= set(au["items"][0]),
        "n=%d" % len(au.get("items", [])))

    st, csv_body = L.get("/records/export?from=%s&to=%s" % (frm, to))
    chk("GET /records/export=200", st == 200, "status=%s" % st)
    chk("导出带 UTF-8 BOM（Excel 不乱码）", csv_body.startswith("\ufeff"))
    lines = [x for x in csv_body.lstrip("\ufeff").split("\r\n") if x]
    chk("导出表头为官网同构列",
        lines and lines[0] == "RequestID,积分消耗,User Prompt,模型,客户端,时间",
        "header=%s" % (lines[0] if lines else None))
    chk("导出行数 == 窗口记录数 + 表头", len(lines) == (sm.get("records") or 0) + 1,
        "csv=%d 记录=%s" % (len(lines), sm.get("records")))
    r1 = L.jget("/api/records?page=1&size=1&from=%s&to=%s" % (frm, to))
    first_id = ((r1.get("items") or [{}])[0]).get("request_id")
    chk("导出与明细同源同序（首行 == 明细首条）",
        len(lines) > 1 and bool(first_id) and first_id in lines[1],
        "api=%s csv=%s" % (first_id, (lines[1][:40] if len(lines) > 1 else None)))

    print("== 10. 安全：开放重定向与凭证外泄 ==")
    L2 = Live(L.base, L.timeout, L.db_path)             # 全新会话，避免已登录被直跳
    st, loc, _, _ = L2.login(user, pwd, nxt="//evil.com", follow=False)
    chk("next=//evil.com 被拒（不出现协议相对跳转）",
        st == 302 and "evil.com" not in (loc or "") and not (loc or "").startswith("//"),
        "status=%s Location=%s" % (st, loc))
    L3 = Live(L.base, L.timeout, L.db_path)
    st, loc, _, _ = L3.login(user, pwd, nxt="/records", follow=False)
    chk("next=/records 站内路径正常放行", st == 302 and loc == "/records",
        "status=%s Location=%s" % (st, loc))
    st, loc = L3.post_raw("/login", {"username": user, "password": pwd, "_csrf": "wrong"})
    chk("错误 CSRF 的登录 POST=400", st == 400, "status=%s" % st)
    st, body = L.get("/logout")
    chk("GET /logout 不执行退出（仅提示）", st == 200 and "退出" in body, "status=%s" % st)
    st, html = L.get("/")
    chk("GET /logout 后仍处于登录态", st == 200 and "概览" in html, "status=%s" % st)

    print("== 11. 多用户：注册入口 / 验证码 / 安全响应头 ==")
    L4 = Live(L.base, L.timeout, L.db_path)             # 全新未登录会话
    st, html = L4.get("/register")
    chk("GET /register 可达", st == 200 and "注册" in html, "status=%s" % st)
    chk("注册页带验证码图", "capimg" in html and "/captcha.png" in html)
    chk("注册页带 CSRF 隐藏域", bool(L4.form_csrf("/register")))
    st, html = L4.get("/login")
    chk("登录页带验证码图", "capimg" in html and 'name="captcha"' in html)
    chk("登录页带自助注册链接", "/register" in html)

    # 出图：真实字节 + 禁缓存 + 确实每次都不一样
    shots = {}
    for purpose in ("login", "register"):
        code, hdr, data = L4.raw("/captcha.png?purpose=" + purpose)
        chk("GET /captcha.png?purpose=%-8s 出 PNG" % purpose,
            code == 200 and data[:4] == b"\x89PNG" and len(data) > 200,
            "status=%s len=%d" % (code, len(data)))
        chk("   └ 禁缓存 no-store", "no-store" in (hdr.get("Cache-Control") or ""))
        chk("   └ 类型 image/png", (hdr.get("Content-Type") or "").startswith("image/png"))
        shots[purpose] = data
    _, _, again = L4.raw("/captcha.png?purpose=login")
    chk("两次取图内容不同（不是一张静态图）", again != shots["login"])
    chk("login 与 register 的图互不相同", shots["login"] != shots["register"])
    # 图必须由服务端单独下发，不能把答案内联进页面
    st, html = L4.get("/login")
    chk("登录页没有内联 data: 图片（答案不走页面源码）",
        "data:image" not in html and "base64," not in html)

    # 安全响应头
    code, hdr, _ = L4.raw("/login")
    for name, want in (("X-Content-Type-Options", "nosniff"),
                       ("X-Frame-Options", "DENY"),
                       ("Referrer-Policy", "same-origin")):
        chk("响应头 %-24s" % name, (hdr.get(name) or "") == want, "=%s" % hdr.get(name))
    chk("响应头含 CSP 且 frame-ancestors 'none'",
        "frame-ancestors 'none'" in (hdr.get("Content-Security-Policy") or ""))
    code, hdr, _ = L4.raw("/captcha.png?purpose=login")
    chk("/captcha 路径带 no-store", "no-store" in (hdr.get("Cache-Control") or ""))

    # 路径穿越式 purpose 必须被收敛到已知用途，而不是 500
    code, _, data = L4.raw("/captcha.png?purpose=../../etc/passwd")
    chk("非法 purpose 不报 500", code == 200 and data[:4] == b"\x89PNG",
        "status=%s" % code)


def main() -> int:
    ap = argparse.ArgumentParser(description="对运行中的用量门户做端到端验收")
    ap.add_argument("--base", default="http://127.0.0.1:8848", help="服务地址")
    ap.add_argument("-u", "--user", default="admin", help="登录用户名")
    ap.add_argument("-p", "--password", default="admin123", help="登录密码")
    ap.add_argument("--from", dest="frm", default="2026-09-08", help="验收窗口起")
    ap.add_argument("--to", dest="to", default="2026-09-14", help="验收窗口止")
    ap.add_argument("--db", default=None,
                    help="SQLite 路径（默认 <repo>/data/usage.sqlite）。"
                         "验证码策略为 always 时用它取答案以完成自动登录；"
                         "指向不存在的文件则跳过需要验证码的登录")
    ap.add_argument("--timeout", type=int, default=20)
    a = ap.parse_args()

    db_path = a.db or os.path.join(
        os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "usage.sqlite")
    print("目标：%s   窗口：%s ~ %s\n验证码答案源：%s\n" % (a.base, a.frm, a.to, db_path))
    run(Live(a.base, a.timeout, db_path), a.user, a.password, a.frm, a.to)
    print("\nRESULT: ok=%d fail=%d" % (OK, FAIL))
    if FAILS:
        print("失败项：%s" % "、".join(FAILS))
    return 1 if FAIL else 0


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