#!/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):
        if isinstance(k, int) or (isinstance(k, str) and k.lstrip("-").isdigit()):
            i = int(k)
            try:
                return list.__getitem__(self, i)
            except IndexError:
                return None
        return self._ext.get(k)

    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()


class Stub:
    """未知宿主对象的占位: 可调用、可属性访问、可下标, 一律返回自身/None。"""

    def __init__(self, name="stub"):
        self._name = name

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

    def __getattr__(self, k):
        if k.startswith("_"):
            raise AttributeError(k)
        return Stub(f"{self._name}.{k}")

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

    def __setitem__(self, k, v):
        pass

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


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


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

    class MathMod:
        pass

    win._ext.update({
        "Math": Stub("Math"),
        "JSON": type("J", (), {"parse": staticmethod(json.loads), "stringify": staticmethod(json.dumps)})(),
        "Object": JArr,
        "Array": JArr,
        "String": str,
        "Number": float,
        "Boolean": bool,
        "Date": int(time.time() * 1000),
        "RegExp": Stub("RegExp"),
        "Error": Stub("Error"),
        "TypeError": Stub("TypeError"),
        "Uint8Array": bytearray,
        "Uint8ClampedArray": bytearray,
        "parseInt": int,
        "parseFloat": float,
        "encodeURIComponent": lambda s: s,
        "decodeURIComponent": lambda s: s,
        "atob": lambda s: s,
        "btoa": lambda s: s,
        "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):
        self.z = z
        self.c = c
        self.E: List[Any] = []          # 寄存器区 + 栈
        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] = []

    # ---- 工具
    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
            h = self.OP.get(op)
            if h is None:
                raise RuntimeError(f"unknown opcode {op} at pc {self.H-1}")
            if h() is True:
                self.done = True
        return self.ret

    # ------------------------------------------------------------------ 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()
            E.append(self.call(fn, args, None))

        @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():
            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]
            E.append([cell[0] if isinstance(cell, list) and cell else JS_UNDEF, 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()
            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)
            sub.key_hits = self.key_hits           # 子 VM 共享命中列表
            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:
                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])

    # tdc.js 顶层只"注册"模块工厂; 这里主动调用寄存器里保存的闭包 (模拟模块系统/宿主调用)
    closures = [cell[0] for cell in vm.E
                if isinstance(cell, list) and cell and isinstance(cell[0], tuple)
                and cell[0] and cell[0][0] == "closure"]
    print("registered closures in registers:", len(closures))
    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())