#!/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())