workbuddy-portal/tools/shots.py

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

"""界面实检：登录后逐页截图，用来目视确认「统一美化」是否真的落地。

用法：
    WB_DATA_DIR=$PWD/data/demo python manage.py serve --port 8849 --no-scheduler
    python tools/shots.py --base http://127.0.0.1:8849 --db data/demo/usage.sqlite
    python tools/shots.py --full                        # 整页长图（默认只截首屏）

产物：data/shots/*.png（该目录在 .gitignore 内，可直接删）。

为什么不用 headless chrome 直出：本项目的页面都要登录态，
`--screenshot` 无法注入会话 Cookie，所以必须用 Playwright 走一次真实登录。

验证码：默认策略是 always，所以本脚本会**从本地库里取答案**（取的是会话里的
captcha id，答案只存在于服务端），这样自动化能跨过验证码这一关。
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sqlite3
import sys
import urllib.parse
import zlib

BASE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE)

# 文件名刻意与原有编号保持一致（docs/USER-GUIDE.md 里就是按这些名字引用的），
# 新增页面排在后面，避免为了两张新图去改一堆文档链接。
CAPTURES = [
    # (文件名, 路径, 标签, 是否需登录)
    ("00-login.png",     "/login",     "登录页   （含图形验证码）", False),
    ("09-register.png",  "/register",  "注册页",                   False),
    ("01-overview.png",  "/",          "概览",                     True),
    ("02-records.png",   "/records",   "数据明细",                 True),
    ("03-tasks.png",     "/tasks",     "任务管理",                 True),
    ("04-config.png",    "/config",    "配置管理",                 True),
    ("05-logs.png",      "/logs",      "日志管理",                 True),
    ("06-users.png",     "/users",     "用户管理",                 True),
    ("07-dashboard.png", "/dashboard", "用量大屏",                 True),
    ("10-profile.png",   "/profile",   "个人中心",                 True),
]


def _find_browser() -> str | None:
    """找一个可用的 Chromium 可执行文件。

    Playwright 的驱动版本与本地已下载的浏览器版本常常错位（例如驱动要
    chromium_headless_shell-1234 而本机只有 -1228），此时 launch() 会直接报
    「Executable doesn't exist」。这里按优先级探测，命中就显式传 executable_path。
    """
    import glob
    home = os.environ.get("LOCALAPPDATA") or os.path.expanduser("~")
    pats = [
        os.path.join(home, "ms-playwright", "chromium-*", "chrome-win64", "chrome.exe"),
        os.path.join(home, "ms-playwright", "chromium_headless_shell-*",
                     "chrome-headless-shell-win64", "chrome-headless-shell.exe"),
        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
        r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
    ]
    for p in pats:
        hits = sorted(glob.glob(p))
        if hits:
            return hits[-1]
    return None


def _session_payload(ctx) -> dict:
    """解出 Flask 会话 cookie 里的载荷，只为拿 captcha 的 id。

    Flask 会话是「签名 + base64，不加密」的，所以这里面能读出内容 ——
    正因如此，验证码答案绝不能放进去（本项目只放一个随机 id）。
    """
    for c in ctx.cookies():
        if not c.get("name", "").startswith("workbuddy_portal_sid"):
            continue
        seg = urllib.parse.unquote(c.get("value", "")).split(".")[0]
        seg += "=" * (-len(seg) % 4)
        try:
            raw = base64.urlsafe_b64decode(seg)
            try:
                raw = zlib.decompress(raw)
            except zlib.error:
                pass
            return json.loads(raw.decode("utf-8"))
        except Exception:                       # noqa: BLE001
            return {}
    return {}


def _captcha_answer(ctx, db_path: str, purpose: str) -> str | None:
    cid = _session_payload(ctx).get("cap_" + purpose)
    if not cid or not db_path or not os.path.exists(db_path):
        return None
    try:
        con = sqlite3.connect(db_path)
        try:
            row = con.execute("SELECT answer FROM captchas WHERE id=?", (cid,)).fetchone()
        finally:
            con.close()
    except sqlite3.Error:
        return None
    return row[0] if row else None


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="http://127.0.0.1:8849")
    ap.add_argument("-u", "--user", default="admin")
    ap.add_argument("-p", "--password", default="admin123")
    ap.add_argument("--out", default=os.path.join(BASE, "data", "shots"))
    ap.add_argument("--db", default=None,
                    help="SQLite 路径（默认 <repo>/data/usage.sqlite），用于取验证码答案")
    ap.add_argument("--full", action="store_true", help="截整页长图")
    ap.add_argument("--browser", default="", help="显式指定 chrome/msedge 可执行文件")
    ap.add_argument("--width", type=int, default=1440)
    ap.add_argument("--height", type=int, default=900)
    a = ap.parse_args()

    db_path = a.db or os.path.join(BASE, "data", "usage.sqlite")
    from playwright.sync_api import sync_playwright

    exe = a.browser or _find_browser()
    os.makedirs(a.out, exist_ok=True)
    problems = []
    with sync_playwright() as pw:
        try:
            br = pw.chromium.launch(executable_path=exe) if exe else pw.chromium.launch()
        except Exception as e:                  # noqa: BLE001
            print("[FAIL] 启动浏览器失败：%s" % e)
            print("       可用 --browser 显式指定，或用 `playwright install chromium` 装齐。")
            return 1
        print("浏览器：%s" % (exe or "playwright 默认"))
        ctx = br.new_context(viewport={"width": a.width, "height": a.height},
                             device_scale_factor=2, locale="zh-CN")
        page = ctx.new_page()
        errors = []
        page.on("console", lambda m: errors.append(m.text) if m.type == "error" else None)
        page.on("pageerror", lambda e: errors.append(str(e)))

        def shoot(name, path, label):
            errors.clear()
            page.goto(a.base + path, wait_until="networkidle")
            page.wait_for_timeout(900)          # 等 ECharts / 表格渲染稳下来
            page.screenshot(path=os.path.join(a.out, name), full_page=a.full)
            js_err = [e for e in errors if "favicon" not in e.lower()]
            if js_err:
                problems.append("%s: %s" % (label, js_err[:3]))
            print("[ok] %-22s %s%s" % (name, label,
                                       "" if not js_err else "  [JS错误] " + " | ".join(js_err[:3])))

        # 1) 未登录的两页
        for name, path, label, need_auth in CAPTURES:
            if need_auth:
                break
            shoot(name, path, label)

        # 2) 登录（策略为 always 时自动解验证码）
        page.goto(a.base + "/login", wait_until="networkidle")
        page.fill('input[name=username]', a.user)
        page.fill('input[name=password]', a.password)
        if page.query_selector('input[name=captcha]'):
            ans = _captcha_answer(ctx, db_path, "login")
            if not ans:
                print("[FAIL] 需要验证码但取不到答案（--db 是否指向本实例的库？）：%s" % db_path)
                br.close()
                return 1
            page.fill('input[name=captcha]', ans)
            print("[ok] 已用库里的答案通过验证码")
        page.click('button[type=submit]')
        page.wait_for_load_state("networkidle")
        if "/login" in page.url:
            print("[FAIL] 登录失败，后续截图无意义")
            br.close()
            return 1

        # 3) 登录后的页面
        for name, path, label, need_auth in CAPTURES:
            if not need_auth:
                continue
            shoot(name, path, label)

        # 4) 大屏页再点几个交互，确认控件联动不炸
        page.goto(a.base + "/dashboard", wait_until="networkidle")
        page.wait_for_timeout(1200)
        for sel in ["#segRange button", ".seg button"]:
            btns = page.query_selector_all(sel)
            if len(btns) > 1:
                errors.clear()
                btns[1].click()
                page.wait_for_timeout(900)
                page.screenshot(path=os.path.join(a.out, "08-dashboard-interact.png"),
                                full_page=a.full)
                js_err = [e for e in errors if "favicon" not in e.lower()]
                print("[ok] 08-dashboard-interact.png  点击 %s 第 2 项%s"
                      % (sel, "" if not js_err else "  [JS错误] %s" % js_err[:2]))
                if js_err:
                    problems.append("大屏交互: %s" % js_err[:2])
                break

        br.close()

    print("\n截图目录：%s" % a.out)
    if problems:
        print("发现问题：")
        for p in problems:
            print("  - " + p)
        return 1
    print("RESULT: 全页面截图完成，无 JS 报错")
    return 0


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