"""Read-only instruction contexts for an exact AMD64 PE immediate candidate.

Uses exception-table bounds and validates byte candidates against disassembly.
No native execution; output is not a proof of a field/command's semantics.
"""
import argparse
import bisect
import hashlib
import struct
from pathlib import Path
from capstone import Cs, CS_ARCH_X86, CS_MODE_64

p = argparse.ArgumentParser(description=__doc__)
p.add_argument('--binary', type=Path, required=True)
p.add_argument('--value', type=lambda s: int(s, 0), required=True)
p.add_argument('--width', type=int, choices=(2, 4), default=4)
p.add_argument('--output', type=Path, required=True)
a = p.parse_args()
b = a.binary.read_bytes()
pe = struct.unpack_from('<I', b, 0x3c)[0]
if b[:2] != b'MZ' or b[pe:pe+4] != b'PE\0\0' or struct.unpack_from('<H', b, pe+4)[0] != 0x8664:
    raise SystemExit('Expected AMD64 PE image')
count = struct.unpack_from('<H', b, pe+6)[0]
optional_size = struct.unpack_from('<H', b, pe+20)[0]
sections = []
for i in range(count):
    at = pe+24+optional_size+i*40
    _, rva, size, raw = struct.unpack_from('<IIII', b, at+8)
    flags = struct.unpack_from('<I', b, at+36)[0]
    sections.append((rva, size, raw, flags))

def offset(rva):
    for start, size, raw, _ in sections:
        if start <= rva < start+size:
            return raw+rva-start
    raise ValueError('RVA outside raw sections')

exceptions_rva, exception_size = struct.unpack_from('<II', b, pe+24+112+3*8)
exceptions = [struct.unpack_from('<III', b, at)[:2]
              for at in range(offset(exceptions_rva), offset(exceptions_rva)+exception_size, 12)]
exceptions.sort()
starts = [x[0] for x in exceptions]
needle = struct.pack('<H' if a.width == 2 else '<I', a.value)
decoder = Cs(CS_ARCH_X86, CS_MODE_64)
lines = [f'; SHA256 {hashlib.sha256(b).hexdigest()}', f'; selected immediate {a.value:#x}']
seen = set()
decoded = {}
for rva, size, raw, flags in sections:
    if not flags & 0x20000000:
        continue
    at = raw
    while True:
        at = b.find(needle, at, raw+size)
        if at < 0:
            break
        candidate = rva+at-raw
        at += 1
        index = bisect.bisect_right(starts, candidate)-1
        if index < 0:
            continue
        begin, end = exceptions[index]
        if not begin <= candidate < end or not 0 < end-begin <= 65536:
            continue
        if begin not in decoded:
            decoded[begin] = list(decoder.disasm(b[offset(begin):offset(begin)+end-begin], begin))
        instructions = decoded[begin]
        for i, instruction in enumerate(instructions):
            if instruction.address <= candidate < instruction.address+instruction.size:
                if a.value not in [int(x, 16) for x in __import__('re').findall(r'0x[0-9a-f]+', instruction.op_str)]:
                    continue
                key = instruction.address
                if key in seen:
                    continue
                seen.add(key)
                lines.append(f'; function {begin:#x}..{end:#x}; candidate {candidate:#x}')
                lines.extend(f'{x.address:x}: {x.mnemonic} {x.op_str}' for x in instructions[max(0, i-6):i+7])
                lines.append('')
        if len(seen) > 128:
            raise SystemExit('Candidate bound reached')
a.output.write_text('\n'.join(lines)+'\n', encoding='utf-8')
print(f'Wrote {len(seen)} instruction contexts to {a.output}')