"""解析 ``frida_probe_passport_wcfg.js`` 的 ``@@WCFG`` 抓证日志。

目标是把动态证据整理成机器可读结果：

``/f/a/p`` 调用 → ``wcfg["a_y_q_z"]`` 写入 → ``WeaponHI.dd(21)`` 读取
→ 可回灌到 ``app_login_fields`` 的 ``passport_account_image``。
"""

from __future__ import annotations

import argparse
import glob
import json
import re
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from core.privacykit_encrypt import diagnose_passport_account_image  # noqa: E402


WCFG_PREFIX = "@@WCFG "
PASSPORT_KEY = "a_y_q_z"
VIMG_RE = re.compile(r"VIMG_[A-Za-z0-9+/_=-]+(?:\$AI_[0-9a-fA-F]{32})?")
FAP_RESPONSE_TAGS = {
    "R1_A_AFTER",
    "R1_B_AFTER",
    "I_A_K1_AFTER",
    "I_A_CTX_K1_AFTER",
    "Y0_A_K1_AFTER",
    "OKHTTP_RESPONSE_BODY_STRING",
}
FAP_REQUEST_TAGS = {
    "R1_A_BEFORE",
    "R1_B_BEFORE",
    "I_A_K1_BEFORE",
    "I_A_CTX_K1_BEFORE",
    "Y0_A_K1_BEFORE",
    "OKHTTP_REQUEST_BUILD",
}


def _latest_wcfg_log() -> Path:
    logs = sorted(_default_wcfg_logs(), key=lambda p: p.stat().st_mtime, reverse=True)
    if not logs:
        raise FileNotFoundError("out 下没有 probe_passport_wcfg_*.log / probe_multi_*.log")
    return logs[0]


def _default_wcfg_logs() -> list[Path]:
    """默认扫描 passport 单进程日志和 multi-attach 子日志。"""

    out_dir = ROOT / "out"
    logs = [*out_dir.glob("probe_passport_wcfg_*.log"), *out_dir.glob("probe_multi_*.log")]
    return sorted({path.resolve() for path in logs if path.is_file()}, key=lambda p: p.stat().st_mtime, reverse=True)


def _resolve_log_glob(pattern: str) -> list[Path]:
    base_pattern = pattern
    if not Path(pattern).is_absolute():
        base_pattern = str(ROOT / pattern)
    paths = [Path(item) for item in glob.glob(base_pattern)]
    return sorted({path.resolve() for path in paths if path.is_file()}, key=lambda p: p.stat().st_mtime, reverse=True)


def _repair_nul_interleaved_text(text: str) -> str:
    """修复 PowerShell Job/Tee 产生的 NUL 交错日志。

    现象：日志开头由 ``Out-File -Encoding UTF8`` 写入，后续 Frida 输出在
    某些环境下会以 UTF-16LE 字节追加；用 UTF-8 解码不会报错，但会变成
    ``@\x00@\x00W...``，导致 ``@@WCFG`` 前缀匹配不到。日志正文是 JSONL，
    NUL 字符没有业务意义，因此检测到明显交错后直接剥离 NUL。
    """

    if "\x00" not in text:
        return text
    nul_count = text.count("\x00")
    if nul_count < max(16, len(text) // 20):
        return text
    repaired = text.replace("\x00", "")
    return repaired if repaired.count(WCFG_PREFIX) >= text.count(WCFG_PREFIX) else text


def _read_text(path: Path) -> str:
    raw = path.read_bytes()
    if raw.startswith(b"\xff\xfe") or raw.startswith(b"\xfe\xff"):
        return _repair_nul_interleaved_text(raw.decode("utf-16", errors="replace"))
    for enc in ("utf-8", "gbk"):
        try:
            return _repair_nul_interleaved_text(raw.decode(enc, errors="strict"))
        except UnicodeDecodeError:
            pass
    return _repair_nul_interleaved_text(raw.decode("utf-8", errors="replace"))


def iter_wcfg_events(log_path: Path):
    for line_no, line in enumerate(_read_text(log_path).splitlines(), 1):
        index = line.find(WCFG_PREFIX)
        if index < 0:
            continue
        payload = line[index + len(WCFG_PREFIX):]
        try:
            ev = json.loads(payload)
        except json.JSONDecodeError:
            continue
        ev["_line_no"] = line_no
        ev["_source_log"] = str(log_path)
        yield ev


def _event_value(ev: dict[str, Any], primary: str = "value") -> tuple[str, bool]:
    """返回事件里的完整值和是否被截断。

    新版 Frida 脚本会输出 ``value_full``；旧日志只有 ``value`` 时，如果包含
    ``...(len=``，则只能作为证据，不能回灌。
    """

    full_key = f"{primary}_full"
    if ev.get(full_key) is not None:
        return str(ev.get(full_key) or ""), False
    value = str(ev.get(primary) or "")
    return value, "...(len=" in value


def _looks_vimg(value: str) -> bool:
    return str(value or "").startswith("VIMG_")


def _looks_final_ticket(value: str) -> bool:
    text = str(value or "")
    return text.startswith("VIMG_") and "$AI_" in text


def _json_string_values(obj: Any):
    if isinstance(obj, str):
        yield obj
    elif isinstance(obj, dict):
        for value in obj.values():
            yield from _json_string_values(value)
    elif isinstance(obj, list):
        for value in obj:
            yield from _json_string_values(value)


def _extract_vimg_values(text: str) -> list[str]:
    """从响应体/返回体里抽取可候选的 ``VIMG_`` 值。"""

    raw = str(text or "")
    if not raw:
        return []
    values: list[str] = []
    if _looks_vimg(raw):
        values.append(raw)
    try:
        decoded = json.loads(raw)
    except Exception:  # noqa: BLE001 - 不是 JSON 时走正则
        decoded = None
    if decoded is not None:
        for item in _json_string_values(decoded):
            if _looks_vimg(item):
                values.append(item)
    values.extend(match.group(0) for match in VIMG_RE.finditer(raw))

    # 保持顺序去重。
    deduped: list[str] = []
    seen: set[str] = set()
    for value in values:
        if value and value not in seen:
            seen.add(value)
            deduped.append(value)
    return deduped


def _value_preview(value: str, limit: int = 80) -> str:
    text = str(value or "")
    if len(text) <= limit:
        return text
    return f"{text[:limit]}...(len={len(text)})"


def _stack_top(stack: str) -> str:
    for line in str(stack or "").splitlines():
        line = line.strip()
        if line.startswith("at "):
            return line
    return ""


def _as_int(value: Any, default: int = 0) -> int:
    try:
        return int(value)
    except (TypeError, ValueError):
        return default


def _is_fap_context(ev: dict[str, Any]) -> bool:
    return "/f/a/p" in str(ev.get("url") or "") or "/f/a/p" in str(ev.get("endpoint_hint") or "")


def _candidate_priority(ev: dict[str, Any], value: str, truncated: bool, source: str = "value") -> int:
    if not _looks_vimg(value) or truncated:
        return 0
    tag = str(ev.get("tag") or "")
    key = str(ev.get("key") or "")
    name = str(ev.get("name") or "")
    final_bonus = 50 if _looks_final_ticket(value) else 0
    if key == PASSPORT_KEY and ("WRITE" in tag or tag.startswith("EDITOR_")):
        return 100 + final_bonus
    if tag == "WEAPON_DD" and _as_int(ev.get("type")) == 21:
        return 90 + final_bonus
    if key == PASSPORT_KEY and ("READ" in tag or "GET" in tag):
        return 80 + final_bonus
    if tag in {"OKHTTP_FORM_ADD", "OKHTTP_FORM_ADD_ENCODED", "OKHTTP2_FORM_ADD", "OKHTTP2_FORM_ADD_ENCODED"} and name == "passport_account_image":
        return 75 + final_bonus
    if tag in {"EDITOR_PUT_STRING_IMPL", "EDITOR_PUT_STRING_IFACE"}:
        return 70 + final_bonus
    if source in {"ret", "body"} and tag in FAP_RESPONSE_TAGS:
        return (65 if _is_fap_context(ev) else 45) + final_bonus
    if tag == "WEAPON_B_UPLOAD_ENCRYPT":
        return 30 + final_bonus
    return 10 + final_bonus


def _timeline_event(ev: dict[str, Any]) -> dict[str, Any]:
    value, truncated = _event_value(ev)
    ret, ret_truncated = _event_value(ev, "ret")
    body, body_truncated = _event_value(ev, "body")
    out: dict[str, Any] = {
        "source_log": ev.get("_source_log"),
        "line_no": ev.get("_line_no"),
        "seq": ev.get("seq"),
        "ts": ev.get("ts"),
        "tag": ev.get("tag"),
    }
    for key in ("key", "name", "type", "prefs", "url", "overload", "endpoint_hint", "method", "code", "process_name", "reason"):
        if ev.get(key) not in (None, ""):
            out[key] = ev.get(key)
    if value:
        out["value_preview"] = _value_preview(value)
        out["value_len"] = int(ev.get("value_len") or len(value))
        out["value_truncated"] = truncated
    if ret:
        out["ret_preview"] = _value_preview(ret)
        out["ret_truncated"] = ret_truncated
    if body:
        out["body_preview"] = _value_preview(body)
        out["body_truncated"] = body_truncated
    if ev.get("stack"):
        out["stack_top"] = _stack_top(str(ev.get("stack") or ""))
    return out


def _candidate_records(ev: dict[str, Any]) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    tag = str(ev.get("tag") or "")
    key = str(ev.get("key") or "")
    name = str(ev.get("name") or "")
    for source in ("value", "ret", "body"):
        source_text, truncated = _event_value(ev, source)
        if not source_text:
            continue
        extracted = _extract_vimg_values(source_text)
        if not extracted and _looks_vimg(source_text):
            extracted = [source_text]
        for value in extracted:
            priority = _candidate_priority(ev, value, truncated, source)
            if not priority and not _looks_vimg(value):
                continue
            diag = diagnose_passport_account_image(value) if not truncated else {"ok": False, "error": "value truncated"}
            records.append(
                {
                    "priority": priority,
                    "source_log": ev.get("_source_log"),
                    "line_no": ev.get("_line_no"),
                    "seq": ev.get("seq"),
                    "ts": ev.get("ts"),
                    "tag": tag,
                    "key": key,
                    "name": name,
                    "source": source,
                    "truncated": truncated,
                    "value": value if priority else "",
                    "value_preview": _value_preview(value),
                    "value_len": len(value),
                    "diagnosis": diag,
                }
            )
    return records


def _extract_wcfg_evidence_from_events(
    events: list[dict[str, Any]],
    *,
    source_log: str,
    source_logs: list[str] | None = None,
) -> dict[str, Any]:
    """从已解析事件里提取 passport 票据和链路证据。"""

    flags = {
        "has_fap_call": False,
        "has_mobile_checker_url": False,
        "has_upload_encrypt": False,
        "has_wcfg_open": False,
        "has_a_y_q_z_write": False,
        "has_a_y_q_z_read": False,
        "has_dd21_read": False,
        "has_fap_request_body": False,
        "has_fap_response_body": False,
        "has_passport_form_add": False,
    }
    candidates: list[dict[str, Any]] = []
    timeline: list[dict[str, Any]] = []
    process_names: list[str] = []
    seen_process_names: set[str] = set()

    interesting_tags = {
        "OKHTTP_BUILDER_URL",
        "OKHTTP_REQUEST_BUILD",
        "OKHTTP_NEWCALL",
        "OKHTTP_RESPONSE_BODY",
        "OKHTTP_RESPONSE_BODY_STRING",
        "WEAPON_B_UPLOAD_ENCRYPT",
        "EDITOR_PUT_STRING_IFACE",
        "EDITOR_PUT_STRING_IMPL",
        "Z0_WRITE_A_BOOL_STRING_STRING",
        "Z0_WRITE_B_STRING_STRING_BOOL",
        "PREFS_GET_STRING",
        "Z0_READ_A_STRING_STRING",
        "WEAPON_DD",
        "PREFS_OPEN",
        "PROCESS_NAME",
        "PROCESS_SKIP",
        "HOOK_SKIPPED",
        "R1_A_BEFORE",
        "R1_A_AFTER",
        "R1_B_BEFORE",
        "R1_B_AFTER",
        "I_A_K1_BEFORE",
        "I_A_K1_AFTER",
        "I_A_CTX_K1_BEFORE",
        "I_A_CTX_K1_AFTER",
        "Y0_A_K1_BEFORE",
        "Y0_A_K1_AFTER",
        "OKHTTP_FORM_ADD",
        "OKHTTP_FORM_ADD_ENCODED",
        "OKHTTP2_FORM_ADD",
        "OKHTTP2_FORM_ADD_ENCODED",
    }

    for ev in events:
        tag = str(ev.get("tag") or "")
        key = str(ev.get("key") or "")
        name = str(ev.get("name") or "")
        url = str(ev.get("url") or "")
        endpoint_hint = str(ev.get("endpoint_hint") or "")
        value, truncated = _event_value(ev)
        ret, _ret_truncated = _event_value(ev, "ret")
        body, _body_truncated = _event_value(ev, "body")
        process_name = str(ev.get("process_name") or "")

        if process_name and process_name not in seen_process_names:
            seen_process_names.add(process_name)
            process_names.append(process_name)

        if "/f/a/p" in url or "/f/a/p" in endpoint_hint:
            flags["has_fap_call"] = True
        if ("/f/a/p" in url or "/f/a/p" in endpoint_hint) and tag in FAP_REQUEST_TAGS and body:
            flags["has_fap_request_body"] = True
        if ("/f/a/p" in url or "/f/a/p" in endpoint_hint) and tag in FAP_RESPONSE_TAGS and (ret or body):
            flags["has_fap_response_body"] = True
        if "mobile/checker" in url:
            flags["has_mobile_checker_url"] = True
        if tag == "WEAPON_B_UPLOAD_ENCRYPT":
            flags["has_upload_encrypt"] = True
        if tag in {"OKHTTP_FORM_ADD", "OKHTTP_FORM_ADD_ENCODED", "OKHTTP2_FORM_ADD", "OKHTTP2_FORM_ADD_ENCODED"} and name == "passport_account_image":
            flags["has_passport_form_add"] = True
        if tag == "PREFS_OPEN" and str(ev.get("name") or "") == "wcfg":
            flags["has_wcfg_open"] = True
        if key == PASSPORT_KEY and ("WRITE" in tag or tag.startswith("EDITOR_")):
            flags["has_a_y_q_z_write"] = True
        if key == PASSPORT_KEY and ("READ" in tag or "GET" in tag):
            flags["has_a_y_q_z_read"] = True
        if tag == "WEAPON_DD" and _as_int(ev.get("type")) == 21:
            flags["has_dd21_read"] = True

        candidates.extend(_candidate_records(ev))

        if (
            tag in interesting_tags
            or key == PASSPORT_KEY
            or _looks_vimg(value)
            or _looks_vimg(ret)
            or _looks_vimg(body)
            or "/f/a/p" in url
            or "/f/a/p" in endpoint_hint
        ):
            timeline.append(_timeline_event(ev))

    candidates.sort(key=lambda item: (int(item["priority"]), int(item.get("seq") or 0)), reverse=True)
    selected = candidates[0] if candidates and int(candidates[0]["priority"]) > 0 else {}
    passport_value = str(selected.get("value") or "")
    app_fields_patch = {}
    if passport_value and not selected.get("truncated") and _looks_final_ticket(passport_value):
        app_fields_patch = {
            "passport_account_image": passport_value,
            "request_passport_account_image": passport_value,
            "checker_passport_account_image": passport_value,
        }

    chain_complete = bool(
        flags["has_fap_call"]
        and flags["has_a_y_q_z_write"]
        and flags["has_dd21_read"]
        and passport_value
        and _looks_final_ticket(passport_value)
    )

    return {
        "source_log": source_log,
        "source_logs": source_logs or ([source_log] if source_log else []),
        "process_names": process_names,
        "event_count": len(events),
        "flags": flags,
        "chain_complete": chain_complete,
        "passport_account_image": passport_value,
        "selected": {k: v for k, v in selected.items() if k != "value"},
        "app_fields_patch": app_fields_patch,
        "candidate_count": len(candidates),
        "candidates": [{k: v for k, v in item.items() if k != "value"} for item in candidates[:20]],
        "timeline": timeline,
    }


def extract_wcfg_evidence(log_path: Path) -> dict[str, Any]:
    """提取单个日志里的 passport 票据和链路证据。"""

    events = list(iter_wcfg_events(log_path))
    return _extract_wcfg_evidence_from_events(events, source_log=str(log_path), source_logs=[str(log_path)])


def extract_wcfg_evidence_many(log_paths: list[Path]) -> dict[str, Any]:
    """聚合多个 Frida 日志，适配 multi-attach 每进程一个日志的场景。"""

    resolved = [path for path in log_paths if path.exists() and path.is_file()]
    events: list[dict[str, Any]] = []
    for path in resolved:
        events.extend(iter_wcfg_events(path))
    source_logs = [str(path) for path in resolved]
    source_log = source_logs[0] if len(source_logs) == 1 else f"<multi:{len(source_logs)} logs>"
    return _extract_wcfg_evidence_from_events(events, source_log=source_log, source_logs=source_logs)


def _masked(value: str, head: int = 14, tail: int = 18) -> str:
    text = str(value or "")
    if not text:
        return ""
    if len(text) <= head + tail + 6:
        return f"<set len={len(text)}>"
    return f"{text[:head]}...{text[-tail:]}(len={len(text)})"


def _summary(result: dict[str, Any], show: bool) -> dict[str, Any]:
    out = dict(result)
    if not show and out.get("passport_account_image"):
        out["passport_account_image"] = _masked(str(out["passport_account_image"]))
    if not show and out.get("app_fields_patch"):
        out["app_fields_patch"] = {k: _masked(str(v)) for k, v in out["app_fields_patch"].items()}
    return out


def update_app_fields(app_fields_path: Path, patch: dict[str, str]) -> dict[str, Any]:
    if not patch:
        raise ValueError("没有可回灌的完整 passport_account_image")
    data = json.loads(app_fields_path.read_text(encoding="utf-8")) if app_fields_path.exists() else {}
    if not isinstance(data, dict):
        raise ValueError("app-fields JSON root must be object")
    data.update(patch)
    app_fields_path.parent.mkdir(parents=True, exist_ok=True)
    app_fields_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
    return data


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="解析 passport/wcfg Frida 抓证日志")
    p.add_argument("--log", default="", help="probe_passport_wcfg 日志；不传取 out 下最新")
    p.add_argument("--log-glob", default="", help="按 glob 聚合多个日志，如 out/probe_multi_*.log")
    p.add_argument("--all-logs", action="store_true", help="聚合 out 下 passport_wcfg/probe_multi 全部日志")
    p.add_argument("--out", default="out/passport_wcfg_evidence_latest.json", help="证据 JSON 输出路径")
    p.add_argument("--app-fields", default="", help="可选：app_login_fields JSON 路径")
    p.add_argument("--update-app-fields", action="store_true", help="把完整最终票据回灌到 --app-fields")
    p.add_argument("--show", action="store_true", help="终端显示完整票据；默认脱敏")
    return p


def main() -> int:
    args = build_parser().parse_args()
    if args.log_glob:
        log_paths = _resolve_log_glob(args.log_glob)
        if not log_paths:
            raise SystemExit(f"--log-glob 未匹配到日志: {args.log_glob}")
        result = extract_wcfg_evidence_many(log_paths)
    elif args.all_logs:
        log_paths = _default_wcfg_logs()
        if not log_paths:
            raise SystemExit("out 下没有 probe_passport_wcfg_*.log / probe_multi_*.log")
        result = extract_wcfg_evidence_many(log_paths)
    else:
        log_path = Path(args.log) if args.log else _latest_wcfg_log()
        if not log_path.is_absolute():
            log_path = ROOT / log_path
        result = extract_wcfg_evidence(log_path)

    out_path = Path(args.out)
    if not out_path.is_absolute():
        out_path = ROOT / out_path
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")

    if args.update_app_fields:
        if not args.app_fields:
            raise SystemExit("--update-app-fields 需要同时传 --app-fields")
        app_fields_path = Path(args.app_fields)
        if not app_fields_path.is_absolute():
            app_fields_path = ROOT / app_fields_path
        update_app_fields(app_fields_path, result.get("app_fields_patch") or {})
        print(f"[OK] 已回灌 {app_fields_path}")

    print(f"[OK] 写入 {out_path}")
    print(json.dumps(_summary(result, args.show), ensure_ascii=False, indent=2))
    return 0


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