pyflashkit 1.0.0__py3-none-any.whl

flashkit/abc/writer.py

@@ -0,0 +1,230 @@
+"""
+ABC bytecode serializer.
+
+Serializes an ``AbcFile`` structure back to raw ABC binary data.
+The output is byte-for-byte identical to the original input when
+no modifications have been made (round-trip fidelity).
+
+Usage::
+
+    from flashkit.abc.parser import parse_abc
+    from flashkit.abc.writer import serialize_abc
+
+    abc = parse_abc(raw_bytes)
+    # ... modify abc ...
+    output = serialize_abc(abc)
+
+Reference: Adobe AVM2 Overview, Chapter 4 (abc file format).
+"""
+
+from __future__ import annotations
+
+import struct
+
+from ..errors import SerializeError
+from .types import AbcFile, TraitInfo
+from .parser import write_u30, write_s32
+from .constants import (
+    CONSTANT_QName, CONSTANT_QNameA,
+    CONSTANT_RTQName, CONSTANT_RTQNameA,
+    CONSTANT_RTQNameL, CONSTANT_RTQNameLA,
+    CONSTANT_Multiname, CONSTANT_MultinameA,
+    CONSTANT_MultinameL, CONSTANT_MultinameLA,
+    CONSTANT_TypeName,
+    METHOD_HasOptional, METHOD_HasParamNames,
+    INSTANCE_ProtectedNs,
+)
+
+
+def _write_traits(traits: list[TraitInfo]) -> bytes:
+    """Serialize a trait list using the raw binary data stored during parse."""
+    out = write_u30(len(traits))
+    for t in traits:
+        out += t.data
+    return out
+
+
+def serialize_abc(abc: AbcFile) -> bytes:
+    """Serialize an AbcFile back to raw ABC bytecode.
+
+    Args:
+        abc: The AbcFile to serialize.
+
+    Returns:
+        Raw ABC bytecode bytes ready to embed in a DoABC/DoABC2 tag.
+
+    Raises:
+        SerializeError: If the AbcFile structure is invalid.
+    """
+    try:
+        return _serialize_abc_inner(abc)
+    except SerializeError:
+        raise
+    except (IndexError, struct.error, ValueError, TypeError,
+            AttributeError) as e:
+        raise SerializeError(f"Failed to serialize ABC: {e}") from e
+
+
+def _serialize_abc_inner(abc: AbcFile) -> bytes:
+    """Internal serializer (no error wrapping)."""
+    out = bytearray()
+    out += struct.pack("<HH", abc.minor_version, abc.major_version)
+
+    # ── Constant pool ────────────────────────────────────────────────────
+
+    # Integers (count 0 means empty pool — only the implicit default entry)
+    # Use raw LEB128 bytes when available for round-trip fidelity, since
+    # the AVM2 spec allows non-minimal s32 encodings.
+    int_extra = abc.int_pool[1:]
+    int_raw = abc._int_pool_raw[1:] if len(abc._int_pool_raw) > 1 else []
+    out += write_u30(len(int_extra) + 1 if int_extra else 0)
+    for i, v in enumerate(int_extra):
+        if i < len(int_raw) and int_raw[i]:
+            out += int_raw[i]
+        else:
+            out += write_s32(v)
+
+    # Unsigned integers
+    # Use raw bytes for round-trip fidelity (AVM2 uint values can exceed
+    # 30 bits despite the spec calling the encoding "u30").
+    uint_extra = abc.uint_pool[1:]
+    uint_raw = abc._uint_pool_raw[1:] if len(abc._uint_pool_raw) > 1 else []
+    out += write_u30(len(uint_extra) + 1 if uint_extra else 0)
+    for i, v in enumerate(uint_extra):
+        if i < len(uint_raw) and uint_raw[i]:
+            out += uint_raw[i]
+        else:
+            out += write_u30(v)
+
+    # Doubles
+    dbl_extra = abc.double_pool[1:]
+    out += write_u30(len(dbl_extra) + 1 if dbl_extra else 0)
+    for v in dbl_extra:
+        out += struct.pack("<d", v)
+
+    # Strings
+    str_extra = abc.string_pool[1:]
+    out += write_u30(len(str_extra) + 1 if str_extra else 0)
+    for s in str_extra:
+        encoded = s.encode("utf-8")
+        out += write_u30(len(encoded))
+        out += encoded
+
+    # Namespaces
+    ns_extra = abc.namespace_pool[1:]
+    out += write_u30(len(ns_extra) + 1 if ns_extra else 0)
+    for ns in ns_extra:
+        out += bytes([ns.kind])
+        out += write_u30(ns.name)
+
+    # Namespace sets
+    nss_extra = abc.ns_set_pool[1:]
+    out += write_u30(len(nss_extra) + 1 if nss_extra else 0)
+    for nss in nss_extra:
+        out += write_u30(len(nss.namespaces))
+        for ns in nss.namespaces:
+            out += write_u30(ns)
+
+    # Multinames
+    mn_extra = abc.multiname_pool[1:]
+    out += write_u30(len(mn_extra) + 1 if mn_extra else 0)
+    for mn in mn_extra:
+        out += bytes([mn.kind])
+        if mn.kind in (CONSTANT_QName, CONSTANT_QNameA):
+            out += write_u30(mn.ns)
+            out += write_u30(mn.name)
+        elif mn.kind in (CONSTANT_RTQName, CONSTANT_RTQNameA):
+            out += write_u30(mn.name)
+        elif mn.kind in (CONSTANT_RTQNameL, CONSTANT_RTQNameLA):
+            pass
+        elif mn.kind in (CONSTANT_Multiname, CONSTANT_MultinameA):
+            out += write_u30(mn.name)
+            out += write_u30(mn.ns_set)
+        elif mn.kind in (CONSTANT_MultinameL, CONSTANT_MultinameLA):
+            out += write_u30(mn.ns_set)
+        elif mn.kind == CONSTANT_TypeName:
+            out += write_u30(mn.ns)    # base type multiname index
+            out += write_u30(mn.name)  # parameter count
+            out += mn.data             # pre-serialized parameter u30s
+        else:
+            raise SerializeError(
+                f"Unknown multiname kind 0x{mn.kind:02X}")
+
+    # ── Methods ──────────────────────────────────────────────────────────
+
+    out += write_u30(len(abc.methods))
+    for mi in abc.methods:
+        out += write_u30(mi.param_count)
+        out += write_u30(mi.return_type)
+        for pt in mi.param_types:
+            out += write_u30(pt)
+        out += write_u30(mi.name)
+        out += bytes([mi.flags])
+
+        if mi.flags & METHOD_HasOptional:
+            out += write_u30(len(mi.options))
+            for val, vkind in mi.options:
+                out += write_u30(val)
+                out += bytes([vkind])
+
+        if mi.flags & METHOD_HasParamNames:
+            for pn in mi.param_names:
+                out += write_u30(pn)
+
+    # ── Metadata ─────────────────────────────────────────────────────────
+
+    out += write_u30(len(abc.metadata))
+    for md in abc.metadata:
+        out += write_u30(md.name)
+        out += write_u30(len(md.items))
+        for k, v in md.items:
+            out += write_u30(k)
+            out += write_u30(v)
+
+    # ── Instances + Classes ──────────────────────────────────────────────
+
+    out += write_u30(len(abc.instances))
+    for inst in abc.instances:
+        out += write_u30(inst.name)
+        out += write_u30(inst.super_name)
+        out += bytes([inst.flags])
+        if inst.flags & INSTANCE_ProtectedNs:
+            out += write_u30(inst.protectedNs)
+        out += write_u30(len(inst.interfaces))
+        for ifc in inst.interfaces:
+            out += write_u30(ifc)
+        out += write_u30(inst.iinit)
+        out += _write_traits(inst.traits)
+
+    for ci in abc.classes:
+        out += write_u30(ci.cinit)
+        out += _write_traits(ci.traits)
+
+    # ── Scripts ──────────────────────────────────────────────────────────
+
+    out += write_u30(len(abc.scripts))
+    for si in abc.scripts:
+        out += write_u30(si.init)
+        out += _write_traits(si.traits)
+
+    # ── Method bodies ────────────────────────────────────────────────────
+
+    out += write_u30(len(abc.method_bodies))
+    for mb in abc.method_bodies:
+        out += write_u30(mb.method)
+        out += write_u30(mb.max_stack)
+        out += write_u30(mb.local_count)
+        out += write_u30(mb.init_scope_depth)
+        out += write_u30(mb.max_scope_depth)
+        out += write_u30(len(mb.code))
+        out += mb.code
+        out += write_u30(len(mb.exceptions))
+        for ei in mb.exceptions:
+            out += write_u30(ei.from_offset)
+            out += write_u30(ei.to_offset)
+            out += write_u30(ei.target)
+            out += write_u30(ei.exc_type)
+            out += write_u30(ei.var_name)
+        out += _write_traits(mb.traits)
+
+    return bytes(out)

flashkit/analysis/__init__.py

@@ -0,0 +1,28 @@
+"""
+Analysis services for ABC content.
+
+This package provides graph-based and index-based analysis of the loaded
+ABC bytecode. Each module builds a specific data structure from the parsed
+ABC data that enables efficient queries.
+
+Modules:
+    inheritance: InheritanceGraph — class hierarchy (parent/child/interface).
+    call_graph: CallGraph — method-to-method call edges from bytecode.
+    references: ReferenceIndex — cross-references (field types, instantiations, imports).
+    strings: StringIndex — string constant search and classification.
+"""
+
+from .inheritance import InheritanceGraph
+from .call_graph import CallGraph, CallEdge
+from .references import ReferenceIndex, Reference
+from .strings import StringIndex, StringUsage
+
+__all__ = [
+    "InheritanceGraph",
+    "CallGraph",
+    "CallEdge",
+    "ReferenceIndex",
+    "Reference",
+    "StringIndex",
+    "StringUsage",
+]

flashkit/analysis/call_graph.py

@@ -0,0 +1,317 @@
+"""
+Call graph extraction from method body bytecode.
+
+Scans AVM2 bytecode instructions in MethodBodyInfo.code to build a
+graph of method-to-method call edges. Each edge records the calling
+method, target multiname, opcode type, and bytecode offset.
+
+Usage::
+
+    from flashkit.workspace import Workspace
+    from flashkit.analysis.call_graph import CallGraph
+
+    ws = Workspace()
+    ws.load_swf("application.swf")
+    graph = CallGraph.from_workspace(ws)
+
+    callers = graph.get_callers("doSomething")
+    callees = graph.get_callees("MyClass.init")
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from collections import defaultdict
+
+from ..abc.types import AbcFile
+from ..abc.disasm import decode_instructions
+from ..abc.constants import (
+    OP_callproperty, OP_callpropvoid, OP_constructprop,
+    OP_getproperty, OP_setproperty, OP_initproperty,
+    OP_getlex, OP_findpropstrict, OP_newclass,
+)
+from ..info.member_info import resolve_multiname, build_method_body_map
+from ..info.class_info import ClassInfo
+
+
+# Opcode categories for edges
+CALL_OPS = {OP_callproperty, OP_callpropvoid}
+CONSTRUCT_OPS = {OP_constructprop}
+PROPERTY_READ_OPS = {OP_getproperty, OP_getlex, OP_findpropstrict}
+PROPERTY_WRITE_OPS = {OP_setproperty, OP_initproperty}
+CLASS_OPS = {OP_newclass}
+
+# All opcodes that reference a multiname in their first operand
+_MULTINAME_OPS = (
+    CALL_OPS | CONSTRUCT_OPS | PROPERTY_READ_OPS
+    | PROPERTY_WRITE_OPS | CLASS_OPS
+)
+
+
+@dataclass
+class CallEdge:
+    """A single call/reference edge in the graph.
+
+    Attributes:
+        caller: Qualified name of the calling method (``"Class.method"``).
+        caller_method_index: Method index in the AbcFile.
+        target: Target multiname string (the called/referenced name).
+        opcode: The opcode that generated this edge.
+        mnemonic: Human-readable opcode name.
+        offset: Bytecode offset within the method body.
+        edge_type: Category string: ``"call"``, ``"construct"``, ``"read"``,
+                   ``"write"``, or ``"class"``.
+    """
+    caller: str
+    caller_method_index: int
+    target: str
+    opcode: int
+    mnemonic: str
+    offset: int
+    edge_type: str
+
+
+def _classify_op(opcode: int) -> str:
+    """Classify an opcode into an edge type string."""
+    if opcode in CALL_OPS:
+        return "call"
+    elif opcode in CONSTRUCT_OPS:
+        return "construct"
+    elif opcode in PROPERTY_READ_OPS:
+        return "read"
+    elif opcode in PROPERTY_WRITE_OPS:
+        return "write"
+    elif opcode in CLASS_OPS:
+        return "class"
+    return "unknown"
+
+
+@dataclass
+class CallGraph:
+    """Graph of method call and reference edges extracted from bytecode.
+
+    Attributes:
+        edges: All edges in the graph.
+        callers_index: Map of target name → list of edges calling it.
+        callees_index: Map of caller name → list of edges it produces.
+    """
+    edges: list[CallEdge] = field(default_factory=list)
+    callers_index: dict[str, list[CallEdge]] = field(
+        default_factory=lambda: defaultdict(list))
+    callees_index: dict[str, list[CallEdge]] = field(
+        default_factory=lambda: defaultdict(list))
+
+    @classmethod
+    def from_workspace(cls, workspace: object) -> CallGraph:
+        """Build a CallGraph from a Workspace.
+
+        Iterates all ABC blocks and their method bodies, decodes
+        instructions, and collects edges for multiname-referencing opcodes.
+
+        Args:
+            workspace: A Workspace instance (with .abc_blocks and .classes).
+
+        Returns:
+            Populated CallGraph.
+        """
+        graph = cls()
+
+        # Build a map from method_index → class_name.method_name
+        # for all classes in the workspace
+        from ..workspace.workspace import Workspace
+        ws: Workspace = workspace  # type: ignore[assignment]
+
+        for abc in ws.abc_blocks:
+            method_name_map = _build_method_name_map(abc, ws.classes)
+            method_body_map = build_method_body_map(abc)
+
+            for body in abc.method_bodies:
+                caller_name = method_name_map.get(
+                    body.method, f"method_{body.method}")
+
+                try:
+                    instructions = decode_instructions(body.code)
+                except Exception:
+                    continue
+
+                for instr in instructions:
+                    if instr.opcode in _MULTINAME_OPS and instr.operands:
+                        mn_index = instr.operands[0]
+                        target = resolve_multiname(abc, mn_index)
+                        if target == "*" or target.startswith("multiname["):
+                            continue
+
+                        edge = CallEdge(
+                            caller=caller_name,
+                            caller_method_index=body.method,
+                            target=target,
+                            opcode=instr.opcode,
+                            mnemonic=instr.mnemonic,
+                            offset=instr.offset,
+                            edge_type=_classify_op(instr.opcode),
+                        )
+                        graph.edges.append(edge)
+                        graph.callers_index[target].append(edge)
+                        graph.callees_index[caller_name].append(edge)
+
+        return graph
+
+    @classmethod
+    def from_abc(cls, abc: AbcFile,
+                 classes: list[ClassInfo] | None = None) -> CallGraph:
+        """Build a CallGraph from a single AbcFile.
+
+        Args:
+            abc: The AbcFile to analyze.
+            classes: Optional class list for method name resolution.
+
+        Returns:
+            Populated CallGraph.
+        """
+        graph = cls()
+        method_name_map = _build_method_name_map(abc, classes or [])
+
+        for body in abc.method_bodies:
+            caller_name = method_name_map.get(
+                body.method, f"method_{body.method}")
+
+            try:
+                instructions = decode_instructions(body.code)
+            except Exception:
+                continue
+
+            for instr in instructions:
+                if instr.opcode in _MULTINAME_OPS and instr.operands:
+                    mn_index = instr.operands[0]
+                    target = resolve_multiname(abc, mn_index)
+                    if target == "*" or target.startswith("multiname["):
+                        continue
+
+                    edge = CallEdge(
+                        caller=caller_name,
+                        caller_method_index=body.method,
+                        target=target,
+                        opcode=instr.opcode,
+                        mnemonic=instr.mnemonic,
+                        offset=instr.offset,
+                        edge_type=_classify_op(instr.opcode),
+                    )
+                    graph.edges.append(edge)
+                    graph.callers_index[target].append(edge)
+                    graph.callees_index[caller_name].append(edge)
+
+        return graph
+
+    def get_callers(self, target: str) -> list[CallEdge]:
+        """Get all edges where *target* is called or referenced.
+
+        Args:
+            target: Target method/property name.
+
+        Returns:
+            List of CallEdge objects referencing this target.
+        """
+        return self.callers_index.get(target, [])
+
+    def get_callees(self, caller: str) -> list[CallEdge]:
+        """Get all edges originating from *caller*.
+
+        Args:
+            caller: Caller method name (``"Class.method"`` format).
+
+        Returns:
+            List of CallEdge objects from this caller.
+        """
+        return self.callees_index.get(caller, [])
+
+    def get_callers_by_type(self, target: str,
+                            edge_type: str) -> list[CallEdge]:
+        """Get edges of a specific type referencing *target*.
+
+        Args:
+            target: Target name.
+            edge_type: One of ``"call"``, ``"construct"``, ``"read"``,
+                       ``"write"``, ``"class"``.
+
+        Returns:
+            Filtered list of CallEdge objects.
+        """
+        return [e for e in self.get_callers(target)
+                if e.edge_type == edge_type]
+
+    def get_instantiators(self, class_name: str) -> list[str]:
+        """Get unique caller names that construct instances of *class_name*.
+
+        Args:
+            class_name: The class being instantiated.
+
+        Returns:
+            Sorted list of unique caller names.
+        """
+        edges = self.get_callers_by_type(class_name, "construct")
+        return sorted(set(e.caller for e in edges))
+
+    def get_unique_callers(self, target: str) -> list[str]:
+        """Get unique caller names for a target (calls only).
+
+        Args:
+            target: Target method name.
+
+        Returns:
+            Sorted list of unique caller names.
+        """
+        edges = self.get_callers_by_type(target, "call")
+        return sorted(set(e.caller for e in edges))
+
+    def get_unique_callees(self, caller: str) -> list[str]:
+        """Get unique target names called by a caller (calls only).
+
+        Args:
+            caller: Caller method name.
+
+        Returns:
+            Sorted list of unique target names.
+        """
+        edges = [e for e in self.get_callees(caller) if e.edge_type == "call"]
+        return sorted(set(e.target for e in edges))
+
+    @property
+    def edge_count(self) -> int:
+        return len(self.edges)
+
+    @property
+    def unique_targets(self) -> int:
+        return len(self.callers_index)
+
+    @property
+    def unique_callers(self) -> int:
+        return len(self.callees_index)
+
+
+def _build_method_name_map(abc: AbcFile,
+                           classes: list[ClassInfo]) -> dict[int, str]:
+    """Build a mapping from method_index → ``"ClassName.methodName"``.
+
+    Uses ClassInfo's resolved methods to build readable names.
+    Falls back to ``"method_N"`` for methods not attached to classes.
+    """
+    name_map: dict[int, str] = {}
+
+    for ci in classes:
+        # Instance constructor
+        name_map[ci.constructor_index] = f"{ci.name}.<init>"
+        # Static initializer
+        name_map[ci.static_init_index] = f"{ci.name}.<cinit>"
+
+        for m in ci.methods:
+            prefix = ""
+            if m.is_getter:
+                prefix = "get "
+            elif m.is_setter:
+                prefix = "set "
+            name_map[m.method_index] = f"{ci.name}.{prefix}{m.name}"
+
+        for m in ci.static_methods:
+            name_map[m.method_index] = f"{ci.name}.static {m.name}"
+
+    return name_map