"""Raw-image xref sweep (gotcha #21: the decomp export has gaps).

Finds every CALL/JMP rel32 in .text whose target is one of the VAs given on the
command line, plus any absolute immediate referencing them.

  python xref.py 0x4902b0 0x490308
"""
import struct, sys

PATH = r"C:\git\bt411\content\BTL4OPT.EXE"
data = open(PATH, "rb").read()
e_lfanew = struct.unpack_from("<I", data, 0x3C)[0]
coff = e_lfanew + 4
num_sec = struct.unpack_from("<H", data, coff + 2)[0]
opt_size = struct.unpack_from("<H", data, coff + 16)[0]
opt = coff + 20
image_base = struct.unpack_from("<I", data, opt + 28)[0]
sec_tbl = opt + opt_size
secs = []
for i in range(num_sec):
    off = sec_tbl + i * 40
    name = data[off:off + 8].rstrip(b"\0").decode("latin1")
    vsize, va, rawsize, rawptr = struct.unpack_from("<IIII", data, off + 8)
    secs.append((name, va, vsize, rawptr, rawsize))


def off_to_va(off):
    for name, va, vsize, rawptr, rawsize in secs:
        if rawptr <= off < rawptr + rawsize:
            return image_base + va + (off - rawptr)
    return None


targets = set(int(a, 16) for a in sys.argv[1:])
hits = []

for name, va, vsize, rawptr, rawsize in secs:
    blob = data[rawptr:rawptr + rawsize]
    base_va = image_base + va
    # rel32 CALL (E8) / JMP (E9)
    for i in range(len(blob) - 5):
        op = blob[i]
        if op not in (0xE8, 0xE9):
            continue
        rel = struct.unpack_from("<i", blob, i + 1)[0]
        tgt = base_va + i + 5 + rel
        if tgt in targets:
            hits.append((base_va + i, "call" if op == 0xE8 else "jmp", tgt, name))
    # absolute immediates (vtable slots, push offset)
    for i in range(len(blob) - 4):
        v = struct.unpack_from("<I", blob, i)[0]
        if v in targets:
            hits.append((base_va + i, "abs32", v, name))

for addr, kind, tgt, sec in sorted(hits):
    print("%08x  %-6s -> %08x   [%s]" % (addr, kind, tgt, sec))
print("\n%d xrefs" % len(hits))