#!/usr/bin/env python3
"""Minimal AVM1 (ActionScript 2) disassembler for EA APT movies.

Enough of the opcode table to READ a predicate: pushes, member access, calls,
branches and comparisons. Not a decompiler and not complete — anything unknown is
printed as a raw opcode with its length so the reader can see there is a gap
rather than silently skipping it.

APT stores the same AVM1 action blocks a SWF DoAction does, so a constant pool
(0x88) followed by pushes that index it is the normal shape.

    python3 avm1_disasm.py movie.bin --pool
    python3 avm1_disasm.py movie.bin --around CheckIsKitLocked --span 900
"""
from __future__ import annotations

import argparse
import struct
import sys

# opcode -> (name, has_payload)
OPS = {
    0x04: "NextFrame", 0x06: "Play", 0x07: "Stop", 0x0A: "Add", 0x0B: "Subtract",
    0x0C: "Multiply", 0x0D: "Divide", 0x0E: "Equals", 0x0F: "Less",
    0x10: "And", 0x11: "Or", 0x12: "Not", 0x13: "StringEquals",
    0x17: "Pop", 0x18: "ToInteger",
    0x1C: "GetVariable", 0x1D: "SetVariable",
    0x21: "StringAdd", 0x22: "GetProperty", 0x23: "SetProperty",
    0x24: "CloneSprite", 0x26: "Trace", 0x28: "EndDrag",
    0x2A: "Throw", 0x2B: "CastOp", 0x2C: "ImplementsOp",
    0x30: "RandomNumber", 0x3A: "Delete", 0x3B: "Delete2",
    0x3C: "DefineLocal", 0x3D: "CallFunction", 0x3E: "Return",
    0x3F: "Modulo", 0x40: "NewObject", 0x41: "DefineLocal2",
    0x42: "InitArray", 0x43: "InitObject", 0x44: "TypeOf",
    0x46: "Enumerate", 0x47: "Add2", 0x48: "Less2", 0x49: "Equals2",
    0x4A: "ToNumber", 0x4B: "ToString", 0x4C: "PushDuplicate", 0x4D: "StackSwap",
    0x4E: "GetMember", 0x4F: "SetMember", 0x50: "Increment", 0x51: "Decrement",
    0x52: "CallMethod", 0x53: "NewMethod", 0x54: "InstanceOf", 0x55: "Enumerate2",
    0x60: "BitAnd", 0x61: "BitOr", 0x62: "BitXor",
    0x66: "StrictEquals", 0x67: "Greater", 0x68: "StringGreater",
    0x69: "Extends",
}
PAYLOAD = {
    0x81: "GotoFrame", 0x83: "GetURL", 0x87: "StoreRegister", 0x88: "ConstantPool",
    0x8A: "WaitForFrame", 0x8B: "SetTarget", 0x8C: "GotoLabel",
    0x8D: "WaitForFrame2", 0x8E: "DefineFunction2", 0x8F: "Try",
    0x94: "With", 0x96: "Push", 0x99: "Jump", 0x9A: "GetURL2",
    0x9B: "DefineFunction", 0x9D: "If", 0x9E: "Call", 0x9F: "GotoFrame2",
}


def parse_push(data: bytes, pool: list[str]) -> list[str]:
    out, i = [], 0
    while i < len(data):
        t = data[i]; i += 1
        try:
            if t == 0:
                e = data.index(b"\0", i); out.append(repr(data[i:e].decode("latin1"))); i = e + 1
            elif t == 1:
                out.append(f"{struct.unpack_from('<f', data, i)[0]:g}"); i += 4
            elif t == 2:
                out.append("null")
            elif t == 3:
                out.append("undefined")
            elif t == 4:
                out.append(f"reg{data[i]}"); i += 1
            elif t == 5:
                out.append("true" if data[i] else "false"); i += 1
            elif t == 6:
                out.append(f"{struct.unpack_from('<d', data, i)[0]:g}"); i += 8
            elif t == 7:
                out.append(str(struct.unpack_from("<i", data, i)[0])); i += 4
            elif t in (8, 9):
                idx = data[i] if t == 8 else struct.unpack_from("<H", data, i)[0]
                i += 1 if t == 8 else 2
                out.append(f"c{idx}:{pool[idx]!r}" if idx < len(pool) else f"c{idx}")
            else:
                out.append(f"?type{t}"); break
        except Exception:
            out.append("<truncated>"); break
    return out


def disasm(buf: bytes, start: int, end: int, pool: list[str]):
    i, lines = start, []
    while i < end and i < len(buf):
        op = buf[i]
        if op == 0:
            i += 1
            continue
        if op < 0x80:
            lines.append((i, OPS.get(op, f"op{op:#04x}"), ""))
            i += 1
            continue
        if i + 3 > len(buf):
            break
        ln = struct.unpack_from("<H", buf, i + 1)[0]
        body = buf[i + 3:i + 3 + ln]
        name = PAYLOAD.get(op, f"op{op:#04x}")
        arg = ""
        if op == 0x96:
            arg = ", ".join(parse_push(body, pool))
        elif op in (0x99, 0x9D) and ln >= 2:
            arg = f"{struct.unpack_from('<h', body, 0)[0]:+d}"
        elif op == 0x88 and ln >= 2:
            arg = f"{struct.unpack_from('<H', body, 0)[0]} entries"
        elif op in (0x8E, 0x9B):
            e = body.index(b"\0") if b"\0" in body else 0
            arg = body[:e].decode("latin1")
        elif op == 0x87 and ln >= 1:
            arg = f"reg{body[0]}"
        else:
            arg = body[:40].decode("latin1", "replace").replace("\0", ".")
        lines.append((i, name, arg))
        i += 3 + ln
    return lines


def constant_pools(buf: bytes):
    """Every ActionConstantPool in the blob, as (offset, [strings])."""
    pools, i = [], 0
    while i < len(buf) - 3:
        if buf[i] == 0x88:
            ln = struct.unpack_from("<H", buf, i + 1)[0]
            body = buf[i + 3:i + 3 + ln]
            if len(body) >= 2:
                cnt = struct.unpack_from("<H", body, 0)[0]
                parts = body[2:].split(b"\0")
                if 0 < cnt <= 4000 and len(parts) >= cnt:
                    pools.append((i, [p.decode("latin1") for p in parts[:cnt]]))
                    i += 3 + ln
                    continue
        i += 1
    return pools


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("file")
    ap.add_argument("--pool", action="store_true")
    ap.add_argument("--around")
    ap.add_argument("--span", type=int, default=700)
    ap.add_argument("--at", type=lambda s: int(s, 0))
    a = ap.parse_args()

    buf = open(a.file, "rb").read()
    pools = constant_pools(buf)
    pool = max((p for _, p in pools), key=len, default=[])

    if a.pool:
        for off, p in pools:
            print(f"== ConstantPool @ {off:#x}: {len(p)} entries ==")
            for n, s in enumerate(p):
                print(f"  c{n:<4d} {s}")
        return 0

    anchor = a.at
    if a.around:
        anchor = buf.find(a.around.encode())
        if anchor < 0:
            print(f"{a.around!r} not found", file=sys.stderr)
            return 1
        print(f"# anchor {a.around!r} @ {anchor:#x}")
    if anchor is None:
        ap.error("--pool, --around or --at required")

    lo = max(0, anchor - a.span)
    for off, name, arg in disasm(buf, lo, anchor + a.span, pool):
        print(f"  {off:#08x}  {name:<18s} {arg}")
    return 0


if __name__ == "__main__":
    sys.exit(main())