#!/usr/bin/env python3
"""wupData信封纯Python改造器 —— 摆脱设备依赖的核心。

背景(逆向结论, 见 work/envelope/dissector2.py):
  信封 = getQUrlData()产出的WUP请求帧, 布局与登录包同族:
    [u32 total][ver][seq] + TAF(tag3..tag10)
    tag7 sBuffer = MAP{"_wup_data": <struct>, "wupudbrequest_v0": <bytes>}
    _wup_data struct:
      t0 头struct(t2=meta JSON 含 session/associationId)
      t1 设备信息 / t2 屏幕指纹   ← 全是静态设备常量
      t3 = uid (INT64, 大端)      ← 与证书内cred的账号必须一致!
      t4 = b64(cert) (STR4[260])  ← 可用 tools/cert_forge 离线铸造替换

  服务端校验规则(实验测定, 2026-08-25):
    - cert 有效即可离线铸造(nonce 新鲜随机即可, 无需设备)
    - uid(t3) 必须与 cert 内 cred 所属账号一致, 否则 40020
    - 信封整体可长期复用(无短TTL), session 等字段服务端不实时校验

用法:
  from envelope_forge import Envelope
  env = Envelope.load()                       # 默认 evidence/cert_keycap.json
  raw = env.patch_uid(1199666911746).patch_cert(forged_cert_bytes).raw
  wup_b64 = base64.b64encode(raw).decode()
"""
from __future__ import annotations

import base64
import json
import struct
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent

INT8, INT16, INT32, INT64 = 0x00, 0x01, 0x02, 0x03
STRING1, STRING4 = 0x06, 0x07
MAP, LIST = 0x08, 0x09
STRUCT_BEGIN, STRUCT_END = 0x0A, 0x0B
ZERO, SIMPLE_LIST = 0x0C, 0x0D


def _read_len_int(d: bytes, p: int) -> tuple[int, int]:
    """tag0优化int(长度), 返回(值, 新偏移)。"""
    dt = d[p] & 0x0F
    if dt == ZERO:
        return 0, p + 1
    if dt == INT8:
        return struct.unpack_from(">b", d, p + 1)[0], p + 2
    if dt == INT16:
        return struct.unpack_from(">h", d, p + 1)[0], p + 3
    if dt == INT32:
        return struct.unpack_from(">i", d, p + 1)[0], p + 5
    raise ValueError(f"长度int类型异常 {dt:#x}@{p}")


def _skip_value(d: bytes, p: int, dt: int) -> int:
    """跳过一个字段值, 返回新偏移。"""
    if dt == ZERO:
        return p
    if dt == INT8:
        return p + 1
    if dt == INT16:
        return p + 2
    if dt == INT32:
        return p + 4
    if dt == INT64:
        return p + 8
    if dt == STRING1:
        return p + 1 + d[p]
    if dt == STRING4:
        return p + 4 + struct.unpack_from(">i", d, p)[0]
    if dt == SIMPLE_LIST:
        p += 1  # 元素类型头
        n, p = _read_len_int(d, p)
        return p + n
    if dt == MAP:
        n, p = _read_len_int(d, p)
        for _ in range(n):
            h = d[p]; p += 1
            kd = h & 0x0F
            if kd == STRING1:
                p += 1 + d[p]
            elif kd == STRING4:
                p += 4 + struct.unpack_from(">i", d, p)[0]
            else:
                raise ValueError(f"map key 类型 {kd:#x}")
            vh = d[p]; p += 1
            p = _skip_value(d, p, vh & 0x0F)
        return p
    if dt == LIST:
        n, p = _read_len_int(d, p)
        # 元素同型: 读第一个元素头后逐个跳
        eh = d[p]; p += 1
        edt = eh & 0x0F
        for _ in range(n):
            p = _skip_value(d, p, edt)
        return p
    if dt == STRUCT_BEGIN:
        while True:
            h = d[p]; p += 1
            sdt = h & 0x0F
            if sdt == STRUCT_END:
                return p
            if ((h >> 4) & 0x0F) == 15:
                p += 1  # 扩展tag
            p = _skip_value(d, p, sdt)
    raise ValueError(f"未支持的TAF类型 {dt:#x}")


class Envelope:
    """一个已解析的wupData信封, 支持原位补丁(uid/证书/session)。"""

    def __init__(self, raw: bytes):
        self.raw = bytearray(raw)
        self._locate()

    # ---------- 构造 ----------
    @classmethod
    def load(cls, path: str | Path | None = None) -> "Envelope":
        path = Path(path) if path else ROOT / "evidence/cert_keycap.json"
        ev = json.loads(path.read_text())
        qurl = next(e["data"] for e in ev
                    if e.get("type") == "qurl_done" and e.get("data"))
        return cls(base64.b64decode(qurl))

    @classmethod
    def from_bytes(cls, raw: bytes) -> "Envelope":
        return cls(raw)

    # ---------- 定位 ----------
    def _locate(self) -> None:
        raw = self.raw
        # 跳过 [u32 total][ver][seq u16]
        p = 7
        self.uid_off = self.cert_off = self.cert_len = None
        self.meta_json_span = None      # (起, 止) t0.t2 字符串值
        self.tag4_span = (8, 12)        # iRequestId 固定位置
        buf_span = None
        while p < len(raw):
            h = raw[p]; p += 1
            tag, dt = (h >> 4) & 0x0F, h & 0x0F
            if tag == 15:
                tag = raw[p]; p += 1
            if dt == STRUCT_END:
                continue
            start = p
            if tag == 7 and dt == SIMPLE_LIST:
                p += 1                      # 元素类型头
                n, p = _read_len_int(bytes(raw), p)
                buf_span = (p, p + n)
                p += n
                continue
            p = _skip_value(bytes(raw), p, dt)
        if not buf_span:
            raise ValueError("信封中未找到 tag7 sBuffer")
        # 解析 sBuffer MAP, 找 "_wup_data"
        d = bytes(raw)
        q = buf_span[0]
        assert d[q] & 0x0F == MAP, "sBuffer 不是MAP"
        n, q = _read_len_int(d, q + 1)
        wd_span = None
        for _ in range(n):
            kh = d[q]; q += 1
            assert kh & 0x0F == STRING1
            klen = d[q]; q += 1
            key = d[q:q + klen].decode(); q += klen
            vh = d[q]; q += 1
            assert vh & 0x0F == SIMPLE_LIST
            q += 1                          # 元素类型头
            vn, q = _read_len_int(d, q)
            if key == "_wup_data":
                wd_span = (q, q + vn)
            q += vn
        if not wd_span:
            raise ValueError("sBuffer 中无 _wup_data")
        # 遍历 _wup_data struct: 记录 t3(INT64)/t4(STR4)/t0.t2(meta json)
        q = wd_span[0]
        h = d[q]; q += 1
        assert h & 0x0F == STRUCT_BEGIN
        while True:
            hh = d[q]; q += 1
            tag, dt = (hh >> 4) & 0x0F, hh & 0x0F
            if tag == 15:
                tag = d[q]; q += 1
            if dt == STRUCT_END:
                break
            if tag == 0 and dt == STRUCT_BEGIN:     # t0 头struct
                while True:
                    h2 = d[q]; q += 1
                    tg2, dt2 = (h2 >> 4) & 0x0F, h2 & 0x0F
                    if tag == 15:
                        tag = d[q]; q += 1
                    if dt2 == STRUCT_END:
                        break
                    vs = q
                    q = _skip_value(d, q, dt2)
                    if tg2 == 2 and dt2 in (STRING1, STRING4):
                        off = vs + (4 if dt2 == STRING4 else 1)
                        self.meta_json_span = (off, q)
            elif tag == 3 and dt == INT64:           # uid
                self.uid_off = q
                q += 8
            elif tag == 4 and dt == STRING4:         # b64 cert
                ln = struct.unpack_from(">i", d, q)[0]
                self.cert_off = q + 4
                self.cert_len = ln
                q = self.cert_off + ln
            else:
                q = _skip_value(d, q, dt)
        if self.uid_off is None or self.cert_off is None:
            raise ValueError("未定位到 t3(uid)/t4(cert)")

    # ---------- 补丁 ----------
    @property
    def uid(self) -> int:
        return struct.unpack_from(">Q", self.raw, self.uid_off)[0]

    def patch_uid(self, uid: int) -> "Envelope":
        struct.pack_into(">Q", self.raw, self.uid_off, uid)
        return self

    @property
    def cert_b64(self) -> bytes:
        return bytes(self.raw[self.cert_off:self.cert_off + self.cert_len])

    def patch_cert(self, cert: bytes) -> "Envelope":
        b64 = base64.b64encode(cert)
        if len(b64) != self.cert_len:
            raise ValueError(
                f"证书b64长度不符: {len(b64)} != 模板 {self.cert_len} "
                f"(cert {len(cert)}B)")
        self.raw[self.cert_off:self.cert_off + self.cert_len] = b64
        return self

    def patch_session(self, session: int) -> "Envelope":
        """宽度保持地改写三处session(iRequestId/meta JSON/v0尾4字节)。"""
        # 1) iRequestId
        struct.pack_into(">I", self.raw, self.tag4_span[0], session & 0xFFFFFFFF)
        # 2) meta JSON 内 "session":N (等宽数字替换)
        if self.meta_json_span:
            s, e = self.meta_json_span
            js = bytes(self.raw[s:e]).decode("utf-8", "replace")
            old = _extract_session(js)
            if old is not None:
                new = str(session % (10 ** len(str(old)))).rjust(len(str(old)), "0")
                js2 = js.replace(f'"session":{old}', f'"session":{new}', 1)
                self.raw[s:e] = js2.encode("utf-8")[: e - s]
        # 3) wupudbrequest_v0 尾4字节
        d = self.raw
        i = d.find(b"wupudbrequest_v0")
        if i >= 0:
            j = i + len(b"wupudbrequest_v0") + 2  # 跳过key头+value头
            j += 2                                 # 元素类型头+长度int头
            if j + 4 <= len(d):
                struct.pack_into(">I", d, j, session & 0xFFFFFFFF)
        return self

    def wup_b64(self) -> str:
        return base64.b64encode(bytes(self.raw)).decode()


def _extract_session(js: str) -> int | None:
    import re
    m = re.search(r'"session":(\d+)', js)
    return int(m.group(1)) if m else None


if __name__ == "__main__":
    env = Envelope.load()
    print(f"信封 {len(env.raw)}B")
    print(f"  uid@{env.uid_off} = {env.uid}")
    print(f"  cert b64@{env.cert_off} len={env.cert_len}: {env.cert_b64[:40]!r}...")
    if env.meta_json_span:
        s, e = env.meta_json_span
        print(f"  metaJSON@{s}-{e}: {bytes(env.raw[s:e]).decode()[:160]}")