# -*- coding: utf-8 -*-
"""
腾讯天御 TDC 混沌 VM (Chaos VM) —— XTEA 密钥的**纯 Python** 静态还原器。

为什么需要它
------------
此前要拿到某一版 tdc.js 的 XTEA 密钥, 只能把 tdc.js 丢进 Node 插桩跑一遍
(ADDU opcode 吐密钥字)。密钥**不是字面量**, 而是 VM 运行时由字节码算出来的,
所以必须"读指令流"而不是"搜字符串"。

本模块完全离线(零 Node、零 JS 运行时), 基于**正确反汇编**(tools/tdc_disasm.py)做
**模式 + 数据流**定位, 复刻密钥构造算法。全部结论都有运行时神谕（真浏览器 TDC.getData）
逐条对照过, 见 har/1ms_dianxuan/PROGRESS.md。

算法 (已证: 纯 Python 结果 == Node extract_xtea_key 逐字一致)
------------------------------------------------------------
1. **4 个密钥字**: 每个字由一个专门的闭包算出:
       word_i = Σ_{k=0..3} ( (ord(src_i[k]) + off_i) & 0xFF ) << (8*k)
   源串 `src_i` 是紧跟在该闭包 `CREATE_CLOSURE` 之后的 4 个 `PUSH_STR_CHAR`
   字面量; 偏移 `off_i` 是打包点的立即数; 存入下标 `i` 由闭包尾部的
   `PUSH_IMM_ARRAY / PUSH_IMM [i] / ... / SET_PROP` 给出。

2. **哪些闭包是真密钥**: 密钥数组会被 **XTEA 闭包捕获**。
   XTEA 轮里出现 `delta = 0x9E3779B9`; 含 delta 的那个闭包体所属的
   `CREATE_CLOSURE` 即 XTEA 闭包, 其 `caps=[r7=r<K>]` 的 `<K>` 就是密钥数组所在的
   父帧寄存器。**只有 caps 右侧等于该寄存器的密钥字闭包**才是真密钥
   (同一份脚本往往还并行算了第二个"诱饵"数组, 用不同寄存器, 必须排除)。

3. **轮加盐**: XTEA 轮内 `key[(sum>>>11)&3]` 查表之后会**条件性**加上一个立即数
   (salt): `GET_PROP / PUSH_IMM <salt> / ADD / ADD / XOR`。加给哪个下标由守卫条件
   `idx == ((sum>>>11)&3)` / `idx == (sum & 3)` 里的 `idx` 常量决定。
   有效密钥 = base 字按盐的下标位置相加。

所有定位都是**模式 + 立即数**, 不写死任何 pc (每份构建都会重排)。
"""

from __future__ import annotations

import os
import re
import sys
from typing import Dict, List, Optional, Tuple

_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _ROOT not in sys.path:            # 允许在任意 cwd 下 import tools.*
    sys.path.insert(0, _ROOT)

from tcaptcha.telemetry.disasm import extract_bytecode, disassemble_cfg  # noqa: E402

from tcaptcha.telemetry.opcodes import derive_opcode_table  # noqa: E402

M32 = 0xFFFFFFFF
XTEA_DELTA = 0x9E3779B9              # 2654435769, XTEA 黄金比常数
KEY_LEN = 4

# 走"下标表达式"时要跳过的算子 (它们构成 `(sum & 3)` / `(sum>>>11)&3`)
_SKIP_OPS = {
    "BITWISE_AND", "SHR_UNSIGNED", "LOAD_REG_IMM", "MAKE_REG_PROP_REF",
    "GET_PROP", "SWAP", "STORE_REG", "ADD", "DUP", "LOAD_REG_INDIRECT",
    "PUSH_IMM_ARRAY", "POP",
}
# 这些算子紧跟在"掩码 PUSH_IMM 3"之前, 说明那个 3 是掩码而不是下标
_MASK_NEIGHBOURS = {"BITWISE_AND", "LOAD_REG_IMM", "SHR_UNSIGNED"}


class _Dis:
    """对反汇编结果做少量索引封装, 便于按 pc 前后游走。"""

    def __init__(self, js_text: str) -> None:
        self.z = extract_bytecode(js_text)
        # opcode 编号每份构建都重排 -> 必须先按构建把语义表推导出来
        names, arity, dynamic, diag = derive_opcode_table(js_text)
        self.names, self.arity, self.dynamic, self.diag = names, arity, dynamic, diag
        self.insts = disassemble_cfg(self.z, names=names, arity=arity, dynamic=dynamic)
        self.by_pc: Dict[int, dict] = {i["pc"]: i for i in self.insts}
        self.pcs: List[int] = sorted(self.by_pc)
        self._idx: Dict[int, int] = {pc: n for n, pc in enumerate(self.pcs)}

    def op(self, pc: Optional[int]) -> Optional[str]:
        i = self.by_pc.get(pc) if pc is not None else None
        return i["name"] if i else None

    def ops(self, pc: Optional[int]) -> List[int]:
        i = self.by_pc.get(pc) if pc is not None else None
        return list(i["operands"]) if i else []

    def note(self, pc: int) -> str:
        return " ".join(self.by_pc[pc].get("note") or [])

    def nxt(self, pc: Optional[int]) -> Optional[int]:
        n = self._idx.get(pc) if pc is not None else None
        return self.pcs[n + 1] if n is not None and n + 1 < len(self.pcs) else None

    def prev(self, pc: Optional[int]) -> Optional[int]:
        n = self._idx.get(pc) if pc is not None else None
        return self.pcs[n - 1] if n is not None and n > 0 else None

    def in_range(self, lo: int, hi: int) -> List[int]:
        import bisect
        a = bisect.bisect_left(self.pcs, lo)
        b = bisect.bisect_right(self.pcs, hi)
        return self.pcs[a:b]


# --------------------------------------------------------------------------- #
# 1. 闭包清单 / 密钥字打包点
# --------------------------------------------------------------------------- #
def _closures(d: _Dis) -> List[dict]:
    """所有 CREATE_CLOSURE -> {pc, target, caps:{dst:src}, str_after}。"""
    out: List[dict] = []
    for i in d.insts:
        if i["name"] != "CREATE_CLOSURE":
            continue
        pc = i["pc"]
        note = " ".join(i["note"] or [])
        m = re.search(r"target=@(\d+)", note)
        if not m:
            continue
        caps: Dict[int, int] = {}
        caps_part = re.search(r"caps=\[([^\]]*)\]", note)
        if caps_part:
            for dst, src in re.findall(r"r(\d+)=r(\d+)", caps_part.group(1)):
                caps[int(dst)] = int(src)
        out.append({"pc": pc, "target": int(m.group(1)), "caps": caps,
                    "str_after": _string_after(d, pc)})
    return out


def _string_after(d: _Dis, pc: int) -> Optional[str]:
    """CREATE_CLOSURE 紧随其后的 `PUSH_EMPTY_STR + PUSH_STR_CHAR*` 字面量。"""
    p = d.nxt(pc)
    if p is None or d.op(p) != "PUSH_EMPTY_STR":
        return None
    chars: List[str] = []
    q = d.nxt(p)
    while q is not None and d.op(q) == "PUSH_STR_CHAR":
        chars.append(chr(d.ops(q)[0]))
        q = d.nxt(q)
    return "".join(chars)


def _pack_sites(d: _Dis, lo: int, hi: int) -> List[dict]:
    """
    在 [lo,hi] 里找"密钥字打包点":
        ... CALL_METHOD[1] / PUSH_IMM c / SUB            (off = -c)
        或  ... CALL_METHOD[1] / PUSH_IMM 0 / PUSH_IMM c / SUB / SUB   (off = +c)
        / PUSH_IMM 8 / LOAD_REG_IMM / MUL / SHL / BITWISE_OR
    返回 [{or_pc, off}]。
    """
    res: List[dict] = []
    for pc in d.in_range(lo, hi):
        if d.op(pc) != "PUSH_IMM" or d.ops(pc) != [8]:
            continue
        a, b, c, e = d.nxt(pc), None, None, None
        b = d.nxt(a)
        c = d.nxt(b)
        e = d.nxt(c)
        if not (d.op(a) == "LOAD_REG_IMM" and d.op(b) == "MUL"
                and d.op(c) == "SHL" and d.op(e) == "BITWISE_OR"):
            continue
        off = _decode_offset(d, pc)
        if off is not None:
            res.append({"or_pc": e, "imm8_pc": pc, "off": off})
    return res


def _decode_offset(d: _Dis, imm8_pc: int) -> Optional[int]:
    """解出紧跟 CALL_METHOD 之后的偏移立即数 (含两种编码)。"""
    p1 = d.prev(imm8_pc)                 # 最后一个 SUB
    if d.op(p1) != "SUB":
        return None
    p2 = d.prev(p1)
    if d.op(p2) == "PUSH_IMM":           # X - c  => off = -c
        return -d.ops(p2)[0]
    if d.op(p2) == "SUB":                # X - (0 - c) => off = +c
        p3 = d.prev(p2)
        p4 = d.prev(p3)
        if d.op(p3) == "PUSH_IMM" and d.op(p4) == "PUSH_IMM" and d.ops(p4) == [0]:
            return d.ops(p3)[0]
    return None


def _store_index(d: _Dis, lo: int, hi: int) -> Optional[int]:
    """闭包尾部的 `PUSH_IMM_ARRAY / PUSH_IMM [i] / MAKE_REG_PROP_REF / LOAD_REG_IMM / SET_PROP`。"""
    hits: List[int] = []
    for pc in d.in_range(lo, hi):
        if d.op(pc) != "PUSH_IMM_ARRAY":
            continue
        a, b, c, e = d.nxt(pc), None, None, None
        b = d.nxt(a)
        c = d.nxt(b)
        e = d.nxt(c)
        if (d.op(a) == "PUSH_IMM" and d.op(b) == "MAKE_REG_PROP_REF"
                and d.op(c) == "LOAD_REG_IMM" and d.op(e) == "SET_PROP"):
            hits.append(d.ops(a)[0])
    return hits[-1] if hits else None


# --------------------------------------------------------------------------- #
# 2. XTEA 闭包定位 + 轮加盐
# --------------------------------------------------------------------------- #
def _find_xtea_closure(d: _Dis, closures: List[dict]) -> Optional[dict]:
    """含 XTEA 常数(delta)的那个闭包体所属的 CREATE_CLOSURE。"""
    delta_pcs = [pc for pc in d.pcs
                 if d.op(pc) == "PUSH_IMM" and d.ops(pc)[0] in (XTEA_DELTA, XTEA_DELTA * 32)]
    # 取"最内层"(区间最小)的闭包体 —— 外层 init 函数的体也会横跨 delta,
    # 但它不捕获密钥数组。
    best: Optional[dict] = None
    best_end: Optional[int] = None
    for cl in closures:
        end = _body_end(d, cl["pc"])
        if end is None or end < cl["target"]:
            continue
        if any(cl["target"] <= pc <= end for pc in delta_pcs):
            if best is None or best_end is None or (end - cl["target"]) < (best_end - best["target"]):
                best, best_end = cl, end
    return best


def _largest_reg(regs: Dict[int, int]) -> Optional[int]:
    return max(regs.values()) if regs else None


def _caps_single(caps: Dict[int, int]) -> Optional[Tuple[int, int]]:
    """字/XTEA 闭包只捕获一个数组: 返回 (dst, src)。"""
    if len(caps) != 1:
        return None
    (dst, src), = caps.items()
    return dst, src


def _defining_closure(d: _Dis, closures: List[dict], pc: int) -> Optional[dict]:
    """含给定 pc 的**最内层**闭包(即创建 pc 处那条 CREATE_CLOSURE 的那个函数)。"""
    best: Optional[dict] = None
    best_span: Optional[int] = None
    for c in closures:
        e = _body_end(d, c["pc"])
        if e is None or e < c["target"]:
            continue
        if c["target"] <= pc <= e:
            span = e - c["target"]
            if best_span is None or span < best_span:
                best, best_span = c, span
    return best


def _body_end(d: _Dis, cl_pc: int) -> Optional[int]:
    """
    闭包体结束 pc (该体的 RET_TRUE)。约定: 闭包被创建后要么紧跟其体尾的 RET_TRUE,
    要么中间夹一条 `PUSH_IMM_ARRAY`(闭包被塞进数组时)。
    """
    e = d.prev(cl_pc)
    if e is not None and d.op(e) == "PUSH_IMM_ARRAY":
        e = d.prev(e)
    return e if (e is not None and d.op(e) == "RET_TRUE") else None


def _guard_index(d: _Dis, eq_pc: int) -> Optional[int]:
    """
    守卫 `idx == ((sum & 3) | ((sum>>>11)&3))` 里的常量 idx。
    从 EQ_LOOSE 反向游走, 跳过下标表达式自身的算子, 第一个落在 0..3 的立即数即 idx。
    掩码 3 靠"它的后继是 BITWISE_AND / LOAD_REG_IMM / SHR_UNSIGNED"识别并跳过。
    """
    pc = d.prev(eq_pc)
    last: Optional[str] = None
    for _ in range(24):
        if pc is None:
            return None
        nm, ops = d.op(pc), d.ops(pc)
        if nm in ("PUSH_IMM", "REPLACE_TOP_IMM") and ops:
            v = ops[0]
            if 0 <= v <= 3:
                if not (nm == "PUSH_IMM" and v == 3 and last in _MASK_NEIGHBOURS):
                    return v
        elif nm not in _SKIP_OPS:
            # 出现意料之外的算子: 继续向前找, 但记录上下文
            pass
        last = nm
        pc = d.prev(pc)
    return None


def _salts(d: _Dis, lo: int, hi: int) -> List[Tuple[int, int]]:
    """
    XTEA 轮内的条件加盐: `GET_PROP / PUSH_IMM <salt> / ADD / ADD / XOR`。
    返回 [(idx, salt)] (按 (idx,salt) 去重)。
    """
    salt_pcs: List[int] = []
    for pc in d.in_range(lo, hi):
        if d.op(pc) != "PUSH_IMM":
            continue
        if d.op(d.prev(pc)) != "GET_PROP" or d.op(d.nxt(pc)) != "ADD":
            continue
        salt_pcs.append(pc)
    if not salt_pcs:
        return []

    # 收集所有 (JUMP_IF_TRUE, EQ_LOOSE) 守卫; block(T) 到 T 之后第一条无条件 JUMP
    guards: List[Tuple[int, int, Optional[int]]] = []   # (target, jump_pc, idx)
    for pc in d.in_range(lo, hi):
        if d.op(pc) != "JUMP_IF_TRUE":
            continue
        tgt_ops = d.ops(pc)
        if not tgt_ops:
            continue
        tgt = tgt_ops[0]
        eq = d.prev(pc)
        if d.op(eq) != "EQ_LOOSE":
            continue
        idx = _guard_index(d, eq)
        if idx is not None:
            guards.append((tgt, pc, idx))

    def block_end(tgt: int) -> int:
        p = tgt
        while p is not None and p <= hi:
            if d.op(p) == "JUMP":
                return p
            p = d.nxt(p)
        return hi

    out: List[Tuple[int, int]] = []
    for spc in salt_pcs:
        salt = d.ops(spc)[0]
        best: Optional[Tuple[int, int]] = None      # (target, idx)
        for tgt, _j, idx in guards:
            if tgt <= spc < block_end(tgt):
                if best is None or tgt > best[0]:
                    best = (tgt, idx)
        if best is not None:
            out.append((best[1], salt))
    # 去重保序
    seen = set()
    uniq: List[Tuple[int, int]] = []
    for x in out:
        if x not in seen:
            seen.add(x)
            uniq.append(x)
    return uniq


# --------------------------------------------------------------------------- #
# 3. 主入口
# --------------------------------------------------------------------------- #
def word_from_chars(src: str, off: int) -> int:
    """word = Σ ((ord(c) + off) & 0xFF) << (8*k)  (小端打包)"""
    w = 0
    for k, ch in enumerate(src):
        w |= ((ord(ch) + off) & 0xFF) << (8 * k)
    return w & M32


def extract_xtea_key_pure(js_text: str, verbose: bool = False):
    """
    纯 Python 从 tdc.js 里还原本次构建的 XTEA 有效密钥 (4 个 32 位字)。

    返回 dict:
      base   : 4 个 base 字 (未折盐)
      salts  : [(idx, salt), ...]
      key    : 有效密钥
      words  : [(idx, word, src, off), ...]  (便于诊断)
    """
    d = _Dis(js_text)
    closures = _closures(d)
    xtea = _find_xtea_closure(d, closures)
    if xtea is None:
        raise ValueError("未找到 XTEA 闭包 (字节码里没有 delta 常数?)")
    xs = _caps_single(xtea["caps"])
    if xs is None:
        raise ValueError("XTEA 闭包没有唯一捕获的数组, 无法确定密钥数组")
    xtea_src = xs[1]                     # XTEA 所在帧里, 密钥数组所在的寄存器

    # 密钥数组是"被 XTEA 闭包捕获"的那一个。字闭包与 XTEA 闭包往往不在同一帧:
    #   父帧 G:  r<R> = 密钥数组
    #     ├─ 字闭包 W: caps=[r7 = 本帧 r<Wsrc>]        (Wsrc ∈ {R, 诱饵寄存器})
    #     └─ XTEA 闭包 : caps=[r7 = 本帧 r<xtea_src>]
    # 所以要把 W 的捕获穿一层: B = 定义 W 的那个闭包(B 帧 -> G 帧), 检查
    #   B.caps[Wsrc] == xtea_src  -> 该字闭包写的就是真密钥数组。
    words: Dict[int, Tuple[int, str, int]] = {}
    picked: List[str] = []
    for cl in closures:
        if cl["str_after"] is None:
            continue
        cs = _caps_single(cl["caps"])
        if cs is None:
            continue
        w_src = cs[1]                    # 该字闭包捕获的、定义帧里的寄存器
        definer = _defining_closure(d, closures, cl["pc"])
        if definer is not None and definer["caps"]:
            real = definer["caps"].get(w_src) == xtea_src
        else:                            # 同帧: 直接比
            real = (w_src == xtea_src)
        if not real:
            continue
        end = _body_end(d, cl["pc"])
        if end is None or end < cl["target"]:
            continue
        packs = _pack_sites(d, cl["target"], end)
        if not packs:
            continue
        idx = _store_index(d, cl["target"], end)
        if idx is None or not (0 <= idx < KEY_LEN):
            continue
        src, off = cl["str_after"], packs[-1]["off"]
        words[idx] = (word_from_chars(src, off), src, off)
        picked.append("%s(idx=%d,off=%+d)" % (src, idx, off))

    if len(words) != KEY_LEN:
        raise ValueError(
            "密钥字不完整: 只定位到 %s (期望 0..3); XTEA 密钥帧寄存器 r%s"
            % (sorted(words), xtea_src))

    base = [words[i][0] for i in range(KEY_LEN)]
    salts = _salts(d, xtea["target"], d.prev(xtea["pc"]) or xtea["target"])

    key = list(base)
    for idx, salt in salts:
        key[idx] = (key[idx] + salt) & M32

    report = {
        "base": base,
        "salts": salts,
        "key": key,
        "words": [(i, words[i][0], words[i][1], words[i][2]) for i in range(KEY_LEN)],
        "xtea_closure_pc": xtea["pc"],
        "key_reg": xtea_src,
    }
    if verbose:
        print("[tdc_key_pure] XTEA closure @%d  key array = r%s" % (xtea["pc"], xtea_src))
        for i, w, s, o in report["words"]:
            print("  word[%d] <- %-6r off=%+d  = %d" % (i, s, o, w))
        for idx, salt in salts:
            print("  salt +%d -> key[%d]" % (salt, idx))
        print("  base = %s" % base)
        print("  key  = %s" % key)
    return report


def extract_xtea_key(js_text: str) -> List[int]:
    """便捷入口: 只取有效密钥 (与 crypto.extract_xtea_key 同口径)。"""
    return list(extract_xtea_key_pure(js_text)["key"])