#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
腾讯天御 Chaos VM —— **纯 Python 具体执行器** (目标: 静态解出 XTEA 密钥, 去掉 Node 依赖)

思路
----
密钥字是 4 个 "ASCII 打包的 4 字符串", 但既不是立即数、也不是任何字面字符串/其子串,
说明它是**运行时算出来的**。所以这里把 VM 的 70 个 opcode 语义 (由 handlers 源码逐一确认)
在 Python 里重新实现, 直接跑字节码, 把寄存器里出现的 "4 个可打印字符拼成的 32bit 数组" 抓出来。

VM 模型 (与源码一致):
  * `E` 同时是**寄存器区** (E[0..n-1]) 与**求值栈** (E[n..]);
  * 寄存器单元是长度为 1 的数组: 值 = E[i][0];
  * `z` = 字节码, `H` = PC, `c` = 宿主全局对象 (window)。

用法:
    .venv/bin/python tools/tdc_vm_run.py tcaptcha/core_js/tdc_sample.js
    KEY=<已知密钥> .venv/bin/python tools/tdc_vm_run.py tdc.js     # 命中即打印所在寄存器
"""
from __future__ import annotations

import json
import math
import os
import re
import sys
import time
from typing import Any, Dict, List, Optional

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from tdc_disasm import extract_bytecode  # noqa: E402


# --------------------------------------------------------------------------- #
# JS 风格容器: 既支持整数下标, 也支持字符串属性
# --------------------------------------------------------------------------- #
class JArr(list):
    def __init__(self, *a):
        super().__init__(*a)
        self._ext: Dict[str, Any] = {}

    def __getitem__(self, k):
        # JS 数组/对象: 缺失的属性是 undefined (不是 null!), 这点必须严格
        if isinstance(k, int) or (isinstance(k, str) and k.lstrip("-").isdigit()):
            i = int(k)
            try:
                return list.__getitem__(self, i)
            except IndexError:
                return JS_UNDEF
        if k in self._ext:
            return self._ext[k]
        return JS_UNDEF

    def __setitem__(self, k, v):
        if isinstance(k, int) or (isinstance(k, str) and k.lstrip("-").isdigit()):
            i = int(k)
            if 0 <= i < len(self):
                list.__setitem__(self, i, v)
            else:
                self._ext[str(i)] = v
        else:
            self._ext[k] = v


JS_UNDEF = object()

VM_VERSION = 0          # 对应源码里的 __TENCENT_CHAOS_VM.v


class Stub:
    """未知宿主对象: 可调用、可属性访问、可下标; **会真正记住被赋的属性**
    (tdc.js 把模块工厂挂到对象上, 若丢弃就永远找不到模块)。"""

    def __init__(self, name="stub"):
        self._name = name
        self._props: Dict[Any, Any] = {}

    def __call__(self, *a, **k):
        return Stub(self._name + "()")

    def __getattr__(self, k):
        if k.startswith("_"):
            raise AttributeError(k)
        return self[k]

    def __getitem__(self, k):
        if k in self._props:
            return self._props[k]
        v = Stub(f"{self._name}.{k}")
        self._props[k] = v
        return v

    def __setitem__(self, k, v):
        self._props[k] = v

    def __repr__(self):
        return f"<{self._name} props={list(self._props)[:5]}>"


class HostFn:
    """宿主内置函数/构造器 (Object/Array/String/...)。"""

    def __init__(self, name, fn):
        self._name = name
        self._fn = fn

    def __call__(self, *a):
        return self._fn(*a)

    def __getitem__(self, k):
        return Stub(f"{self._name}[{k}]")

    def __repr__(self):
        return f"<host {self._name}>"


def _mk_array(*a):
    if len(a) == 1 and isinstance(a[0], int):
        return JArr([JS_UNDEF] * a[0])
    return JArr(a)


def _mk_object(*a):
    return JArr()


def _to_host(**kw):
    o = JArr()
    o._ext.update(kw)
    return o


def make_host() -> Any:
    """最小宿主 (c = window)。未知项退化为 Stub, 保证不因缺属性而崩。"""
    win = _to_host()

    win._ext.update({
        "Math": Stub("Math"),
        "JSON": type("J", (), {"parse": staticmethod(json.loads), "stringify": staticmethod(json.dumps)})(),
        "Object": HostFn("Object", _mk_object),
        "Array": HostFn("Array", _mk_array),
        "String": HostFn("String", lambda *a: "" if not a else str(a[0])),
        "Number": HostFn("Number", lambda *a: 0 if not a else float(a[0])),
        "Boolean": HostFn("Boolean", lambda *a: bool(a[0]) if a else False),
        "Date": HostFn("Date", lambda *a: int(time.time() * 1000)),
        "RegExp": Stub("RegExp"),
        "Error": Stub("Error"),
        "TypeError": Stub("TypeError"),
        "Uint8Array": HostFn("Uint8Array", lambda *a: JArr([0] * (a[0] if a and isinstance(a[0], int) else 0))),
        "Uint8ClampedArray": HostFn("Uint8ClampedArray", lambda *a: JArr([0] * (a[0] if a and isinstance(a[0], int) else 0))),
        "parseInt": HostFn("parseInt", lambda *a: int(float(a[0])) if a else 0),
        "parseFloat": HostFn("parseFloat", lambda *a: float(a[0]) if a else 0.0),
        "encodeURIComponent": HostFn("encodeURIComponent", lambda *a: a[0] if a else ""),
        "decodeURIComponent": HostFn("decodeURIComponent", lambda *a: a[0] if a else ""),
        "atob": HostFn("atob", lambda *a: a[0] if a else ""),
        "btoa": HostFn("btoa", lambda *a: a[0] if a else ""),
        "undefined": JS_UNDEF,
        "NaN": float("nan"),
    })
    return win


def getprop(base, key):
    """JS 的 base[key] (非严格语义: 对原始值取属性返回 undefined)"""
    if isinstance(base, (Stub, JArr)):
        return base[key]
    if isinstance(base, dict):
        return base.get(str(key), JS_UNDEF)
    if isinstance(base, list):
        return base[int(key)] if isinstance(key, (int, str)) and str(key).lstrip("-").isdigit() and 0 <= int(key) < len(base) else JS_UNDEF
    if isinstance(base, str):
        return base[int(key)] if isinstance(key, (int, str)) and str(key).lstrip("-").isdigit() and 0 <= int(key) < len(base) else JS_UNDEF
    return JS_UNDEF          # number/bool/null/undefined: 属性访问为 undefined


def setprop(base, key, val):
    """JS 的 base[key] = val (非严格语义: 对原始值赋值静默忽略)"""
    if isinstance(base, Stub):
        return
    if isinstance(base, JArr):
        base[key] = val
        return
    if isinstance(base, dict):
        base[str(key)] = val
        return
    if isinstance(base, list):
        i = int(key)
        if 0 <= i < len(base):
            base[i] = val
        else:
            base.append(val)
        return
    return                    # 原始值: 静默忽略


def as_int(v):
    if isinstance(v, bool):
        return 1 if v else 0
    if v is JS_UNDEF or v is None:
        return 0
    if isinstance(v, (int, float)):
        return int(v) & 0xFFFFFFFF
    if isinstance(v, str):
        return len(v) and (ord(v[0]) & 0xFFFFFFFF) or 0
    return 0


def truthy(v):
    if v is JS_UNDEF or v is None or v is False:
        return False
    if isinstance(v, (int, float)):
        return v != 0
    if isinstance(v, str):
        return len(v) > 0
    return True


# --------------------------------------------------------------------------- #
class VM:
    def __init__(self, z: List[int], c, max_steps: int = 20_000_000, fresh: bool = True):
        self.z = z
        self.c = c
        self.fresh = fresh       # True=最外层调用(E 未传入), False=闭包调用
        # 与 VM 源码默认一致: E = [[this],[{}]]; 顶层是普通调用 -> this = 全局对象 (c)
        self.E: List[Any] = [JArr([c]), JArr([JArr()])]
        self.B: List[Any] = []          # try 栈
        self.w = None                   # 异常标志
        self.H = 0
        self.steps = 0
        self.max_steps = max_steps
        self.done = False
        self.ret = None
        self.key_hits: List[Any] = []
        self.trace = bool(os.environ.get("TRACE"))
        self.trace_full = bool(os.environ.get("TRACE_FULL"))
        self.pcs: List[int] = []         # 全量 PC 轨迹 (TRACE_FULL 时记录)
        self.events: List[Any] = []      # 记录调用/闭包/返回等关键事件

    # ---- 工具
    def pu(self):
        return self.E.pop()

    def look(self):
        self.record_cell(self.E[-1])

    def record_cell(self, v):
        if isinstance(v, list) and len(v) == 4 and all(isinstance(x, int) for x in v):
            if all(32 <= (x & 0xFF) < 127 and 32 <= ((x >> 24) & 0xFF) < 127 for x in v):
                packed = b"".join(x.to_bytes(4, "big") for x in v)
                if all(32 <= b < 127 for b in packed):
                    self.key_hits.append(list(v))

    # ---- 主循环
    def run(self):
        z = self.z
        while not self.done and self.H < len(z):
            self.steps += 1
            if self.steps > self.max_steps:
                raise RuntimeError("step limit exceeded")
            op = z[self.H]; self.H += 1
            if self.trace_full:
                self.pcs.append(self.H - 1)
            if self.trace and op in (1, 6, 19, 34, 41, 49, 57, 61):
                self.events.append((self.H - 1, op))
            h = self.OP.get(op)
            if h is None:
                raise RuntimeError(f"unknown opcode {op} at pc {self.H-1}")
            r = h()
            if os.environ.get("DBGCALL") and self.E and self.E[-1] is None:
                print(f"[push None @pc={self.H-1}] op={op} steps={self.steps}", file=sys.stderr)
            if os.environ.get("DBGCALL") and self.E and isinstance(self.E[-1], list) and len(self.E[-1]) == 2 and self.E[-1][0] is None:
                print(f"[产生 [None,{self.E[-1][1]!r}] @pc={self.H-1}] op={op} steps={self.steps}", file=sys.stderr)
            if r is True:
                self.done = True
        return self._finish()

    def _finish(self):
        """复刻 VM 的返回语义 (源码):
        `return A ? (E.pop(), E.slice(3 + VM.v)) : E.pop()`
        A = !E (E 未传入 => 最外层调用); 闭包调用时 E 已传入 => 返回**栈顶** E.pop()。
        """
        E = self.E
        if not E:
            return JS_UNDEF
        if self.fresh:
            E.pop()
            start = 3 + VM_VERSION
            return JArr(E[start:]) if len(E) > start else JArr()
        return E.pop()

    # ------------------------------------------------------------------ ops
    def _run(self):
        z, E = self.z, self.E
        OP = {}

        def op(n):
            def deco(f):
                OP[n] = f
                return f
            return deco

        @op(0)      # PUSH_STR_CHAR
        def _0():
            E[-1] = (E[-1] if isinstance(E[-1], str) else "") + chr(z[self.H]); self.H += 1

        @op(1)      # IS_EXCEPTION -> 结束
        def _1():
            if truthy(self.w):
                self.ret = self.w
                return True
            return None

        @op(2)
        def _2():
            E.append(E[-1])

        def binop(f):
            def g():
                b = E.pop(); E[-2] = f(E[-2], b)
            return g

        OP[3] = binop(lambda a, b: as_int(a) ^ as_int(b))
        OP[4] = binop(lambda a, b: a is b or a == b)
        OP[5] = binop(lambda a, b: (as_int(a) % as_int(b)) if as_int(b) else float("nan"))
        OP[13] = binop(lambda a, b: a == b)
        OP[15] = binop(lambda a, b: as_int(a) > as_int(b))
        OP[20] = binop(lambda a, b: as_int(a) >= as_int(b))
        OP[23] = binop(lambda a, b: (as_int(a) * as_int(b)) & 0xFFFFFFFF)
        OP[30] = binop(lambda a, b: as_int(a) & as_int(b))
        OP[33] = binop(lambda a, b: (as_int(a) - as_int(b)) & 0xFFFFFFFF)
        OP[35] = binop(lambda a, b: as_int(a) >> (as_int(b) & 31))
        OP[36] = binop(lambda a, b: (as_int(a) << (as_int(b) & 31)) & 0xFFFFFFFF)
        OP[39] = binop(lambda a, b: ((as_int(a) + as_int(b)) & 0xFFFFFFFF) if (isinstance(a, int) and isinstance(b, int)) else (a + b))
        OP[52] = binop(lambda a, b: (as_int(a) / as_int(b)) if as_int(b) else float("nan"))
        OP[62] = binop(lambda a, b: as_int(a) | as_int(b))
        OP[65] = binop(lambda a, b: (as_int(a) & 0xFFFFFFFF) >> (as_int(b) & 31))
        OP[18] = binop(lambda a, b: isinstance(b, (list, dict, JArr)) and (a in (b._ext if isinstance(b, JArr) else b)))

        @op(6)      # CALL_FUNC n
        def _6():
            n = z[self.H]; self.H += 1
            args = E[-n:] if n else []
            del E[len(E) - n:]
            fn = E.pop()
            if os.environ.get("DBGCALL"):
                print(f"[CALL_FUNC @{self.H-2}] n={n} callee={type(fn).__name__} {str(fn)[:70]}", file=sys.stderr)
            E.append(self.call(fn, args, self.c))   # 真实语义: fn.apply(c, B) -> this = c

        @op(8)      # GET_GLOBAL_PROP
        def _8():
            E[-1] = getprop(self.c, E[-1])

        @op(9)      # MAKE_PROP_REF
        def _9():
            a = E.pop(); b = E.pop()
            E.append([getprop(b[0], b[1]), a])

        @op(10)
        def _10():
            E[-1] = type(E[-1]).__name__

        @op(11)
        def _11():
            a = E.pop(); E.append(True)

        @op(12)
        def _12():
            E[-1] = not truthy(E[-1])

        @op(14)     # SET_PROP
        def _14():
            a = E[-2]; setprop(a[0], a[1], E[-1])

        @op(17)     # SWAP n
        def _17():
            n = z[self.H]; self.H += 1
            i = len(E) - 2 - n
            b = E[i]; E[i] = E.pop(); E.append(b)

        @op(19)     # RET_TRUE
        def _19():
            self.ret = True
            return True

        @op(22)     # PUSH_IMM
        def _22():
            E.append(z[self.H]); self.H += 1

        @op(25)     # LOAD_REG_INDIRECT
        def _25():
            cell = E[E.pop()[0]]
            E.append(cell[0] if cell else JS_UNDEF)

        @op(26)     # ITER_NEXT
        def _26():
            it = E[-1]
            if it:
                E.append(it.pop(0)); E.append(True)
            else:
                E.append(JS_UNDEF); E.append(False)

        @op(28)
        def _28():
            E.append([self.c, E.pop()])

        @op(29)
        def _29():
            if os.environ.get("DBGCALL"):
                print(f"[PUSH_NULL @pc={self.H-1}] steps={self.steps}", file=sys.stderr)
            E.append(None)

        @op(31)     # PUSH_IMM_ARRAY
        def _31():
            E.append(JArr([z[self.H]])); self.H += 1

        @op(32)     # MAKE_REG_PROP_REF
        def _32():
            a = E.pop()
            idx = E.pop()
            cell = E[idx[0]] if isinstance(idx, list) else E[idx]
            val = cell[0] if isinstance(cell, list) and cell else JS_UNDEF
            if os.environ.get("DBGCALL") and (val is None or val is JS_UNDEF):
                print(f"[MAKE_REG_PROP_REF @{self.H-1}] reg={idx[0] if isinstance(idx,list) else idx} val={val!r} cell={str(cell)[:40]}", file=sys.stderr)
            E.append([val, a])

        @op(34)     # CALL_METHOD n
        def _34():
            n = z[self.H]; self.H += 1
            args = E[-n:] if n else []
            del E[len(E) - n:]
            ref = E.pop()
            if os.environ.get("DBGCALL"):
                callee = getprop(ref[0], ref[1])
                print(f"[CALL_METHOD @{self.H-2}] n={n} ref={str(ref)[:90]} arity/type={type(ref).__name__}", file=sys.stderr)
                print(f"    stack tail: {[str(x)[:36] for x in E[-6:]]}", file=sys.stderr)
            E.append(self.call(getprop(ref[0], ref[1]), args, ref[0]))

        @op(37)     # JUMP_IF_TRUE
        def _37():
            t = z[self.H]; self.H += 1
            if truthy(E[-1]):
                self.H = t

        @op(38)     # FOR_IN_KEYS
        def _38():
            o = E.pop()
            E.append(JArr(list(o._ext.keys()) if isinstance(o, JArr) else list(o.keys()) if isinstance(o, dict) else []))

        @op(41)     # NEW_METHOD n
        def _41():
            n = z[self.H]; self.H += 1
            args = E[-n:] if n else []
            del E[len(E) - n:]
            args.insert(0, None)
            fn = E.pop()
            E.append(self.new(fn, args))

        @op(42)
        def _42():
            E.append(JS_UNDEF)

        @op(45)
        def _45():
            E.append(False)

        @op(46)     # STORE_REG
        def _46():
            idx = E[-2][0]
            cell = E[idx]
            if not isinstance(cell, list):
                cell = []
                E[idx] = cell
            v = E[-1]
            if cell:
                cell[0] = v
            else:
                cell.append(v)
            self.record_cell(v)

        @op(47)     # ALLOC_REG
        def _47():
            a = z[self.H]; self.H += 1
            while len(E) <= a:
                E.append(JS_UNDEF)
            if not isinstance(E[a], list):
                E[a] = []

        @op(48)
        def _48():
            self.B.pop()

        @op(49)     # THROW
        def _49():
            raise RuntimeError("VM throw")

        @op(51)
        def _51():
            E.pop()

        @op(53)     # SET_STACK_LEN
        def _53():
            n = z[self.H]; self.H += 1
            while len(E) < n:
                E.append(JS_UNDEF)
            del E[n:]

        @op(54)
        def _54():
            a = E.pop(); b = E.pop()
            E.append(JArr([a, b]))

        @op(55)
        def _55():
            E.append("")

        @op(56)     # PUSH_TRY
        def _56():
            self.H += 2

        @op(57)     # CREATE_CLOSURE
        def _57():
            target = z[self.H]; self.H += 1
            cap = z[self.H]; self.H += 1
            par = z[self.H]; self.H += 1
            caps = []
            for _ in range(cap):
                q = z[self.H]; self.H += 1
                r = z[self.H]; self.H += 1
                caps.append((q, r))
            params = []
            for _ in range(par):
                params.append(z[self.H]); self.H += 1
            E.append(("closure", target, caps, params))

        @op(59)     # GET_PROP
        def _59():
            a = E.pop()
            E.append(getprop(a[0], a[1]))

        @op(60)     # JUMP
        def _60():
            self.H = z[self.H]

        @op(61)     # NEW_FUNC n
        def _61():
            n = z[self.H]; self.H += 1
            args = E[-n:] if n else []
            del E[len(E) - n:]
            args.insert(0, None)
            fn = E.pop()
            E.append(self.new(fn, args))

        @op(64)     # REPLACE_TOP_IMM
        def _64():
            E[-1] = z[self.H]; self.H += 1

        @op(66)
        def _66():
            E.append(True)

        @op(68)     # LOAD_REG_IMM
        def _68():
            a = z[self.H]; self.H += 1
            cell = E[a] if a < len(E) else None
            E.append(cell[0] if isinstance(cell, list) and cell else JS_UNDEF)

        @op(69)
        def _69():
            self.w = None

        self.OP = OP
        return self.run()

    # ---- 调用/构造
    def call(self, fn, args, this):
        if isinstance(fn, tuple) and fn and fn[0] == "closure":
            sub = VM(self.z, self.c, self.max_steps - self.steps, fresh=False)
            sub.key_hits = self.key_hits           # 子 VM 共享命中列表
            sub.pcs = self.pcs                     # 共享 PC 轨迹 (轨迹是全局的)
            sub.events = self.events               # 共享关键事件
            sub.trace = self.trace
            sub.trace_full = self.trace_full
            if os.environ.get("DBGCALL"):
                print(f"[invoke closure target=@{fn[1]}] caps={fn[2]} params={fn[3]} args={[str(a)[:20] for a in args]}", file=sys.stderr)
                for (q, r) in fn[2]:
                    src = self.E[r] if r < len(self.E) else "<oob>"
                    print(f"    cap r{q} <- parent r{r} = {str(src)[:50]}", file=sys.stderr)
            sub.E = [JArr() for _ in range(32)]
            sub.E[0] = JArr([this]); sub.E[1] = JArr([JArr(args)]); sub.E[2] = JArr([fn])
            for (q, r) in fn[2]:
                sub.E[q] = self.E[r]
            for i, r in enumerate(fn[3]):
                sub.E[r] = JArr([args[i] if i < len(args) else JS_UNDEF])
            sub.H = fn[1]
            try:
                v = sub._run()
                self.steps += sub.steps
                return v
            except Exception as exc:
                if os.environ.get("DBGCALL"):
                    import traceback
                    print(f"[closure @{fn[1]} 执行异常] {type(exc).__name__}: {exc}", file=sys.stderr)
                    traceback.print_exc(limit=3)
                return JS_UNDEF
        if isinstance(fn, HostFn):
            return fn(*args)
        if callable(fn) and not isinstance(fn, (JArr, Stub, tuple)):
            try:
                return fn(*args)
            except Exception:
                return JS_UNDEF
        if isinstance(fn, Stub):
            return Stub("call")
        return JS_UNDEF

    def new(self, fn, args):
        return self.call(fn, args, {})


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)
    known = os.environ.get("KEY")
    known_key = [int(x) for x in known.split(",")] if known else None
    print(f"bytecode={len(z)}  known_key={known_key}")

    vm = VM(z, make_host(), max_steps=30_000_000)
    try:
        vm._run()
    except Exception as e:
        print("VM halted with:", type(e).__name__, str(e)[:200])
    print("steps:", vm.steps, "done:", vm.done, "ret:", str(vm.ret)[:80])
    names = {1: "IS_EXCEPTION", 6: "CALL_FUNC", 19: "RET_TRUE", 34: "CALL_METHOD",
             41: "NEW_METHOD", 49: "THROW", 57: "CREATE_CLOSURE", 61: "NEW_FUNC"}
    print("关键事件 (pc, op):")
    for pc, op in vm.events[:40]:
        print(f"   @{pc:05d} {names.get(op, op)}")
    print(f"   ... 共 {len(vm.events)} 条")

    # tdc.js 顶层只"注册"模块工厂; 这里把寄存器/数组里保存的闭包全收集出来 = 模块表
    def collect_closures(obj, out, seen):
        oid = id(obj)
        if oid in seen:
            return
        seen.add(oid)
        if isinstance(obj, tuple) and obj and obj[0] == "closure":
            out.append(obj)
            return
        if isinstance(obj, Stub):
            for v in obj._props.values():
                collect_closures(v, out, seen)
        elif isinstance(obj, (JArr, list)):
            for v in list(obj):
                collect_closures(v, out, seen)
            if isinstance(obj, JArr):
                for v in obj._ext.values():
                    collect_closures(v, out, seen)
        elif isinstance(obj, dict):
            for v in obj.values():
                collect_closures(v, out, seen)

    closures: List[Any] = []
    collect_closures([vm.E, vm.c], closures, set())
    print("模块表: 收集到闭包", len(closures), "个 -> target:", [c[1] for c in closures][:80])
    for i, cl in enumerate(closures):
        try:
            vm.call(cl, [Stub("arg")], {})
        except Exception as e:
            print(f"   closure#{i} (target=@{cl[1]}) raised {type(e).__name__}: {str(e)[:80]}")
    print("total steps after invoking closures:", vm.steps)
    hits = [h for h in vm.key_hits if known_key is None or h == known_key]
    print("candidate 4-char packed arrays:", len(hits))
    for h in hits[:10]:
        print("   ", h, "->", b"".join(x.to_bytes(4, "big") for x in h))
    return 0


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