"""find_alloc_size.py <exe_path> <hex_size>
Search the .text section for `push <size>` followed by a call (operator new
or HeapAlloc) — finds code sites that allocate buffers of that size."""
import struct, sys

def main():
    path = sys.argv[1]
    target = int(sys.argv[2], 0)

    with open(path, 'rb') as f:
        data = f.read()
    pe_off = struct.unpack_from('<I', data, 0x3C)[0]
    n_sec = struct.unpack_from('<H', data, pe_off + 4 + 2)[0]
    opt_off = pe_off + 4 + 20
    img_base = struct.unpack_from('<I', data, opt_off + 28)[0]
    opt_sz = struct.unpack_from('<H', data, pe_off + 4 + 16)[0]
    sec_off = opt_off + opt_sz
    for i in range(n_sec):
        s = sec_off + i*40
        name = data[s:s+8].rstrip(b'\x00').decode()
        if name == '.text':
            text_va = img_base + struct.unpack_from('<I', data, s+12)[0]
            raddr = struct.unpack_from('<I', data, s+20)[0]
            rsize = struct.unpack_from('<I', data, s+16)[0]
            text = data[raddr:raddr+rsize]
            break

    # Patterns:
    #   68 XX XX XX XX  e8 YY YY YY YY  =  push imm32; call rel32
    needle1 = bytes([0x68]) + struct.pack('<I', target)
    # And:
    #   6A XX           = push imm8 (only if target fits in byte; unlikely for 0x41000)

    print(f"Searching .text for `push 0x{target:08x}` followed by call...")
    off = 0
    found = 0
    while True:
        off = text.find(needle1, off)
        if off < 0: break
        # Check if next instruction (at +5) is a call (E8) or call indirect (FF 15 / FF 14)
        next_byte = text[off + 5] if off + 5 < len(text) else 0
        marker = ""
        if next_byte == 0xE8:
            target_rel = struct.unpack_from('<i', text, off + 6)[0]
            call_target = text_va + off + 5 + 5 + target_rel
            marker = f"-> call 0x{call_target:08x}"
        elif next_byte == 0xFF:
            marker = "-> call [...]"
        else:
            marker = f"-> next op 0x{next_byte:02x}"
        va = text_va + off
        print(f"  0x{va:08x}: push 0x{target:08x}  {marker}")
        found += 1
        off += 5

    print(f"\nTotal: {found} sites")

if __name__ == '__main__':
    main()