Build RE baseline and initial Rust workspace

tools/py/rt3_rekitlib.py

@@ -0,0 +1,543 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import csv
+import json
+import re
+import subprocess
+from bisect import bisect_right
+from functools import lru_cache
+from pathlib import Path
+
+
+PENDING_TEMPLATE_STORE_DEFAULT_SEEDS = [
+    0x0059B2E0,
+    0x0059B710,
+    0x0059B740,
+    0x0059C470,
+    0x0059C540,
+    0x0059C590,
+    0x0059C5B0,
+    0x0059C5E0,
+    0x0059C5F0,
+]
+
+PENDING_TEMPLATE_STORE_ADJACENT_MIN = 0x0059B000
+PENDING_TEMPLATE_STORE_ADJACENT_MAX = 0x0059D000
+DESTRUCTOR_SWITCH_ADDR = 0x0059B2E0
+HEAP_FREE_ADDR = 0x0058F3C0
+
+
+def parse_hex(text: str) -> int:
+    value = text.strip().lower()
+    if value.startswith("0x"):
+        value = value[2:]
+    return int(value, 16)
+
+
+def fmt_addr(value: int) -> str:
+    return f"0x{value:08x}"
+
+
+def display_string(text: str) -> str:
+    return text.encode("unicode_escape").decode("ascii")
+
+
+def clean_json_payload(text: str) -> str:
+    stripped = text.strip()
+    if not stripped:
+        raise ValueError("rizin returned empty output")
+    starts = [index for index in (stripped.find("["), stripped.find("{")) if index >= 0]
+    if not starts:
+        raise ValueError("rizin did not return JSON")
+    return stripped[min(starts) :]
+
+
+def run_rizin_json(exe_path: Path, command: str) -> object:
+    result = subprocess.run(
+        [
+            "rizin",
+            "-q",
+            "-e",
+            "scr.color=false",
+            "-c",
+            command,
+            str(exe_path),
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return json.loads(clean_json_payload(result.stdout))
+
+
+def run_objdump_excerpt(exe_path: Path, address: int, radius: int = 0x20) -> str:
+    start = max(address - radius, 0)
+    stop = address + radius
+    result = subprocess.run(
+        [
+            "llvm-objdump",
+            "-d",
+            "--no-show-raw-insn",
+            f"--start-address={fmt_addr(start)}",
+            f"--stop-address={fmt_addr(stop)}",
+            str(exe_path),
+        ],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    lines = [
+        line.rstrip()
+        for line in result.stdout.splitlines()
+        if re.match(r"^\s*[0-9a-fA-F]+:", line)
+    ]
+    return "\n".join(lines)
+
+
+def load_curated_rows(path: Path) -> dict[int, dict[str, str]]:
+    if not path.exists():
+        return {}
+    with path.open(newline="", encoding="utf-8") as handle:
+        rows = csv.DictReader(handle)
+        return {parse_hex(row["address"]): dict(row) for row in rows}
+
+
+class FunctionIndex:
+    def __init__(self, rows: list[dict[str, object]], curated_names: dict[int, str]):
+        self.rows = sorted(rows, key=lambda row: int(row["offset"]))
+        self.by_start = {int(row["offset"]): row for row in self.rows}
+        self.starts = [int(row["offset"]) for row in self.rows]
+        self.curated_names = curated_names
+
+    def get_exact(self, address: int) -> dict[str, object] | None:
+        return self.by_start.get(address)
+
+    def find_containing(self, address: int) -> dict[str, object] | None:
+        index = bisect_right(self.starts, address) - 1
+        if index < 0:
+            return None
+        row = self.rows[index]
+        start = int(row["offset"])
+        end = int(row.get("maxbound", start + int(row.get("size", 0))))
+        if start <= address < end:
+            return row
+        return None
+
+    def preferred_name(self, row: dict[str, object]) -> str:
+        start = int(row["offset"])
+        return self.curated_names.get(start, str(row["name"]))
+
+
+class BranchAnalyzer:
+    def __init__(self, exe_path: Path, output_dir: Path):
+        self.exe_path = exe_path.resolve()
+        self.output_dir = output_dir.resolve()
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        curated_map = self.output_dir / "function-map.csv"
+        self.curated_rows = load_curated_rows(curated_map)
+        self.curated_names = {
+            address: row["name"] for address, row in self.curated_rows.items()
+        }
+        self.function_index = FunctionIndex(
+            self._load_function_rows(),
+            self.curated_names,
+        )
+        self.strings = list(run_rizin_json(self.exe_path, "izzj"))
+        self.strings_by_addr = {int(entry["vaddr"]): entry for entry in self.strings}
+
+    def _load_function_rows(self) -> list[dict[str, object]]:
+        rows = list(run_rizin_json(self.exe_path, "aaa; aflj"))
+        known_starts = {int(row["offset"]) for row in rows}
+        missing_curated = sorted(
+            address for address in self.curated_names if address not in known_starts
+        )
+        if not missing_curated:
+            return rows
+
+        define_cmd = "aaa; " + "; ".join(
+            f"af @ {fmt_addr(address)}" for address in missing_curated
+        ) + "; aflj"
+        return list(run_rizin_json(self.exe_path, define_cmd))
+
+    @lru_cache(maxsize=None)
+    def xrefs_to(self, address: int) -> list[dict[str, object]]:
+        return list(run_rizin_json(self.exe_path, f"aaa; axtj @ {fmt_addr(address)}"))
+
+    @lru_cache(maxsize=None)
+    def function_pdfj(self, address: int) -> dict[str, object]:
+        payload = run_rizin_json(self.exe_path, f"aaa; s {fmt_addr(address)}; pdfj")
+        if not isinstance(payload, dict):
+            raise TypeError(f"unexpected pdfj payload for {fmt_addr(address)}")
+        return payload
+
+    @lru_cache(maxsize=None)
+    def excerpt(self, address: int) -> str:
+        return run_objdump_excerpt(self.exe_path, address)
+
+    def fallback_function(self, address: int) -> dict[str, object] | None:
+        curated = self.curated_rows.get(address)
+        if curated is None:
+            return None
+
+        size = int(curated["size"])
+        return {
+            "offset": address,
+            "name": curated["name"],
+            "size": size,
+            "maxbound": address + size,
+            "calltype": curated["calling_convention"],
+            "signature": "",
+            "codexrefs": self.xrefs_to(address),
+            "callrefs": [],
+            "datarefs": [],
+            "synthetic": True,
+        }
+
+    def resolve_target_function(self, address: int) -> dict[str, object] | None:
+        exact = self.function_index.get_exact(address)
+        if exact is not None:
+            return exact
+
+        fallback = self.fallback_function(address)
+        if fallback is not None:
+            return fallback
+
+        return self.function_index.find_containing(address)
+
+    def format_callers(self, row: dict[str, object]) -> list[dict[str, object]]:
+        callers: list[dict[str, object]] = []
+        for ref in row.get("codexrefs", []):
+            if ref.get("type") != "CALL":
+                continue
+            call_site = int(ref["from"])
+            caller = self.function_index.find_containing(call_site)
+            callers.append({"call_site": call_site, "function": caller})
+        callers.sort(key=lambda entry: entry["call_site"])
+        return callers
+
+    def format_callees(self, row: dict[str, object]) -> list[dict[str, object]]:
+        callees: list[dict[str, object]] = []
+        seen: set[tuple[int, int]] = set()
+        for ref in row.get("callrefs", []):
+            if ref.get("type") != "CALL":
+                continue
+            call_site = int(ref["from"])
+            callee_site = int(ref["to"])
+            callee = self.function_index.find_containing(callee_site)
+            if callee is None:
+                continue
+            key = (call_site, int(callee["offset"]))
+            if key in seen:
+                continue
+            seen.add(key)
+            callees.append({"call_site": call_site, "function": callee})
+        callees.sort(key=lambda entry: (int(entry["function"]["offset"]), entry["call_site"]))
+        return callees
+
+    def format_data_refs(self, row: dict[str, object]) -> list[dict[str, object]]:
+        refs: list[dict[str, object]] = []
+        seen: set[tuple[int, int, str]] = set()
+        for ref in row.get("datarefs", []):
+            from_addr = int(ref["from"])
+            to_addr = int(ref["to"])
+            ref_type = str(ref.get("type", "DATA"))
+            key = (from_addr, to_addr, ref_type)
+            if key in seen:
+                continue
+            seen.add(key)
+            refs.append(
+                {
+                    "from": from_addr,
+                    "to": to_addr,
+                    "type": ref_type,
+                    "string": self.strings_by_addr.get(to_addr),
+                }
+            )
+        refs.sort(key=lambda entry: (entry["to"], entry["from"]))
+        return refs
+
+    def describe_caller(self, call_site: int, function: dict[str, object] | None) -> str:
+        if function is None:
+            return fmt_addr(call_site)
+        return (
+            f"{fmt_addr(call_site)}@{fmt_addr(int(function['offset']))}:"
+            f"{self.function_index.preferred_name(function)}"
+        )
+
+    def describe_callee(self, call_site: int, function: dict[str, object] | None) -> str:
+        if function is None:
+            return fmt_addr(call_site)
+        return (
+            f"{fmt_addr(call_site)}->{fmt_addr(int(function['offset']))}:"
+            f"{self.function_index.preferred_name(function)}"
+        )
+
+    def describe_data_ref(self, entry: dict[str, object]) -> str:
+        target = fmt_addr(int(entry["to"]))
+        string_entry = entry["string"]
+        if string_entry is not None:
+            target += f':"{display_string(str(string_entry["string"]))}"'
+        return f"{fmt_addr(int(entry['from']))}->{target}"
+
+    def collect_key_constants(self, pdfj: dict[str, object]) -> list[str]:
+        values: list[int] = []
+        seen: set[int] = set()
+        for op in pdfj.get("ops", []):
+            for key in ("val", "ptr"):
+                raw = op.get(key)
+                if not isinstance(raw, int):
+                    continue
+                if raw < 0x10 or raw > 0x100000:
+                    continue
+                if 0x00400000 <= raw <= 0x01000000:
+                    continue
+                if raw in seen:
+                    continue
+                seen.add(raw)
+                values.append(raw)
+        return [fmt_addr(value) for value in values[:10]]
+
+    def collect_key_strings(self, data_refs: list[dict[str, object]]) -> list[str]:
+        strings: list[str] = []
+        seen: set[str] = set()
+        for entry in data_refs:
+            string_entry = entry["string"]
+            if string_entry is None:
+                continue
+            text = display_string(str(string_entry["string"]))
+            if text in seen:
+                continue
+            seen.add(text)
+            strings.append(text)
+        return strings[:8]
+
+    def discover_adjacent_functions(self, addresses: list[int]) -> list[int]:
+        discovered = set(addresses)
+        for address in addresses:
+            function = self.resolve_target_function(address)
+            if function is None:
+                continue
+            for entry in self.format_callers(function):
+                caller = entry["function"]
+                if caller is None:
+                    continue
+                caller_start = int(caller["offset"])
+                if PENDING_TEMPLATE_STORE_ADJACENT_MIN <= caller_start < PENDING_TEMPLATE_STORE_ADJACENT_MAX:
+                    discovered.add(caller_start)
+            for entry in self.format_callees(function):
+                callee = entry["function"]
+                callee_start = int(callee["offset"])
+                if PENDING_TEMPLATE_STORE_ADJACENT_MIN <= callee_start < PENDING_TEMPLATE_STORE_ADJACENT_MAX:
+                    discovered.add(callee_start)
+        return sorted(discovered)
+
+    def build_function_rows(self, addresses: list[int]) -> list[dict[str, str]]:
+        rows: list[dict[str, str]] = []
+        for query_address in self.discover_adjacent_functions(addresses):
+            function = self.resolve_target_function(query_address)
+            if function is None:
+                continue
+            callers = self.format_callers(function)
+            callees = self.format_callees(function)
+            data_refs = self.format_data_refs(function)
+            pdfj = self.function_pdfj(int(function["offset"]))
+            rows.append(
+                {
+                    "query_address": fmt_addr(query_address),
+                    "function_address": fmt_addr(int(function["offset"])),
+                    "name": self.function_index.preferred_name(function),
+                    "size": str(function["size"]),
+                    "calling_convention": str(function.get("calltype", "unknown")),
+                    "caller_count": str(len(callers)),
+                    "callers": "; ".join(
+                        self.describe_caller(entry["call_site"], entry["function"])
+                        for entry in callers
+                    ),
+                    "callee_count": str(len(callees)),
+                    "callees": "; ".join(
+                        self.describe_callee(entry["call_site"], entry["function"])
+                        for entry in callees
+                    ),
+                    "data_ref_count": str(len(data_refs)),
+                    "data_refs": "; ".join(self.describe_data_ref(entry) for entry in data_refs),
+                    "key_constants": "; ".join(self.collect_key_constants(pdfj)),
+                    "key_strings": "; ".join(self.collect_key_strings(data_refs)),
+                    "entry_excerpt": self.excerpt(int(function["offset"])).replace("\n", " | "),
+                }
+            )
+        return rows
+
+    def _case_groups(self, switch_address: int) -> list[dict[str, object]]:
+        pdfj = self.function_pdfj(switch_address)
+        groups: list[dict[str, object]] = []
+        current: dict[str, object] | None = None
+        case_pattern = re.compile(r"^case\.0x[0-9a-f]+\.(\d+)$")
+        default_pattern = re.compile(r"^case\.default\.0x[0-9a-f]+$")
+        for op in pdfj.get("ops", []):
+            labels: list[str] = []
+            for flag in op.get("flags", []):
+                match = case_pattern.match(flag)
+                if match:
+                    labels.append(match.group(1))
+                    continue
+                if default_pattern.match(flag):
+                    labels.append("default")
+            if labels:
+                current = {
+                    "start": int(op["offset"]),
+                    "cases": labels,
+                    "ops": [op],
+                }
+                groups.append(current)
+                continue
+            if current is not None:
+                current["ops"].append(op)
+        return groups
+
+    def _infer_case_shape(self, ops: list[dict[str, object]]) -> tuple[str, str, str]:
+        disasm_lines = [str(op["disasm"]) for op in ops]
+        text = "\n".join(disasm_lines)
+        direct_offsets = {
+            int(match.group(1), 16)
+            for match in re.finditer(r"\[(?:edi|eax|ecx)\+0x([0-9a-f]+)\]", text)
+        }
+        indexed_offsets = {
+            int(match.group(1), 16)
+            for match in re.finditer(r"\[(?:edi|eax|ecx)\+0x([0-9a-f]+)\]", text)
+            if "*4" in text
+        }
+        free_calls = sum(
+            1
+            for op in ops
+            if op.get("jump") == HEAP_FREE_ADDR or "fcn.0058f3c0" in str(op.get("disasm", ""))
+        )
+        has_loop = any("*4" in line for line in disasm_lines)
+
+        fields = ["top-level payload"]
+        for offset in sorted(direct_offsets):
+            fields.append(f"payload+0x{offset:02x}")
+        if has_loop:
+            for offset in sorted(indexed_offsets):
+                fields.append(f"vector@payload+0x{offset:02x}")
+
+        unique_fields = []
+        seen_fields: set[str] = set()
+        for field in fields:
+            if field in seen_fields:
+                continue
+            seen_fields.add(field)
+            unique_fields.append(field)
+
+        if has_loop and len(direct_offsets) >= 3:
+            shape = "pointer vectors with paired side tables"
+        elif has_loop:
+            shape = "indexed pointer vector"
+        elif len(direct_offsets) >= 4:
+            shape = "fixed pointer tuple"
+        elif len(direct_offsets) >= 2:
+            shape = "paired nested pointers"
+        elif free_calls <= 1:
+            shape = "top-level payload pointer"
+        else:
+            shape = "single nested pointer"
+
+        cleanup_summary = f"{free_calls} heap free call(s)"
+        excerpt = " | ".join(disasm_lines[:8])
+        return shape, ", ".join(unique_fields), cleanup_summary + f"; {excerpt}"
+
+    def build_record_kind_rows(self) -> list[dict[str, str]]:
+        rows: list[dict[str, str]] = []
+        for group in self._case_groups(DESTRUCTOR_SWITCH_ADDR):
+            shape, freed_fields, notes = self._infer_case_shape(group["ops"])
+            case_group = ",".join(group["cases"])
+            for case_label in group["cases"]:
+                rows.append(
+                    {
+                        "record_kind": case_label,
+                        "case_group": case_group,
+                        "owning_function": fmt_addr(DESTRUCTOR_SWITCH_ADDR),
+                        "owning_name": self.curated_names.get(
+                            DESTRUCTOR_SWITCH_ADDR,
+                            "multiplayer_transport_destroy_pending_template_dispatch_record",
+                        ),
+                        "inferred_payload_shape": shape,
+                        "freed_fields": freed_fields,
+                        "notes": notes,
+                    }
+                )
+        rows.sort(
+            key=lambda row: (
+                row["record_kind"] == "default",
+                int(row["record_kind"]) if row["record_kind"] != "default" else 0,
+            )
+        )
+        return rows
+
+    def write_csv(self, path: Path, rows: list[dict[str, str]]) -> None:
+        if not rows:
+            return
+        with path.open("w", newline="", encoding="utf-8") as handle:
+            writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
+            writer.writeheader()
+            writer.writerows(rows)
+
+    def write_pending_template_store_markdown(self, function_rows: list[dict[str, str]]) -> None:
+        by_address = {parse_hex(row["function_address"]): row for row in function_rows}
+        sections = {
+            "Init": [0x0059B710, 0x0059C5B0],
+            "Destroy": [0x0059B2E0, 0x0059B740, 0x0059C5E0],
+            "Lookup": [0x0059C540, 0x0059C590],
+            "Prune / Remove": [0x0059C470],
+            "Dispatch / Update": [0x0059C220, 0x0059C5F0],
+        }
+
+        lines = [
+            "# Pending-Template Store Management",
+            "",
+            f"- Target binary: `{self.exe_path}`",
+            "- Scope: companion pending-template dispatch store and its adjacent management helpers.",
+            "",
+        ]
+
+        for title, addresses in sections.items():
+            lines.extend([f"## {title}", ""])
+            for address in addresses:
+                row = by_address.get(address)
+                if row is None:
+                    continue
+                lines.append(f"### `{row['function_address']}` `{row['name']}`")
+                lines.append("")
+                lines.append(f"- Size: `{row['size']}`")
+                lines.append(f"- Calling convention: `{row['calling_convention']}`")
+                lines.append(f"- Callers: {row['callers'] or 'none'}")
+                lines.append(f"- Direct callees: {row['callees'] or 'none'}")
+                lines.append(f"- Data refs: {row['data_refs'] or 'none'}")
+                lines.append(f"- Key constants: {row['key_constants'] or 'none'}")
+                lines.append(f"- Key strings: {row['key_strings'] or 'none'}")
+                lines.append("")
+                lines.append("Entry excerpt:")
+                lines.append("")
+                lines.append("```asm")
+                lines.append(self.excerpt(address))
+                lines.append("```")
+                lines.append("")
+
+        (self.output_dir / "pending-template-store-management.md").write_text(
+            "\n".join(lines) + "\n",
+            encoding="utf-8",
+        )
+
+
+def export_pending_template_store(
+    exe_path: Path,
+    output_dir: Path,
+    seed_addresses: list[int],
+) -> None:
+    analyzer = BranchAnalyzer(exe_path, output_dir)
+    function_rows = analyzer.build_function_rows(seed_addresses)
+    record_rows = analyzer.build_record_kind_rows()
+    analyzer.write_csv(output_dir / "pending-template-store-functions.csv", function_rows)
+    analyzer.write_csv(output_dir / "pending-template-store-record-kinds.csv", record_rows)
+    analyzer.write_pending_template_store_markdown(function_rows)