#!/usr/bin/env python3
"""真实 Chrome(CDP) 抓包：虎牙网页 Cookie 字段的来源分析（无需登录）。

两轮运行模式（同一个 user-data-dir）：
  round 1: 全新 profile 访问 www.huya.com 门户 + 直播间，记录每个请求携带的
           Cookie 头、每个响应的 Set-Cookie，定位每个字段的"首次出现位置"。
  round 2: 同一 profile 重启浏览器再访问，验证哪些字段是浏览器重启后
           首请求就直接携带的（浏览器持久态，而非登录/页面产生）。

用法:
  .venv/bin/python tools/huya_browser_harvest.py round1 --profile /tmp/huya-browser-p1
  .venv/bin/python tools/huya_browser_harvest.py round2 --profile /tmp/huya-browser-p1
  .venv/bin/python tools/huya_browser_harvest.py report --out evidence/browser_cookie_harvest

无第三方依赖：标准库 + websockets（直接走 CDP 协议）。
"""
from __future__ import annotations

import argparse
import asyncio
import json
import subprocess
import sys
import time
import urllib.parse
import urllib.request
from pathlib import Path

import websockets

CHROME = (
    Path.home()
    / "Library/Caches/ms-playwright/chromium-1234/chrome-mac-arm64/"
    "Google Chrome for Testing.app/Contents/MacOS/Google Chrome for Testing"
)

PAGES = [
    ("portal", "https://www.huya.com/"),
    ("live", "https://www.huya.com/30596253"),
]

# 与 core/huya/app_login.py::WEB_DEVICE_COOKIE_KEYS 一致
WEB_DEVICE_KEYS = (
    "guid", "udb_guiddata", "udb_deviceid", "game_did",
    "_qimei_uuid42", "udb_anobiztoken", "__yamid_new",
)

OUT_DIR = Path("evidence/browser_cookie_harvest")


class Cdp:
    def __init__(self, ws):
        self.ws = ws
        self._id = 0
        self._pending: dict[int, asyncio.Future] = {}
        self.listeners = []
        self.events = []

    async def __aenter__(self):
        self._task = asyncio.create_task(self._reader())
        return self

    async def __aexit__(self, *exc):
        self._task.cancel()

    async def _reader(self):
        try:
            async for raw in self.ws:
                msg = json.loads(raw)
                if "id" in msg and msg["id"] in self._pending:
                    self._pending.pop(msg["id"]).set_result(msg)
                elif "method" in msg:
                    self.events.append(msg)
                    for fn in list(self.listeners):
                        try:
                            fn(msg)
                        except Exception:  # noqa: BLE001, S110 - 监听器容错
                            pass
        except Exception:  # noqa: BLE001, S110 - socket closing on shutdown
            pass

    async def call(self, method, params=None):
        self._id += 1
        fut = asyncio.get_running_loop().create_future()
        self._pending[self._id] = fut
        await self.ws.send(json.dumps({"id": self._id, "method": method, "params": params or {}}))
        resp = await fut
        if "error" in resp:
            raise RuntimeError(f"{method}: {resp['error']}")
        return resp.get("result", {})

    async def eval_js(self, expr):
        r = await self.call("Runtime.evaluate", {
            "expression": expr, "returnByValue": True, "awaitPromise": True,
        })
        return r.get("result", {}).get("value")


def launch(profile: Path, port: int, headless: bool = False) -> tuple[subprocess.Popen, str]:
    args = [
        str(CHROME),
        "--no-first-run", "--no-default-browser-check",
        "--disable-background-networking", "--disable-component-update",
        "--disable-sync", "--metrics-recording-only", "--disable-extensions",
        "--disable-features=OptimizationHints,MediaRouter",
        f"--user-data-dir={profile}",
        f"--remote-debugging-port={port}",
        "--window-size=1400,900",
        "about:blank",
    ]
    if headless:
        args.insert(2, "--headless=new")
    proc = subprocess.Popen(
        args, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True
    )
    ws_url = None
    deadline = time.time() + 20
    while time.time() < deadline:
        # 1) 从 DevTools 输出的 ws url（browser 级）
        line = proc.stderr.readline()
        if "DevTools listening on" in line:
            browser_ws = line.strip().split(" ")[-1]
        else:
            browser_ws = None
        # 2) page target 的 ws url（Network 域需要挂在 page target 上）
        try:
            with urllib.request.urlopen(
                f"http://127.0.0.1:{port}/json/list", timeout=2
            ) as resp:
                targets = json.load(resp)
            pages = {t for t in targets if t.get("type") == "page" and t.get("webSocketDebuggerUrl")}
            if pages:
                ws_url = pages[0]["webSocketDebuggerUrl"]
                break
        except Exception:  # noqa: BLE001, S110 - 端口未就绪时轮询重试
            pass
        # 3) 兜底：browser 级 endpoint
        if browser_ws:
            ws_url = browser_ws
            break
    if not ws_url:
        proc.kill()
        raise RuntimeError("Chrome 调试端口未就绪")
    return proc, ws_url


async def navigate(cdp: Cdp, url: str, settle: float = 10.0):
    """导航并等待 load 事件 + 脚本静置时间。"""
    loaded = asyncio.get_running_loop().create_future()

    def on_event(msg):
        if msg.get("method") == "Page.loadEventFired" and not loaded.done():
            loaded.set_result(True)

    cdp.listeners.append(on_event)
    await cdp.call("Page.navigate", {"url": url})
    try:
        await asyncio.wait_for(loaded, timeout=30)
    except TimeoutError:
        pass
    finally:
        cdp.listeners.remove(on_event)
    await asyncio.sleep(settle)


def cookie_names_from_header(headers: dict) -> list[str]:
    raw = headers.get("Cookie") or headers.get("cookie") or ""
    if not raw:
        return []
    return [part.strip().split("=", 1)[0] for part in raw.split(";") if part.strip()]


async def capture_round(profile: Path, tag: str, port: int):
    """抓取一轮：导航序列 + 全量 Cookie 快照。"""
    proc, ws_url = launch(profile, port)
    try:
        async with websockets.connect(ws_url, max_size=64 * 1024 * 1024) as raw_ws:
            cdp = Cdp(raw_ws)
            async with cdp:
                await cdp.call("Network.enable")
                await cdp.call("Page.enable")
                await cdp.call("Runtime.enable")

                requests: dict[str, dict] = {}  # requestId -> 请求信息
                req_seq = 0

                def on_req(msg):
                    nonlocal req_seq
                    if msg.get("method") != "Network.requestWillBeSent":
                        return
                    params = msg["params"]
                    req = params["request"]
                    url = req.get("url", "")
                    rid = params.get("requestId", "")
                    entry = requests.setdefault(rid, {"url": url, "seq": req_seq})
                    ack_order = entry.get("seq")
                    if ack_order is None:
                        entry["seq"] = req_seq
                    else:
                        entry["seq"] = min(ack_order, req_seq)
                    req_seq += 1
                    entry.update({
                        "url": url,
                        "host": urllib.parse.urlparse(url).netloc,
                        "type": params.get("type", ""),
                    })
                    # 部分请求的 Cookie 头在 requestWillBeSent 就直接可见
                    seen = set(entry.get("cookie_names", []))
                    for name in cookie_names_from_header(req.get("headers", {})):
                        if name not in seen:
                            entry.setdefault("cookie_names", []).append(name)
                            seen.add(name)

                cdp.listeners.append(on_req)

                def on_req_extra(msg):
                    if msg.get("method") != "Network.requestWillBeSentExtraInfo":
                        return
                    params = msg["params"]
                    rid = params.get("requestId", "")
                    entry = requests.setdefault(rid, {"url": "?"})
                    # associatedCookies: 该请求实际附带（含 httpOnly/持久）的 Cookie
                    seen = set(entry.get("cookie_names", []))
                    for ck in params.get("associatedCookies", []):
                        name = ck.get("cookie", {}).get("name", "")
                        if name and name not in seen:
                            entry.setdefault("cookie_names", []).append(name)
                            seen.add(name)
                    # headers 里的 Cookie 头（同一信息的冗余来源）
                    for name in cookie_names_from_header(params.get("headers", {})):
                        if name not in seen:
                            entry.setdefault("cookie_names", []).append(name)
                            seen.add(name)

                cdp.listeners.append(on_req_extra)

                resp_extra: list[dict] = []

                def on_resp_extra(msg):
                    if msg.get("method") != "Network.responseReceivedExtraInfo":
                        return
                    params = msg["params"]
                    rid = params.get("requestId", "")
                    entry = requests.setdefault(rid, {"url": "?"})
                    url = entry.get("url", "?")
                    headers = params.get("headers", {})
                    set_cookie = headers.get("Set-Cookie") or headers.get("set-cookie") or []
                    if isinstance(set_cookie, str):
                        set_cookie = [set_cookie]
                    for sc in set_cookie:
                        name = sc.split("=", 1)[0].strip()
                        resp_extra.append({"cookie": name, "url": url, "meta": sc[:160]})

                cdp.listeners.append(on_resp_extra)

                doc_cookies: dict[str, dict] = {}
                snapshots: list[dict] = []
                for name, url in PAGES:
                    ts = time.strftime("%H:%M:%S")
                    doc_cookies[name] = {
                        "before": await cdp.eval_js("document.cookie"),
                        "ts": ts,
                    }
                    await navigate(cdp, url)
                    doc_cookies[name]["after"] = await cdp.eval_js("document.cookie")
                    snapshots.append({"page": name, "url": url, "time": ts,
                                      "cookies": await cdp.call("Network.getAllCookies")})
                    print(f"[harvest] {tag}/{name} done: {url}  doc_cookie_len="
                          f"{len(doc_cookies[name]['after'] or '')}", flush=True)

                # 按事件到达次序输出请求列表
                req_list = sorted(requests.values(), key=lambda r: r.get("seq", 10**9))

                out = OUT_DIR / tag
                out.mkdir(parents=True, exist_ok=True)
                (out / "requests.json").write_text(
                    json.dumps(req_list, ensure_ascii=False, indent=1)
                )
                (out / "set_cookies.json").write_text(
                    json.dumps(resp_extra, ensure_ascii=False, indent=1)
                )
                (out / "doc_cookies.json").write_text(
                    json.dumps(doc_cookies, ensure_ascii=False, indent=1)
                )
                (out / "snapshots.json").write_text(
                    json.dumps(snapshots, ensure_ascii=False, indent=1)
                )
                print(f"[harvest] {tag} 已写出到 {out}/")
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()


async def probe_anon(profile: Path, port: int):
    """复用已有 profile：复放 udblgn 匿名流程，解密请求/响应体。

    抓取会话内 anonymousLogin 与 web/middle 的 postData 与响应体，
    然后带全套 Cookie 用 requests 独立重放 anonymousLogin，检查
    服务端是否在匿名流程中下发 guid/udb_guiddata/匿名 token 等字段。
    """
    proc, ws_url = launch(profile, port)
    try:
        async with websockets.connect(ws_url, max_size=64 * 1024 * 1024) as raw_ws:
            cdp = Cdp(raw_ws)
            async with cdp:
                await cdp.call("Network.enable")
                await cdp.call("Page.enable")
                await cdp.call("Runtime.enable")

                anon: list[dict] = []  # anonymousLogin / middle 请求详情
                rid_map: dict[str, dict] = {}

                def on_anon_req(msg):
                    if msg.get("method") != "Network.requestWillBeSent":
                        return
                    params = msg["params"]
                    req = params["request"]
                    url = req.get("url", "")
                    if "anonymousLogin" not in url and "/web/middle/" not in url:
                        return
                    entry = {
                        "url": url,
                        "method": req.get("method", ""),
                        "postData": req.get("postData", ""),
                        "headers": {k: v for k, v in req.get("headers", {}).items()
                                    if k.lower() in ("cookie", "content-type", "referer", "origin")},
                        "body": None,
                        "status": None,
                    }
                    rid_map[params.get("requestId", "")] = entry
                    anon.append(entry)
                    print(f"[probe] 抓到匿名请求: {req.get('method')} {url[:120]}", flush=True)

                cdp.listeners.append(on_anon_req)

                def on_anon_resp(msg):
                    if msg.get("method") != "Network.responseReceived":
                        return
                    rid = msg["params"].get("requestId", "")
                    entry = rid_map.get(rid)
                    if entry is None:
                        return
                    entry["status"] = msg["params"].get("response", {}).get("status")

                cdp.listeners.append(on_anon_resp)

                await navigate(cdp, PAGES[0][1], settle=14.0)

                # 抓响应体
                for rid, entry in rid_map.items():
                    try:
                        body = await cdp.call("Network.getResponseBody", {"requestId": rid})
                        entry["body"] = body.get("body", "")
                        print(f"[probe] 响应体 {len(entry['body'])}B status={entry['status']}",
                              flush=True)
                    except Exception as exc:  # noqa: BLE001
                        print(f"[probe] 响应体获取失败: {exc}", flush=True)

                # 全量 Cookie -> jar（含 httpOnly；Domain 匹配规则简化：host 以 domain 结尾即适用）
                cookies = (await cdp.call("Network.getAllCookies")).get("cookies", [])
                jar = {}
                for c in cookies:
                    dom = str(c.get("domain", "")).lstrip(".")
                    if "huya.com" not in dom:
                        continue
                    jar.setdefault(dom, {})[c["name"]] = c["value"]
                print(f"[probe] cookie jar 域名: {sorted(jar)}")

                out = OUT_DIR / "probe"
                out.mkdir(parents=True, exist_ok=True)

                def cookie_header(host: str) -> str:
                    parts = []
                    for dom, kv in jar.items():
                        if host.endswith(dom) or host == dom:
                            parts.extend(f"{k}={v}" for k, v in kv.items())
                    return "; ".join(parts)

                import requests as pyrequests

                replay = []
                for entry in anon:
                    host = urllib.parse.urlparse(entry["url"]).netloc
                    headers = {"User-Agent": (
                        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
                        "AppleWebKit/537.36 (KHTML, like Gecko) "
                        "Chrome/149.0.0.0 Safari/537.36"
                    )}
                    ck = cookie_header(host)
                    if ck:
                        headers["Cookie"] = ck
                    if entry["postData"]:
                        headers.setdefault("Content-Type", "application/x-www-form-urlencoded")
                    try:
                        resp = await asyncio.to_thread(
                            pyrequests.request,
                            entry["method"] or "GET",
                            entry["url"],
                            headers=headers,
                            data=entry["postData"] or None,
                            timeout=15,
                            allow_redirects=False,
                        )
                        replay.append({
                            "url": entry["url"][:160],
                            "method": entry["method"],
                            "status": resp.status_code,
                            "set_cookie": resp.headers.get("Set-Cookie", ""),
                            "location": resp.headers.get("Location", ""),
                            "body_head": resp.text[:600],
                            "body_len": len(resp.content),
                        })
                        print(f"[probe] 重放 {entry['method']} {entry['url'][:90]} "
                              f"-> {resp.status_code} setcookie={bool(resp.headers.get('Set-Cookie'))}",
                              flush=True)
                    except Exception as exc:  # noqa: BLE001
                        replay.append({"url": entry["url"][:160], "error": str(exc)})
                        print(f"[probe] 重放失败: {exc}", flush=True)

                (out / "anon_requests.json").write_text(
                    json.dumps(anon, ensure_ascii=False, indent=1)
                )
                (out / "replay.json").write_text(
                    json.dumps(replay, ensure_ascii=False, indent=1)
                )
                print(f"[probe] 已写出到 {out}/")
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()


def build_report():
    """汇总 round1/round2，生成 Cookie 字段来源矩阵报告。"""
    r1 = OUT_DIR / "round1"
    r2 = OUT_DIR / "round2"
    if not (r1 / "requests.json").exists():
        print("先运行 round1", file=sys.stderr)
        sys.exit(2)

    req1 = json.loads((r1 / "requests.json").read_text())
    req2 = json.loads((r2 / "requests.json").read_text()) if (r2 / "requests.json").exists() else []
    setc1 = json.loads((r1 / "set_cookies.json").read_text())
    snap2 = json.loads((r2 / "snapshots.json").read_text()) if (r2 / "snapshots.json").exists() else []
    doc1 = json.loads((r1 / "doc_cookies.json").read_text())

    # 每个 cookie 名在 round1 的首次出现位置（按请求顺序）
    first_seen: dict[str, dict] = {}
    for i, req in enumerate(req1):
        for name in req.get("cookie_names", []):
            if name not in first_seen:
                first_seen[name] = {"idx": i, "host": req["host"], "url": req["url"], "type": req["type"]}

    # Set-Cookie 写入的 cookie
    set_from: dict[str, str] = {}
    for sc in setc1:
        set_from.setdefault(sc["cookie"], sc["url"])

    # round2 首请求携带（持久化证据）
    persistent: set[str] = set()
    if req2:
        for req in req2:
            for name in req.get("cookie_names", []):
                persistent.add(name)

    # 最终状态（round2 快照）
    final_cookies: dict[str, dict] = {}
    if snap2:
        for snap in snap2:
            for c in snap.get("cookies", {}).get("cookies", []):
                final_cookies.setdefault(c["name"], c)

    def fmt_first(info: dict | None):
        if not info:
            return "未见来源"
        return f'#{info["idx"]} {info["host"]} ({info["type"]})'

    lines = [
        "# 浏览器 Cookie 字段来源分析（真实 Chrome 抓包，未登录）",
        "",
        f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}",
        f"页面序列: {', '.join(f'{n}={u}' for n, u in PAGES)}",
        "",
        "## 9.1 网页设备态必需字段（WEB_DEVICE_COOKIE_KEYS 对齐）",
        "",
        "| 字段 | 首次出现(round1) | 是否 Set-Cookie 下发 | round2 重启后首请求携带 | 最终 httpOnly |",
        "| --- | --- | --- | --- | --- |",
    ]
    for name in WEB_DEVICE_KEYS:
        info = first_seen.get(name)
        set_url = set_from.get(name)
        persisted = "✅" if name in persistent else "—"
        http_only = "✅" if final_cookies.get(name, {}).get("httpOnly") else "—"
        lines.append(
            f"| `{name}` | {fmt_first(info)} | {set_url or '—'} | {persisted} | {http_only} |"
        )

    lines += [
        "",
        "## round1 全部请求（前 60，含 Cookie 头字段数）",
        "",
        "| # | host | type | Cookie 字段 |",
        "| --- | --- | --- | --- |",
    ]
    for i, req in enumerate(req1[:60]):
        names = ", ".join(req.get("cookie_names", [])[:12])
        names_all = req.get("cookie_names", [])
        extra = f" +{len(names_all)-12}" if len(names_all) > 12 else ""
        lines.append(f"| {i} | {req['host']} | {req['type']} | {names}{extra} |")

    lines += [
        "",
        "## Set-Cookie 写入清单（round1）",
        "",
    ]
    for sc in setc1:
        lines.append(f"- `{sc['cookie']}` <- {sc['url'][:100]}  `{sc['meta'][:100]}`")

    lines += [
        "",
        "## round1 之后 document.cookie（页面脚本可得，非 httpOnly）",
        "",
    ]
    seen = set()
    for page, v in doc1.items():
        after = v.get("after") or ""
        names = [p.split("=", 1)[0].strip() for p in after.split(";") if p.strip()]
        diff = [n for n in names if n not in seen]
        seen.update(names)
        lines.append(f"- {page}: before={bool(v.get('before'))} after={len(names)} 字段"
                     f" 新增={', '.join(diff) if diff else '无'}")
    lines += ["", "## round2 重启后首请求即携带的字段（浏览器持久态）", ""]
    if req2:
        pk = []
        for req in req2[:8]:
            names = req.get("cookie_names", [])
            pk.append(f"- {req['host']} ({req['type']}): {', '.join(names[:15])}")
        lines += pk or ["（无持久 Cookie）"]
    else:
        lines += ["（未运行 round2）"]
    lines += ["", "## 其他出现过的字段（round1 全字段集合）", ""]
    all_names = sorted(
        {n for r in req1 for n in r.get("cookie_names", [])} | set(set_from) | set(seen)
    )
    lines.append("`" + "`, `".join(all_names) + "`")

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    (OUT_DIR / "report.md").write_text("\n".join(lines) + "\n")
    print(f"[harvest] 报告已写出: {OUT_DIR / 'report.md'}")
    print("\n".join(lines[:40]))


def main():
    global OUT_DIR
    ap = argparse.ArgumentParser()
    ap.add_argument("mode", choices=["round1", "round2", "probe", "report"])
    ap.add_argument("--profile", default="/tmp/huya-browser-p1")
    ap.add_argument("--port", type=int, default=9333)
    ap.add_argument("--out", default=str(OUT_DIR))
    args = ap.parse_args()
    OUT_DIR = Path(args.out)
    profile = Path(args.profile)
    if args.mode in ("round1", "round2"):
        if args.mode == "round1" and (profile / "Default").exists():
            print("round1 要求全新 profile；请删除", profile)
            sys.exit(2)
        if args.mode == "round2" and not (profile / "Default").exists():
            print("round2 需要先运行 round1（同一 profile）")
            sys.exit(2)
        asyncio.run(capture_round(profile, args.mode, args.port))
    elif args.mode == "probe":
        # 匿名流程只在"无 udb_deviceid 持久态"的全新 profile 触发；
        # probe 模式内部使用临时全新 profile，跑完自动清理。
        import shutil
        import tempfile

        tmp = tempfile.mkdtemp(prefix="huya-probe-")
        try:
            asyncio.run(probe_anon(Path(tmp), args.port))
        finally:
            shutil.rmtree(tmp, ignore_errors=True)
    else:
        build_report()


if __name__ == "__main__":
    main()