#!/usr/bin/env python3
"""Map REDscript/native registration strings to nearby code pointer candidates."""

from __future__ import annotations

import argparse
import re
import struct
from pathlib import Path

from find_pe_string_xrefs import Section, parse_pe


def section_for_rva(sections: list[Section], rva: int) -> Section | None:
    for section in sections:
        if section.contains_rva(rva):
            return section
    return None


def is_readable_data(section: Section) -> bool:
    return bool(section.characteristics & 0x40000000) and not section.is_code


def iter_ascii_strings(data: bytes, section: Section, min_len: int) -> list[tuple[int, str]]:
    start = section.raw_offset
    end = min(len(data), section.raw_offset + section.raw_size)
    pattern = re.compile(rb"[\x20-\x7e]{%d,}" % min_len)
    return [
        (section.virtual_address + match.start(), match.group(0).decode("ascii", "replace"))
        for match in pattern.finditer(data[start:end])
    ]


def read_i32(data: bytes, offset: int) -> int | None:
    if offset < 0 or offset + 4 > len(data):
        return None
    return struct.unpack_from("<i", data, offset)[0]


def find_rip_disp_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 = read_i32(data, offset)
            if disp is None:
                continue
            instr_end_va = image_base + section.virtual_address + (offset - section.raw_offset) + 4
            target_rva = target_vas.get(instr_end_va + disp)
            if target_rva is not None:
                hits[target_rva].append(section.virtual_address + (offset - section.raw_offset))
    return hits


def iter_nearby_rip_code_targets(
    data: bytes,
    sections: list[Section],
    image_base: int,
    code_section: Section,
    center_disp_rva: int,
    radius: int,
) -> list[tuple[int, int, str]]:
    # Common x64 RIP-relative LEA/MOV forms used by the generated registration
    # thunks to pass function pointers and type descriptors around.
    patterns: tuple[tuple[bytes, int, str], ...] = (
        (b"\x48\x8d\x05", 3, "lea rax"),
        (b"\x48\x8d\x0d", 3, "lea rcx"),
        (b"\x48\x8d\x15", 3, "lea rdx"),
        (b"\x48\x8d\x1d", 3, "lea rbx"),
        (b"\x48\x8d\x35", 3, "lea rsi"),
        (b"\x48\x8d\x3d", 3, "lea rdi"),
        (b"\x4c\x8d\x05", 3, "lea r8"),
        (b"\x4c\x8d\x0d", 3, "lea r9"),
        (b"\x4c\x8d\x15", 3, "lea r10"),
        (b"\x4c\x8d\x1d", 3, "lea r11"),
        (b"\x4c\x8d\x25", 3, "lea r12"),
        (b"\x4c\x8d\x2d", 3, "lea r13"),
        (b"\x4c\x8d\x35", 3, "lea r14"),
        (b"\x4c\x8d\x3d", 3, "lea r15"),
    )

    center_off = code_section.rva_to_offset(center_disp_rva)
    start = max(code_section.raw_offset, center_off - radius)
    end = min(len(data), code_section.raw_offset + code_section.raw_size, center_off + radius)
    out: list[tuple[int, int, str]] = []

    for offset in range(start, end):
        for prefix, disp_start, label in patterns:
            if data[offset : offset + len(prefix)] != prefix:
                continue

            disp = read_i32(data, offset + disp_start)
            if disp is None:
                continue

            instr_rva = code_section.virtual_address + (offset - code_section.raw_offset)
            instr_end_va = image_base + instr_rva + disp_start + 4
            target_rva = instr_end_va + disp - image_base
            target_section = section_for_rva(sections, target_rva)
            if target_section and target_section.is_code:
                out.append((instr_rva, target_rva, label))

    return sorted(set(out))


def likely_registered_code(ref_rva: int, candidates: list[tuple[int, int, str]]) -> tuple[int, int, str] | None:
    for candidate in candidates:
        instr_rva, _target_rva, label = candidate
        if instr_rva > ref_rva and label == "lea rax":
            return candidate

    return None


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("pe", type=Path)
    parser.add_argument("--pattern", required=True, help="case-insensitive regex for registration strings")
    parser.add_argument("--min-len", type=int, default=5)
    parser.add_argument("--radius", type=int, default=192)
    parser.add_argument("--max-strings", type=int, default=200)
    args = parser.parse_args()

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

    strings: list[tuple[int, str]] = []
    for section in sections:
        if is_readable_data(section):
            strings.extend((rva, text) for rva, text in iter_ascii_strings(data, section, args.min_len) if regex.search(text))
    strings = sorted(strings, key=lambda item: (item[1].lower(), item[0]))

    xrefs = find_rip_disp_xrefs(data, sections, image_base, {rva for rva, _ in strings})
    print(f"image_base=0x{image_base:x}")
    print(f"matches={len(strings)} pattern={args.pattern!r}")

    emitted = 0
    for string_rva, text in strings:
        refs = xrefs.get(string_rva, [])
        if not refs:
            continue

        print(f"name={text!r} string_rva=0x{string_rva:x} xrefs={len(refs)}")
        for ref in refs[:12]:
            code_section = section_for_rva(sections, ref)
            if not code_section:
                continue
            candidates = iter_nearby_rip_code_targets(data, sections, image_base, code_section, ref, args.radius)
            print(f"  xref_disp_rva=0x{ref:x} nearby_code_ptrs={len(candidates)}")
            likely = likely_registered_code(ref, candidates)
            if likely:
                instr_rva, target_rva, label = likely
                print(f"    likely_registered_code {label} instr_rva=0x{instr_rva:x} target_rva=0x{target_rva:x}")
            for instr_rva, target_rva, label in candidates[:16]:
                print(f"    {label} instr_rva=0x{instr_rva:x} target_rva=0x{target_rva:x}")
            if len(candidates) > 16:
                print(f"    ... +{len(candidates) - 16}")

        emitted += 1
        if emitted >= args.max_strings:
            break

    return 0


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