# -*- coding: utf-8 -*-
"""
tdc_program.py —— TDC Chaos VM 「程序结构层」分析器。

它补的是 VM_SPEC.md §5/§6 自己点出的缺口:
  opcode 表 / 闭包调用约定 早已吃透, 但**模块系统 + 宿主调用序列**一直没被静态说明:
  顶层只"装配模块", 真正干活的代码在闭包被宿主调用时才跑。

本工具把这些做出确定规格:
  * `closures`   —— 全部 CREATE_CLOSURE 站点: (pc, target, caps[(dst,src)], params[reg])
  * `entries`    —— 子程序入口集合 = {0} ∪ {所有 closure target}
  * `strings`    —— 从字节码折叠出的字面串 (PUSH_EMPTY_STR + 连续 PUSH_STR_CHAR)
  * `owner(pc)`  —— 某 pc 属于哪个子程序 (启发式: 覆盖它的、且 ≤pc 的最大入口)
  * `call_graph` —— 子程序 → 它在自己体内创建的闭包 target (静态调用图的"创建边")
  * `xfuncs`     —— 同一 target 被多个站点创建的"共享体" (VM 里的复用函数)

注意 (踩过的坑):
  `_decode_one` 对 CREATE_CLOSURE 的 `operands` **不含 target**(target 只写在 note 里),
  必须自己按 `target = z[pc+1]` 取, 否则会把 cap 寄存器号当成 target。

用法:
    .venv/bin/python tools/tdc_program.py tcaptcha/core_js/tdc_sample.js --summary
    .venv/bin/python tools/tdc_program.py <js> --seq --max 30
    .venv/bin/python tools/tdc_program.py <js> --find getData
    .venv/bin/python tools/tdc_program.py <js> --at 1725
"""
from __future__ import annotations

import argparse
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from tcaptcha.tdc_opcodes import derive_opcode_table  # noqa: E402
from tools.tdc_disasm import disassemble_cfg, extract_bytecode  # noqa: E402

MIN_STR_LEN = 3


@dataclass
class ClosureSite:
    pc: int
    target: int
    caps: List[Tuple[int, int]] = field(default_factory=list)   # (新闭包寄存器, 父闭包寄存器)
    params: List[int] = field(default_factory=list)             # 形参落到的寄存器号


@dataclass
class Program:
    path: str
    z: List[int]
    names: Dict[int, str]
    arity: Dict[int, int]
    dynamic: set
    insts: Dict[int, dict]
    closures: List[ClosureSite]
    strings: List[Tuple[int, str]]
    entries: List[int]
    _owner: Dict[int, int] = field(default_factory=dict)

    # -------------------------------------------------------------- 基本查询
    def name(self, pc: int) -> str:
        i = self.insts.get(pc)
        return i["name"] if i else "?"

    def owner(self, pc: int) -> int:
        """pc 属于哪个子程序入口。

        注意: 代码布局是「函数体之间夹着父函数的初始化序列 + 内联的子函数体」,
        所以**不能**用"覆盖它的最大入口 ≤pc"来判 (会把父函数续段误判给夹在中间的兄弟函数)。
        这里用**边界感知走法**: 从入口开始走控制流, 一旦走到**另一个函数入口**就停下
        (函数只能被调用进入, 不会被穿过)。
        """
        if not self._owner:
            self._build_owners()
        return self._owner.get(pc, 0)

    def _next_pc(self, pc: int, inst: dict) -> int:
        """下一条 fall-through 指令的 pc。CREATE_CLOSURE 是动态长度, 不能用固定 arity。"""
        if inst["name"] == "CREATE_CLOSURE":
            capn, parn = self.z[pc + 2], self.z[pc + 3]
            return pc + 1 + 3 + 2 * capn + parn
        return pc + 1 + self.arity.get(inst["op"], 0)

    def function_code(self, entry: int) -> List[dict]:
        """边界感知地取某个子程序**自己的**指令序列。"""
        starts = set(self.entries)
        out: Dict[int, dict] = {}
        work = [entry]
        while work:
            pc = work.pop()
            if pc in out or pc >= len(self.z) or pc < 0:
                continue
            if pc != entry and pc in starts:      # 进入别的函数体 -> 不跨越
                continue
            inst = self.insts.get(pc)
            if inst is None:
                continue
            out[pc] = inst
            name = inst["name"]
            ops = inst["operands"] or []
            if name == "JUMP":
                work.append(ops[0])
                continue
            if name == "JUMP_IF_TRUE":
                work.append(ops[0])
                work.append(self._next_pc(pc, inst))
                continue
            if name == "PUSH_TRY" and self.arity.get(inst["op"], 0) == 2:
                work.append(ops[0])
            if name == "CREATE_CLOSURE":
                work.append(self.z[pc + 1])
            if name in ("RET_TRUE", "THROW"):
                continue
            work.append(self._next_pc(pc, inst))
        return [out[k] for k in sorted(out)]

    def _build_owners(self) -> None:
        for e in self.entries:
            for inst in self.function_code(e):
                self._owner.setdefault(inst["pc"], e)

    def closures_of(self, entry: int) -> List[ClosureSite]:
        return [c for c in self.closures if self.owner(c.pc) == entry]

    def call_graph(self) -> Dict[int, List[int]]:
        g: Dict[int, List[int]] = {}
        for c in self.closures:
            g.setdefault(self.owner(c.pc), []).append(c.target)
        return {k: sorted(set(v)) for k, v in g.items()}

    def shared_bodies(self) -> Dict[int, List[int]]:
        """target -> 创建它的站点列表 (长度>1 说明这些站点共用同一个代码体)。"""
        m: Dict[int, List[int]] = {}
        for c in self.closures:
            m.setdefault(c.target, []).append(c.pc)
        return {k: v for k, v in m.items() if len(v) > 1}

    def module_registry(self, entry: int = 0) -> List[Tuple[int, int]]:
        """识别顶层「模块注册表装配」模式, 返回 [(模块 id, 工厂闭包 entry), ...]。

        模式 (源码级 `registry[id] = function(module, exports, require){...}` 的编译产物):
            DUP; PUSH_IMM [id]; REVERSE_TOP2; JUMP -> @F; @F: CREATE_CLOSURE ...; SET_PROP

        这是整个程序的骨架: 顶层只把 228 个模块工厂按 id 填进一个数组, 入口再 drive 它。
        """
        code = self.function_code(entry)
        pairs: List[Tuple[int, int]] = []
        for k, inst in enumerate(code):
            if inst["name"] != "CREATE_CLOSURE":
                continue
            for j in range(k - 1, max(-1, k - 4), -1):
                if code[j]["name"] != "JUMP":
                    continue
                ops = code[j]["operands"] or []
                if not ops or ops[0] != inst["pc"]:
                    continue
                for m in range(j - 1, max(-1, j - 6), -1):
                    if code[m]["name"] == "PUSH_IMM" and code[m]["operands"]:
                        # 记的是**工厂函数体入口** (CREATE_CLOSURE 的 target), 不是创建站点
                        pairs.append((int(code[m]["operands"][0]), self.z[inst["pc"] + 1]))
                        break
                break
        seen, out = set(), []
        for mid, tgt in sorted(pairs):
            if mid in seen:
                continue
            seen.add(mid)
            out.append((mid, tgt))
        return out

    # ---------------------------------------------------------- 轻量符号求值
    def symbolic_assignments(self, entry: int, limit: int = 30) -> List[Tuple[str, str]]:
        """对一个小函数(如采集器模块)做**线性符号求值**, 还原 `rN = <表达式>`。

        只覆盖这类小模块实际用到的形状: 字面串 / BIND_THIS / MAKE_(REG_)PROP_REF /
        GET_PROP / GET_GLOBAL_PROP / NEW_FUNC / CALL_FUNC / CALL_METHOD / STORE_REG /
        LOAD_REG_IMM / REVERSE_TOP2 / DUP / 算术。

        目的: 把"这一槽取的是什么值"变成可读路径, 例如 `c.navigator.userAgent`。
        说明: 遇到分支时按 pc 顺序近似处理, 结果用于**定位语义**, 不作为执行依据。
        """
        regs: Dict[int, str] = {}
        out: List[Tuple[str, str]] = []
        st: List[Tuple[str, str]] = []          # ("reg", n) | ("sym", text)

        def top() -> Tuple[str, str]:
            return st[-1] if st else ("sym", "?")

        def popn(n: int) -> List[Tuple[str, str]]:
            r = st[-n:] if n else []
            del st[len(st) - n:]
            return r

        for inst in self.function_code(entry):
            n, ops = inst["name"], inst["operands"] or []
            if n == "PUSH_IMM_ARRAY":
                st.append(("reg", ops[0]))
            elif n == "PUSH_EMPTY_STR":
                st.append(("sym", ""))
            elif n == "PUSH_STR_CHAR":
                k, v = top()
                st[-1] = ("sym", v + chr(ops[0] & 0xFF))
            elif n == "PUSH_IMM":
                st.append(("sym", str(ops[0])))
            elif n == "PUSH_UNDEFINED":
                st.append(("sym", "undefined"))
            elif n == "PUSH_NULL":
                st.append(("sym", "null"))
            elif n == "PUSH_TRUE":
                st.append(("sym", "true"))
            elif n == "PUSH_FALSE":
                st.append(("sym", "false"))
            elif n == "LOAD_REG_IMM":
                st.append(("sym", regs.get(ops[0], f"r{ops[0]}")))
            elif n == "LOAD_REG_INDIRECT":
                k, v = st.pop() if st else ("sym", "?")
                st.append(("sym", regs.get(int(v), f"r{v}") if str(v).isdigit() else v))
            elif n == "BIND_THIS":
                k, v = st.pop() if st else ("sym", "?")
                st.append(("sym", f"c.{v}"))
            elif n == "MAKE_PROP_REF":
                key = top()[1]
                popn(1)
                base = top()[1]
                popn(1)
                st.append(("sym", f"{base}.{key}"))
            elif n == "MAKE_REG_PROP_REF":
                key = top()[1]
                popn(1)
                k, v = st.pop() if st else ("sym", "?")
                st.append(("sym", f"{regs.get(v, f'r{v}')}.{key}"))
            elif n == "GET_GLOBAL_PROP":
                st[-1] = ("sym", f"c.{top()[1]}")
            elif n in ("GET_PROP", "GET_ELEM"):
                pass
            elif n in ("ADD", "SUB", "MUL", "DIV", "MOD", "BIT_AND", "BIT_OR", "BIT_XOR",
                       "SHL", "SHR", "SHR_U", "EQ", "EQ_STRICT", "GREATER"):
                if len(st) >= 2:
                    b = top()[1]
                    popn(1)
                    a = top()[1]
                    popn(1)
                    st.append(("sym", f"({a} {n} {b})"))
            elif n in ("NEW_FUNC", "NEW_METHOD"):
                args = popn(int(ops[0]) if ops else 0)
                fn = st.pop()[1] if st else "?"
                st.append(("sym", f"new {fn}({','.join(x[1] for x in args)})"))
            elif n == "CALL_FUNC":
                args = popn(int(ops[0]) if ops else 0)
                fn = st.pop()[1] if st else "?"
                st.append(("sym", f"{fn}({','.join(x[1] for x in args)})"))
            elif n == "CALL_METHOD":
                args = popn(int(ops[0]) if ops else 0)
                tgt = st.pop() if st else ("sym", "?")
                st.append(("sym", f"{tgt[1]}({','.join(x[1] for x in args)})"))
            elif n == "STORE_REG":
                v = st.pop()[1] if st else "?"
                k, ref = st.pop() if st else ("sym", "?")
                if k == "reg":
                    regs[ref] = v
                    out.append((f"r{ref}", v))
            elif n == "DUP":
                if st:
                    st.append(st[-1])
            elif n == "REVERSE_TOP2":
                if len(st) >= 2:
                    st[-1], st[-2] = st[-2], st[-1]
            elif n in ("POP", "SET_PROP", "ALLOC_REG", "SET_STACK_LEN", "REPLACE_TOP_IMM"):
                if n in ("POP", "SET_PROP") and st:
                    st.pop()
            if len(out) >= limit:
                break
        return out

    def find_strings(self, pattern: str) -> List[Tuple[int, str, int]]:
        """返回 (pc, 串, 所属子程序入口)。"""
        rx = re.compile(pattern, re.I)
        return [(pc, s, self.owner(pc)) for pc, s in self.strings if rx.search(s)]


def _fold_strings(insts: Dict[int, dict]) -> List[Tuple[int, str]]:
    """PUSH_EMPTY_STR 后跟一串 PUSH_STR_CHAR -> 还原字面串 (tdc_disasm 不做这步)。"""
    out: List[Tuple[int, str]] = []
    cur_pc: Optional[int] = None
    buf: List[str] = []
    for pc in sorted(insts):
        i = insts[pc]
        n = i["name"]
        if n == "PUSH_EMPTY_STR":
            cur_pc, buf = pc, []
            continue
        if n == "PUSH_STR_CHAR" and cur_pc is not None:
            ops = i["operands"] or []
            if ops:
                buf.append(chr(ops[0] & 0xFF))
            continue
        if cur_pc is not None:
            if len(buf) >= MIN_STR_LEN:
                out.append((cur_pc, "".join(buf)))
            cur_pc, buf = None, []
    if cur_pc is not None and len(buf) >= MIN_STR_LEN:
        out.append((cur_pc, "".join(buf)))
    return out


def load(js_path: str) -> Program:
    js = open(js_path, encoding="utf-8", errors="replace").read()
    z = extract_bytecode(js)
    names, arity, dynamic, _diag = derive_opcode_table(js)
    insts = {i["pc"]: i for i in
             disassemble_cfg(z, 0, True, names=names, arity=arity, dynamic=dynamic)
             if i["op"] is not None}

    closures: List[ClosureSite] = []
    for pc, i in insts.items():
        if i["name"] != "CREATE_CLOSURE":
            continue
        target, cap, par = z[pc + 1], z[pc + 2], z[pc + 3]
        base = pc + 4
        caps = [(z[base + 2 * k], z[base + 2 * k + 1]) for k in range(cap)]
        params = [z[base + 2 * cap + k] for k in range(par)]
        closures.append(ClosureSite(pc=pc, target=target, caps=caps, params=params))

    strings = _fold_strings(insts)
    entries = sorted({0} | {c.target for c in closures})
    return Program(path=js_path, z=z, names=names, arity=arity, dynamic=dynamic,
                   insts=insts, closures=sorted(closures, key=lambda c: c.pc),
                   strings=strings, entries=entries)


# --------------------------------------------------------------------------- CLI
def main() -> int:
    ap = argparse.ArgumentParser(description="TDC Chaos VM 程序结构分析")
    ap.add_argument("path", nargs="?", default=str(ROOT / "tcaptcha/core_js/tdc_sample.js"))
    ap.add_argument("--summary", action="store_true", help="总体规模")
    ap.add_argument("--seq", action="store_true", help="按子程序列出闭包站点")
    ap.add_argument("--max", type=int, default=40, help="--seq/--find 的截断条数")
    ap.add_argument("--find", default=None, help="按正则找字面串")
    ap.add_argument("--at", type=int, default=None, help="看某个 pc 属于谁")
    ap.add_argument("--shared", action="store_true", help="列出被多个站点共享的闭包体")
    ap.add_argument("--func", type=int, default=None, help="打印某个子程序**自己的**指令序列")
    ap.add_argument("--asm-out", default=None, help="--func 时把反汇编写到该文件")
    ap.add_argument("--modules", action="store_true", help="导出模块注册表 (id -> 工厂 entry)")
    args = ap.parse_args()

    p = load(args.path)
    shown = False

    if args.summary or not any([args.seq, args.find is not None, args.at is not None,
                                args.shared, args.func is not None]):
        shown = True
        print(f"file            : {p.path}")
        print(f"bytecode        : {len(p.z)} ints")
        print(f"reachable instrs: {len(p.insts)}")
        print(f"closure sites   : {len(p.closures)}")
        print(f"subprograms     : {len(p.entries)}  (0 为顶层)")
        print(f"distinct bodies : {len({c.target for c in p.closures})}")
        print(f"literal strings : {len(p.strings)}  (>= {MIN_STR_LEN} 字符)")

    if args.seq:
        shown = True
        g = p.call_graph()
        print(f"\n--- 子程序与它们创建的闭包 (前 {args.max} 个入口) ---")
        for e in p.entries[: args.max]:
            kids = g.get(e, [])
            st = [s for pc, s in p.strings if p.owner(pc) == e]
            head = ", ".join(f"{s[:22]!r}" for s in st[:4])
            print(f"  @{e:05d}  ->  {kids}")
            if head:
                print(f"           串: {head}")

    if args.find is not None:
        shown = True
        hits = p.find_strings(args.find)
        print(f"\n--- 字面串匹配 {args.find!r}: {len(hits)} 条 ---")
        for pc, s, own in hits[: args.max]:
            print(f"  @{pc:05d} [子程序 @{own:05d}] {s[:90]!r}")

    if args.at is not None:
        shown = True
        pc = args.at
        print(f"\n--- pc @{pc} ---")
        print(f"  指令 : {p.insts.get(pc)}")
        print(f"  所属 : @{p.owner(pc)}")

    if args.func is not None:
        shown = True
        code = p.function_code(args.func)
        lines = []
        for x in code:
            extra = (" " + " ".join(x["note"])) if x["note"] else (
                (" " + str(x["operands"])) if x["operands"] else "")
            lines.append(f"@{x['pc']:05d}: {x['name']:<18}{extra}")
        print(f"\n--- 子程序 @{args.func} 自身指令: {len(lines)} 条 ---")
        print("\n".join(lines[: args.max]))
        if args.asm_out:
            Path(args.asm_out).write_text("\n".join(lines), encoding="utf-8")
            print(f"    -> {args.asm_out}")

    if args.modules:
        shown = True
        reg = p.module_registry(0)
        print(f"\n--- 模块注册表: {len(reg)} 个模块 (registry[id] = 工厂闭包) ---")
        for mid, tgt in reg[: args.max]:
            print(f"  registry[{mid:3d}] = 工厂 @{tgt}")
        if len(reg) > args.max:
            print(f"  ... 共 {len(reg)} 个")

    if args.shared:
        shown = True
        sh = p.shared_bodies()
        print(f"\n--- 共享闭包体: {len(sh)} 个 target 被多处创建 ---")
        for t, sites in sorted(sh.items(), key=lambda kv: -len(kv[1]))[: args.max]:
            print(f"  target @{t:05d}  x{len(sites)}  sites={sites[:8]}{'...' if len(sites) > 8 else ''}")

    return 0 if shown else 1


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