#!/usr/bin/env python3
"""
hunt_pe.py -- static protocol hunt for the Borland PE game binaries
(RPL4OPT/BTL4OPT under 32RTM), the PE counterpart of hunt_actions.py (LE).

Find an error-string's data VA, locate the code xref (32-bit immediate),
and disassemble the surrounding function to read off the protocol constants
(expected reply actions, request sizes) that the DOS-side source never had.

    python hunt_pe.py <exe> <needle-string> [--back N] [--fwd N]
"""
import struct
import sys

import capstone


def pe_sections(data):
    e_lfanew = struct.unpack_from('<I', data, 0x3c)[0]
    machine, nsec = struct.unpack_from('<HH', data, e_lfanew + 4)
    opt_size = struct.unpack_from('<H', data, e_lfanew + 20)[0]
    image_base = struct.unpack_from('<I', data, e_lfanew + 24 + 28)[0]
    secs = []
    off = e_lfanew + 24 + opt_size
    for i in range(nsec):
        name = data[off:off + 8].rstrip(b'\0').decode('latin1')
        vsize, va, rsize, roff = struct.unpack_from('<4I', data, off + 8)
        flags = struct.unpack_from('<I', data, off + 36)[0]
        secs.append({'name': name, 'va': va, 'vsize': vsize,
                     'roff': roff, 'rsize': rsize, 'flags': flags})
        off += 40
    return image_base, secs


def file_to_va(secs, base, foff):
    for s in secs:
        if s['roff'] <= foff < s['roff'] + s['rsize']:
            return base + s['va'] + (foff - s['roff'])
    return None


def va_to_file(secs, base, va):
    rva = va - base
    for s in secs:
        if s['va'] <= rva < s['va'] + max(s['vsize'], s['rsize']):
            return s['roff'] + (rva - s['va'])
    return None


def hunt(path, needle, back=0x120, fwd=0x60):
    data = open(path, 'rb').read()
    base, secs = pe_sections(data)
    print(f"{path}: base {base:#x}, sections "
          + " ".join(f"{s['name']}@{s['va']:#x}" for s in secs))

    hits = []
    i = data.find(needle.encode('latin1'))
    while i != -1:
        hits.append(i); i = data.find(needle.encode('latin1'), i + 1)
    if not hits:
        raise SystemExit(f"string {needle!r} not found")

    md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_32)
    md.detail = False
    for h in hits:
        # xrefs point at the start of the full string literal, which may begin
        # before the needle -- scan back to the previous NUL
        s0 = h
        while s0 > 0 and data[s0 - 1] != 0:
            s0 -= 1
        va = file_to_va(secs, base, s0)
        print(f"\nstring @ file {s0:#x} va {va:#x}: "
              f"{data[s0:s0+70].rstrip(bytes(1))!r}")
        if va is None:
            continue
        pat = struct.pack('<I', va)
        j = data.find(pat)
        while j != -1:
            code_va = file_to_va(secs, base, j)
            sec = next((s for s in secs if s['roff'] <= j < s['roff'] + s['rsize']), None)
            if sec and code_va and (sec['flags'] & 0x20000000 or sec['name'] in ('.text', 'CODE')):
                start = max(j - back, sec['roff'])
                code = data[start:j + fwd]
                print(f"--- xref imm @ file {j:#x} va {code_va:#x} ({sec['name']}) ---")
                for ins in md.disasm(code, file_to_va(secs, base, start)):
                    mark = "  <-- xref" if ins.address <= code_va < ins.address + ins.size else ""
                    print(f"  {ins.address:#010x}  {ins.mnemonic:8} {ins.op_str}{mark}")
            j = data.find(pat, j + 1)


if __name__ == '__main__':
    argv = [a for a in sys.argv[1:] if not a.startswith('--')]
    kw = {}
    for k in ('back', 'fwd'):
        if f'--{k}' in sys.argv:
            kw[k] = int(sys.argv[sys.argv.index(f'--{k}') + 1], 0)
    hunt(argv[0], argv[1], **kw)