#!/usr/bin/env python3
"""Disassemble a PE .text range by RVA and annotate direct branch targets."""

from __future__ import annotations

import argparse
import math
import struct
from pathlib import Path

from capstone import Cs, CS_ARCH_X86, CS_MODE_64
from capstone.x86 import X86_OP_IMM, X86_OP_MEM, X86_REG_RIP

from find_pe_string_xrefs import Section, parse_pe


def find_section(sections: list[Section], rva: int) -> Section:
    for section in sections:
        if section.contains_rva(rva):
            return section
    raise ValueError(f"RVA 0x{rva:x} is not inside any PE section")


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


def parse_rva(text: str) -> int:
    return int(text, 16 if text.lower().startswith("0x") else 16)


def format_memory_annotation(data: bytes, sections: list[Section], target_rva: int) -> str:
    section = find_section_or_none(sections, target_rva)
    if section is None:
        return f"rip_rva=0x{target_rva:x}"

    offset = section.rva_to_offset(target_rva)
    if offset < 0 or offset >= len(data):
        return f"rip_rva=0x{target_rva:x}"

    available = data[offset : min(len(data), offset + 8)]
    fields = [f"rip_rva=0x{target_rva:x}"]
    if len(available) >= 4:
        u32 = struct.unpack_from("<I", available)[0]
        f32 = struct.unpack_from("<f", available)[0]
        fields.append(f"u32=0x{u32:08x}")
        if math.isfinite(f32) and abs(f32) < 1.0e12:
            fields.append(f"f32={f32:.9g}")
    if len(available) >= 8:
        u64 = struct.unpack_from("<Q", available)[0]
        fields.append(f"u64=0x{u64:016x}")
    return " ".join(fields)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("pe", type=Path)
    parser.add_argument("rva", type=parse_rva)
    parser.add_argument("--size", type=lambda value: int(value, 0), default=0x200)
    parser.add_argument("--before", type=lambda value: int(value, 0), default=0)
    args = parser.parse_args()

    data = args.pe.read_bytes()
    image_base, sections = parse_pe(data)
    start_rva = max(0, args.rva - args.before)
    section = find_section(sections, start_rva)
    end_rva = min(section.virtual_address + section.raw_size, start_rva + args.size)
    start_offset = section.rva_to_offset(start_rva)
    code = data[start_offset : start_offset + (end_rva - start_rva)]

    disassembler = Cs(CS_ARCH_X86, CS_MODE_64)
    disassembler.detail = True

    print(f"image_base=0x{image_base:x} section={section.name} range=0x{start_rva:x}..0x{end_rva:x}")
    for insn in disassembler.disasm(code, image_base + start_rva):
        rva = insn.address - image_base
        annotation = ""
        if insn.group(1) or insn.group(2):  # jump/call
            for operand in insn.operands:
                if operand.type == X86_OP_IMM:
                    annotation = f" ; target_rva=0x{operand.imm - image_base:x}"
                    break
        elif insn.operands:
            memory_annotations = []
            for operand in insn.operands:
                if operand.type == X86_OP_MEM and operand.mem.base == X86_REG_RIP:
                    target_rva = (insn.address + insn.size + operand.mem.disp) - image_base
                    memory_annotations.append(format_memory_annotation(data, sections, target_rva))
            if memory_annotations:
                annotation = " ; " + " ; ".join(memory_annotations)
        marker = ">>" if rva == args.rva else "  "
        print(f"{marker} 0x{rva:08x}: {insn.mnemonic:<8} {insn.op_str}{annotation}")

    return 0


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