Add PE value reference helper

tools/find_pe_value_refs.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""Find PE code blocks containing selected immediate byte values.
+
+This is a lightweight triage helper, not a real disassembler. It groups .text
+bytes into contiguous blocks split by int3 padding and reports blocks whose raw
+bytes contain values of interest. It is useful when looking for native code that
+touches compact flags such as worldTrafficLanePersistentFlags.
+"""
+
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import re
+from pathlib import Path
+
+from find_pe_string_xrefs import Section, parse_pe
+
+
+DEFAULT_PATTERNS = {
+    "flag_road_u16": (0x0010, 2),
+    "flag_intersection_u16": (0x0020, 2),
+    "flag_traffic_disabled_u16": (0x0080, 2),
+    "flag_gpsonly_u16": (0x0200, 2),
+    "flag_no_ai_driving_u16": (0x2000, 2),
+    "flag_highway_u16": (0x4000, 2),
+    "flag_no_autodrive_u16": (0x8000, 2),
+    "flag_road_u32": (0x0010, 4),
+    "flag_intersection_u32": (0x0020, 4),
+    "flag_traffic_disabled_u32": (0x0080, 4),
+    "flag_gpsonly_u32": (0x0200, 4),
+    "flag_no_ai_driving_u32": (0x2000, 4),
+    "flag_highway_u32": (0x4000, 4),
+    "flag_no_autodrive_u32": (0x8000, 4),
+    "lane_size_u32": (0x00A0, 4),
+    "traffic_data_size_u32": (0x0110, 4),
+}
+
+
+@dataclasses.dataclass(frozen=True)
+class Pattern:
+    name: str
+    needle: bytes
+
+
+@dataclasses.dataclass(frozen=True)
+class Hit:
+    rva: int
+    pattern: str
+
+
+def parse_pattern(text: str) -> Pattern:
+    name, sep, spec = text.partition("=")
+    if not sep or not name:
+        raise argparse.ArgumentTypeError("patterns must be name=hex[:size]")
+
+    value_text, _, size_text = spec.partition(":")
+    value = int(value_text, 0)
+    size = int(size_text, 0) if size_text else 4
+    if size not in (1, 2, 4, 8):
+        raise argparse.ArgumentTypeError("pattern size must be 1, 2, 4, or 8")
+    if value < 0 or value >= (1 << (size * 8)):
+        raise argparse.ArgumentTypeError("pattern value does not fit size")
+
+    return Pattern(name=name, needle=value.to_bytes(size, "little"))
+
+
+def default_patterns() -> list[Pattern]:
+    return [Pattern(name, value.to_bytes(size, "little")) for name, (value, size) in DEFAULT_PATTERNS.items()]
+
+
+def iter_code_blocks(data: bytes, section: Section, min_size: int) -> list[tuple[int, bytes]]:
+    start = section.raw_offset
+    end = min(len(data), section.raw_offset + section.raw_size)
+    blob = data[start:end]
+    blocks: list[tuple[int, bytes]] = []
+    block_start = 0
+
+    for match in re.finditer(rb"\xcc{2,}", blob):
+        if match.start() - block_start >= min_size:
+            blocks.append((section.virtual_address + block_start, blob[block_start : match.start()]))
+        block_start = match.end()
+
+    if len(blob) - block_start >= min_size:
+        blocks.append((section.virtual_address + block_start, blob[block_start:]))
+
+    return blocks
+
+
+def find_hits(block_rva: int, block: bytes, patterns: list[Pattern], max_offsets: int) -> tuple[dict[str, int], list[Hit]]:
+    counts: dict[str, int] = {}
+    hits: list[Hit] = []
+    for pattern in patterns:
+        offset = block.find(pattern.needle)
+        while offset != -1:
+            counts[pattern.name] = counts.get(pattern.name, 0) + 1
+            if len([hit for hit in hits if hit.pattern == pattern.name]) < max_offsets:
+                hits.append(Hit(block_rva + offset, pattern.name))
+            offset = block.find(pattern.needle, offset + 1)
+    return counts, hits
+
+
+def score_counts(counts: dict[str, int]) -> int:
+    score = 0
+    for name, count in counts.items():
+        weight = 1
+        if name.endswith("_u32"):
+            weight = 2
+        if name in {"flag_highway_u16", "flag_highway_u32", "flag_gpsonly_u16", "flag_gpsonly_u32"}:
+            weight += 2
+        if name in {"lane_size_u32", "traffic_data_size_u32"}:
+            weight += 3
+        score += min(count, 8) * weight
+
+    names = set(counts)
+    if names & {"flag_highway_u16", "flag_highway_u32"} and names & {"flag_gpsonly_u16", "flag_gpsonly_u32"}:
+        score += 8
+    if names & {"flag_road_u16", "flag_road_u32"} and names & {"flag_highway_u16", "flag_highway_u32"}:
+        score += 4
+    return score
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pe", type=Path)
+    parser.add_argument("--pattern", action="append", type=parse_pattern, default=[])
+    parser.add_argument("--no-defaults", action="store_true")
+    parser.add_argument("--require", action="append", default=[], help="Require a pattern name; repeatable")
+    parser.add_argument("--min-score", type=int, default=10)
+    parser.add_argument("--min-block-size", type=int, default=16)
+    parser.add_argument("--max-results", type=int, default=80)
+    parser.add_argument("--max-offsets", type=int, default=4)
+    args = parser.parse_args()
+
+    patterns = ([] if args.no_defaults else default_patterns()) + args.pattern
+    data = args.pe.read_bytes()
+    image_base, sections = parse_pe(data)
+
+    results: list[tuple[int, int, int, dict[str, int], list[Hit]]] = []
+    for section in sections:
+        if not section.is_code:
+            continue
+        for block_rva, block in iter_code_blocks(data, section, args.min_block_size):
+            counts, hits = find_hits(block_rva, block, patterns, args.max_offsets)
+            if not counts:
+                continue
+            if any(required not in counts for required in args.require):
+                continue
+            score = score_counts(counts)
+            if score < args.min_score:
+                continue
+            results.append((score, block_rva, block_rva + len(block), counts, hits))
+
+    results.sort(key=lambda item: (-item[0], item[1]))
+    print(f"image_base=0x{image_base:x}")
+    print(f"blocks={len(results)} min_score={args.min_score} require={args.require}")
+    for score, start_rva, end_rva, counts, hits in results[: args.max_results]:
+        count_text = ", ".join(f"{name}:{counts[name]}" for name in sorted(counts))
+        print(f"block rva=0x{start_rva:x}..0x{end_rva:x} size=0x{end_rva - start_rva:x} score={score}")
+        print(f"  counts {count_text}")
+        for hit in sorted(hits, key=lambda item: item.rva):
+            print(f"  hit rva=0x{hit.rva:x} pattern={hit.pattern}")
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())