#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
腾讯天御 TDC 混沌虚拟机 (Chaos VM) —— **纯 Python** 字节码提取 + 反汇编器

为什么需要它
------------
旧 `tools/tdc_vm_disasm.js` 依赖 Node + Babel，且用**线性扫描**解码。实测该字节码数组
里混有数据区/不可达区，线性扫描一进入数据区就失步，产出 `UNKNOWN_7`(opcode 7 其实是空洞)、
`UNKNOWN_5096` 这类把操作数当操作码的垃圾 —— 换句话说旧的"8151 条指令流"是坏的。

本工具:
  1. 纯 Python 从 tdc.js 文本里取出字节码数组 `z`（复刻其内联的 base64 解码 + 游程展开），
     不需要任何 JS 运行时;
  2. 用**可达性驱动的 CFG 反汇编**（从 @0 出发，沿 fall-through / JUMP / JUMP_IF_TRUE /
     PUSH_TRY 的 catch 目标推进），只在真正可执行的地址上解码 —— 数据区自然留空，
     不再失步。

opcode 表来自对 VM 源码 AST 的推导（`G = [ ... ]` 中第 i 个 handler 即 opcode i，
操作数个数 = handler 体内 `z[H++]` 的出现次数），已固化在下方。

用法:
    .venv/bin/python tools/tdc_disasm.py [tdc.js]        # 默认 tcaptcha/core_js/tdc_sample.js
"""
from __future__ import annotations

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

# --------------------------------------------------------------------------- #
# opcode 表 (从 VM 源码 G 数组推导)
#   名字: 语义 (部分沿用旧工具命名);  arity: 固定操作数字节数
#   57 = CREATE_CLOSURE: 动态 (3 + 2*capCount + paramCount)
# --------------------------------------------------------------------------- #
OP_NAMES: Dict[int, str] = {
    0: "PUSH_STR_CHAR", 1: "IS_EXCEPTION", 2: "DUP", 3: "XOR", 4: "EQ_STRICT",
    5: "MOD", 6: "CALL_FUNC", 8: "GET_GLOBAL_PROP", 9: "MAKE_PROP_REF",
    10: "TYPEOF", 11: "DELETE", 12: "LOGICAL_NOT", 13: "EQ_LOOSE", 14: "SET_PROP",
    15: "GREATER_THAN", 17: "SWAP", 18: "IN", 19: "RET_TRUE", 20: "GREATER_EQ",
    22: "PUSH_IMM", 23: "MUL", 25: "LOAD_REG_INDIRECT", 26: "ITER_NEXT",
    28: "BIND_THIS", 29: "PUSH_NULL", 30: "BITWISE_AND", 31: "PUSH_IMM_ARRAY",
    32: "MAKE_REG_PROP_REF", 33: "SUB", 34: "CALL_METHOD", 35: "SHR", 36: "SHL",
    37: "JUMP_IF_TRUE", 38: "FOR_IN_KEYS", 39: "ADD", 41: "NEW_METHOD",
    42: "PUSH_UNDEFINED", 45: "PUSH_FALSE", 46: "STORE_REG", 47: "ALLOC_REG",
    48: "POP_TRY", 49: "THROW", 51: "POP", 52: "DIV", 53: "SET_STACK_LEN",
    54: "REVERSE_TOP2", 55: "PUSH_EMPTY_STR", 56: "PUSH_TRY", 57: "CREATE_CLOSURE",
    59: "GET_PROP", 60: "JUMP", 61: "NEW_FUNC", 62: "BITWISE_OR", 64: "REPLACE_TOP_IMM",
    65: "SHR_UNSIGNED", 66: "PUSH_TRUE", 68: "LOAD_REG_IMM", 69: "CLEAR_EXCEPTION",
}

# 固定操作数字节数 (未列出的 opcode = 0 字节)
ARITY: Dict[int, int] = {
    0: 1, 6: 1, 17: 1, 22: 1, 31: 1, 34: 1, 37: 1, 41: 1, 47: 1, 53: 1,
    56: 2, 60: 1, 61: 1, 64: 1, 68: 1,
}
DYNAMIC_ARITY = {57}          # 需要读操作数才能定长
OP_JUMP = 60
OP_JUMP_IF_TRUE = 37
OP_PUSH_TRY = 56              # operands = (catchTarget, catchVarReg)
OP_CREATE_CLOSURE = 57

_B64_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/="


# --------------------------------------------------------------------------- #
# 1. 纯 Python 取出字节码数组 z (复刻 tdc.js 内联的 I() 解码 + 游程展开)
# --------------------------------------------------------------------------- #
def _b64_decode_custom(text: str) -> str:
    """复刻 tdc.js 内联的 I(A): 标准 base64 码表, 忽略非码表字符 (含换行)。"""
    g = re.sub(r"[=]+$", "", str(text))
    idx = {ch: i for i, ch in enumerate(_B64_ALPHABET)}
    out: List[str] = []
    b = 0
    n = 0
    for ch in g:
        w = idx.get(ch, -1)
        if w == -1:                      # ~w falsy -> 跳过 (换行/空白)
            continue
        b = (64 * b + w) if (n % 4) else w
        old = n
        n += 1
        if old % 4:                       # 每组 4 个字符产出 3 字节
            out.append(chr(255 & (b >> ((-2 * n) & 6))))
    return "".join(out)


def _unescape_js_string(raw: str) -> str:
    """还原 JS 字符串字面量里的转义 (重点: `\\n` -> 真换行)。

    注意: `n` 本身是合法 base64 字符, 若把 `\\n` 当成字面 `n` 会**污染**解码结果,
    必须先把转义还原成真实字符 (非码表字符随后被解码器忽略)。
    """
    out: List[str] = []
    i = 0
    while i < len(raw):
        c = raw[i]
        if c == "\\" and i + 1 < len(raw):
            nxt = raw[i + 1]
            out.append({"n": "\n", "r": "\r", "t": "\t", "\\": "\\",
                        '"': '"', "'": "'", "/": "/"}.get(nxt, nxt))
            i += 2
            continue
        out.append(c)
        i += 1
    return "".join(out)


def extract_bytecode(js_text: str) -> List[int]:
    """从 tdc.js 文本里取出混沌 VM 的字节码数组 z (纯 Python, 无 JS 运行时)。

    结构 (实测): `__TENCENT_CHAOS_VM(0, function(A){...}(["<base64>",[<int>,...]]), ...)`
    这里以 `__TENCENT_CHAOS_VM` 为锚点定位 `}(["...` 实参, 不再用"最长字符串"猜测。
    """
    anchor = js_text.rfind("__TENCENT_CHAOS_VM")
    if anchor < 0:
        raise ValueError("未找到 __TENCENT_CHAOS_VM 锚点")

    # 找锚点之后的 `}([` -> 数组起始的 `[`
    k = js_text.find("}([", anchor)
    if k < 0:
        k = js_text.find("[", js_text.find("function", anchor))
    start = js_text.find("[", k)
    if start < 0 or start + 1 >= len(js_text) or js_text[start + 1] != '"':
        raise ValueError("未能定位字节码实参数组")

    q1 = start + 1
    q2 = js_text.find('"', q1 + 1)          # base64 字面量内不含 `"`
    if q2 < 0:
        raise ValueError("base64 字面量未闭合")
    b64 = _unescape_js_string(js_text[q1 + 1:q2])

    # 取内层整数数组 (括号配对)
    i = q2 + 1
    while i < len(js_text) and js_text[i] in " \t\r\n,":
        i += 1
    if i >= len(js_text) or js_text[i] != "[":
        raise ValueError("base64 之后不是整数数组")
    depth = 0
    j = i
    while j < len(js_text):
        if js_text[j] == "[":
            depth += 1
        elif js_text[j] == "]":
            depth -= 1
            if depth == 0:
                break
        j += 1
    # 该整数数组会混用科学计数法 (如 `1e9` / `16e3`), 需按数值解析
    ints = [int(float(x)) for x in js_text[i + 1:j].split(",") if x.strip() != ""]

    # 3) 复刻 VM 的游程展开
    chars = _b64_decode_custom(b64)                 # G
    w: List[int] = ints

    q: List[int] = []
    h = 0
    m = w.pop(0) if w else -1
    g = w.pop(0) if w else 0

    def flush() -> None:
        nonlocal h, m, g
        while h == m:
            q.append(g)
            h += 1
            m = w.pop(0) if w else -1
            g = w.pop(0) if w else 0

    for ch in chars:
        flush()
        q.append(ord(ch))
        h += 1
    flush()
    return q


# --------------------------------------------------------------------------- #
# 2. 可达性驱动的 CFG 反汇编 (只在可执行地址解码)
# --------------------------------------------------------------------------- #
def _decode_one(z: List[int], pc: int) -> Tuple[Optional[int], List[int], Optional[List[str]], int]:
    """解码 pc 处一条指令 -> (opcode, operands, note, next_fallthrough_pc)。

    返回值 (None,...) 表示该地址无法构成合法指令。
    """
    if pc >= len(z):
        return None, [], None, pc
    op = z[pc]
    if op > 69 or op not in OP_NAMES:      # 空洞 opcode / 越界 -> 非法
        return None, [], None, pc + 1
    pos = pc + 1
    operands: List[int] = []
    note: Optional[List[str]] = None

    if op in DYNAMIC_ARITY:
        # 57: target, capCount, paramCount, [reg, reg]*capCount, [reg]*paramCount
        if pos + 3 > len(z):
            return None, [], None, pos
        target, cap, param = z[pos], z[pos + 1], z[pos + 2]
        pos += 3
        need = 2 * cap + param
        if pos + need > len(z):
            return None, [], None, pos
        operands = z[pos:pos + need]
        pos += need
        caps = [f"r{operands[2*i]}=r{operands[2*i+1]}" for i in range(cap)]
        params = [f"p{i}->r{operands[2*cap+i]}" for i in range(param)]
        note = [f"target=@{target}", f"caps=[{','.join(caps)}]", f"params=[{','.join(params)}]"]
    else:
        n = ARITY.get(op, 0)
        if pos + n > len(z):
            return None, [], None, pos
        operands = z[pos:pos + n]
        pos += n
        if op in (OP_JUMP, OP_JUMP_IF_TRUE) and operands:
            note = [f"target=@{operands[0]}"]
        elif op == OP_PUSH_TRY and len(operands) == 2:
            note = [f"catchTarget=@{operands[0]}", f"catchVarReg=r{operands[1]}"]
    return op, operands, note, pos


def disassemble_cfg(z: List[int], entry: int = 0) -> List[dict]:
    """从 entry 出发按可达性反汇编; 返回按 pc 排序的指令列表。"""
    visited: Dict[int, dict] = {}
    work = [entry]
    while work:
        pc = work.pop()
        if pc in visited or pc >= len(z) or pc < 0:
            continue
        op, operands, note, nxt = _decode_one(z, pc)
        if op is None:
            # 非法地址 (落在数据区): 记录但不继续
            visited[pc] = {"pc": pc, "op": None, "name": "DATA/INVALID", "operands": [], "note": None}
            continue
        visited[pc] = {"pc": pc, "op": op, "name": OP_NAMES[op], "operands": operands, "note": note}
        if op == OP_JUMP:
            if operands:
                work.append(operands[0])
            continue                        # 无条件跳转: 无 fall-through
        if op == OP_JUMP_IF_TRUE:
            if operands:
                work.append(operands[0])
            work.append(nxt)                # 条件跳转: 有 fall-through
            continue
        if op == OP_PUSH_TRY and operands:
            work.append(operands[0])        # catch 目标可达
        if op == OP_CREATE_CLOSURE:
            # 57 的操作数 [0] 是闭包体入口: 该段代码在闭包被调用时才执行, 但它是**代码**,
            # 必须作为遍历根加入, 否则大半个程序都算不到。
            work.append(z[pc + 1])
        work.append(nxt)
    return [visited[k] for k in sorted(visited)]


# --------------------------------------------------------------------------- #
def main() -> int:
    path = sys.argv[1] if len(sys.argv) > 1 else "tcaptcha/core_js/tdc_sample.js"
    js = open(path, encoding="utf-8", errors="replace").read()
    z = extract_bytecode(js)
    insts = disassemble_cfg(z)
    invalid = [i for i in insts if i["op"] is None]
    print(f"bytecode length : {len(z)}")
    print(f"reachable instrs: {len(insts)}   (invalid/data slots: {len(invalid)})")
    op_used = sorted({i["op"] for i in insts if i["op"] is not None})
    holes = [o for o in op_used if o not in OP_NAMES]
    print(f"opcodes used    : {op_used}")
    print(f"illegal opcodes : {holes}")
    out = path.rsplit(".", 1)[0] + ".disasm.txt"
    with open(out, "w", encoding="utf-8") as f:
        for i in insts:
            if i["op"] is None:
                f.write(f"@{i['pc']:05d}: {i['name']}\n")
            else:
                extra = (" " + " ".join(i["note"])) if i["note"] else ""
                argc = (" " + str(i["operands"])) if (i["operands"] and not i["note"]) else ""
                f.write(f"@{i['pc']:05d}: {i['name']:<18}{extra}{argc}\n")
    print(f"written         : {out}")
    for i in insts[:12]:
        print("  ", i)
    return 0


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