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