#!/usr/bin/env python3
"""Find static GPS/map route candidates in a Cyberpunk 2077 executable."""

from __future__ import annotations

import argparse
import dataclasses
import json
import re
import struct
from pathlib import Path


@dataclasses.dataclass(frozen=True)
class Section:
    name: str
    virtual_address: int
    virtual_size: int
    raw_offset: int
    raw_size: int
    characteristics: int

    @property
    def is_code(self) -> bool:
        return bool(self.characteristics & 0x20) and bool(self.characteristics & 0x20000000)

    @property
    def is_readable_data(self) -> bool:
        return bool(self.characteristics & 0x40000000) and not self.is_code

    def contains_rva(self, rva: int) -> bool:
        return self.virtual_address <= rva < self.virtual_address + max(self.virtual_size, self.raw_size)

    def rva_to_offset(self, rva: int) -> int:
        return self.raw_offset + (rva - self.virtual_address)


def parse_pe(data: bytes) -> tuple[int, list[Section]]:
    if data[:2] != b"MZ":
        raise ValueError("not an MZ executable")

    pe_offset = struct.unpack_from("<I", data, 0x3C)[0]
    if data[pe_offset : pe_offset + 4] != b"PE\0\0":
        raise ValueError("not a PE executable")

    coff_offset = pe_offset + 4
    section_count = struct.unpack_from("<H", data, coff_offset + 2)[0]
    optional_size = struct.unpack_from("<H", data, coff_offset + 16)[0]
    optional_offset = coff_offset + 20
    if struct.unpack_from("<H", data, optional_offset)[0] != 0x20B:
        raise ValueError("expected PE32+ executable")

    image_base = struct.unpack_from("<Q", data, optional_offset + 24)[0]
    section_offset = optional_offset + optional_size
    sections: list[Section] = []
    for index in range(section_count):
        offset = section_offset + index * 40
        name = data[offset : offset + 8].split(b"\0", 1)[0].decode("ascii", "replace")
        virtual_size, virtual_address, raw_size, raw_offset = struct.unpack_from("<IIII", data, offset + 8)
        characteristics = struct.unpack_from("<I", data, offset + 36)[0]
        sections.append(Section(name, virtual_address, virtual_size, raw_offset, raw_size, characteristics))

    return image_base, sections


def iter_strings(data: bytes, section: Section, min_len: int) -> list[tuple[int, str, str]]:
    results: list[tuple[int, str, str]] = []
    start = section.raw_offset
    end = min(len(data), section.raw_offset + section.raw_size)
    blob = data[start:end]

    ascii_re = re.compile(rb"[\x20-\x7e]{%d,}" % min_len)
    for match in ascii_re.finditer(blob):
        text = match.group(0).decode("ascii", "replace")
        results.append((section.virtual_address + match.start(), "ascii", text))

    utf16_re = re.compile((rb"(?:[\x20-\x7e]\x00){%d,}" % min_len))
    for match in utf16_re.finditer(blob):
        raw = match.group(0)
        text = raw.decode("utf-16le", "replace")
        results.append((section.virtual_address + match.start(), "utf16", text))

    return results


def find_xrefs(data: bytes, sections: list[Section], image_base: int, target_rvas: set[int]) -> dict[int, list[int]]:
    target_vas = {image_base + rva: rva for rva in target_rvas}
    hits = {rva: [] for rva in target_rvas}
    for section in sections:
        if not section.is_code:
            continue

        start = section.raw_offset
        end = min(len(data), section.raw_offset + section.raw_size)
        for offset in range(start, max(start, end - 4)):
            disp = struct.unpack_from("<i", data, offset)[0]
            instr_end_va = image_base + section.virtual_address + (offset - section.raw_offset) + 4
            target_va = instr_end_va + disp
            target_rva = target_vas.get(target_va)
            if target_rva is not None:
                hits[target_rva].append(section.virtual_address + (offset - section.raw_offset))
    return hits


def load_address_symbols(path: Path, pattern: re.Pattern[str]) -> list[tuple[int, str]]:
    payload = json.loads(path.read_text())
    symbols: list[tuple[int, str]] = []
    for item in payload.get("Addresses", []):
        symbol = item.get("symbol")
        offset = item.get("offset")
        if not symbol or not offset or not pattern.search(symbol):
            continue

        section, _, value = offset.partition(":")
        if section != "0001":
            continue

        symbols.append((int(value, 16), symbol))
    return sorted(symbols)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("exe", type=Path)
    parser.add_argument("--addresses", type=Path)
    parser.add_argument(
        "--pattern",
        default=r"gps|route|path|mappin|track|quest|journal|navigation|map",
        help="case-insensitive regex for string and symbol filtering",
    )
    parser.add_argument("--min-len", type=int, default=5)
    parser.add_argument("--max-strings", type=int, default=300)
    args = parser.parse_args()

    data = args.exe.read_bytes()
    image_base, sections = parse_pe(data)
    pattern = re.compile(args.pattern, re.IGNORECASE)

    matches: list[tuple[int, str, str]] = []
    for section in sections:
        if section.is_readable_data:
            matches.extend(item for item in iter_strings(data, section, args.min_len) if pattern.search(item[2]))

    matches.sort(key=lambda item: (item[0], item[1], item[2]))
    target_rvas = {rva for rva, _, _ in matches}
    xrefs = find_xrefs(data, sections, image_base, target_rvas)

    print(f"image_base=0x{image_base:x}")
    print(f"string_matches={len(matches)} pattern={args.pattern!r}")
    for rva, encoding, text in matches[: args.max_strings]:
        hit_text = ", ".join(f"0x{hit:x}" for hit in xrefs.get(rva, [])[:12])
        if len(xrefs.get(rva, [])) > 12:
            hit_text += f", ... +{len(xrefs[rva]) - 12}"
        print(f"string rva=0x{rva:x} enc={encoding} xrefs={len(xrefs.get(rva, []))} text={text!r}")
        if hit_text:
            print(f"  xref_disp_rvas={hit_text}")

    if args.addresses:
        print()
        symbols = load_address_symbols(args.addresses, pattern)
        print(f"address_symbols={len(symbols)}")
        for rva, symbol in symbols[: args.max_strings]:
            print(f"symbol rva=0x{rva:x} name={symbol}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())