"""
Patch V: vtable-setter in init_archive (0x443010)

Root cause: init_archive explicitly clears vtable ([ESI+0]=0) and never sets it.
load_archives later calls 0x443140 (cleanup) for each archive object; with vtable=0,
0x443140 unmaps the file, closes handles, and zeroes field_0xC (directory ptr).
open_archive_file checks field_0xC != 0, so it fails, and get_archive returns 0 =>
"corrupt or unavailable game file" error fires.

Fix: insert MOV [ESI+0], 1 at the init_archive success return path. With vtable=1,
0x443140's vtable-check (TEST EAX; JNZ) skips cleanup, field_0xC is preserved, and
get_archive succeeds.

Mechanism:
  FO 0x42514 (VA 0x443114): replace 5-byte MOV EAX,1 with JMP to cave at FO 0x66FCE
  Cave (FO 0x66FCE, VA 0x467BCE):
    c7 46 00 01 00 00 00  MOV [ESI+0], 1   <- set vtable sentinel
    b8 01 00 00 00        MOV EAX, 1       <- return value = success
    5e 5b c2 08 00        POP ESI; POP EBX; RET 8
"""
import struct, sys

EXE = "HAVOC_NOCD.EXE"
exe = bytearray(open(EXE, "rb").read())
BASE = 0x400000
FO2VA = lambda fo: fo + 0x400C00

# ---- locate .text section header ----
e_lfanew = struct.unpack_from('<I', exe, 0x3c)[0]
coff_off = e_lfanew + 4
num_sects = struct.unpack_from('<H', exe, coff_off + 2)[0]
opt_size  = struct.unpack_from('<H', exe, coff_off + 16)[0]
sects_off = coff_off + 20 + opt_size
for i in range(num_sects):
    sh = sects_off + i * 40
    if exe[sh:sh+8].rstrip(b'\x00') == b'.text':
        text_sh = sh
        text_rawoff = struct.unpack_from('<I', exe, sh+20)[0]
        text_rawsiz = struct.unpack_from('<I', exe, sh+16)[0]
        text_vrva   = struct.unpack_from('<I', exe, sh+12)[0]
        text_vsiz   = struct.unpack_from('<I', exe, sh+8)[0]
        break
print(".text: rawoff=0x%x rawsiz=0x%x vrva=0x%x vsiz=0x%x" %
      (text_rawoff, text_rawsiz, text_vrva, text_vsiz))

ORIG_TEXT_VSIZ = 0x66bbd
cave_fo = text_rawoff + ORIG_TEXT_VSIZ        # FO 0x66fbd (Patch R area)
patch_v_cave_fo = cave_fo + 17                # FO 0x66fce (Patch V area)
patch_v_cave_va = FO2VA(patch_v_cave_fo)

# ---- Verify cave is clean ----
cave_bytes = bytes(exe[patch_v_cave_fo:patch_v_cave_fo+17])
assert all(b in (0x00, 0xCC) for b in cave_bytes), \
    "Cave area not clean: %s" % cave_bytes.hex(' ')
print("Cave at FO 0x%05x VA 0x%07x: clean" % (patch_v_cave_fo, patch_v_cave_va))

# ---- Verify patch site ----
PATCH_FO = 0x42514
expected_at_site = bytes.fromhex('b801000000')  # MOV EAX, 1
actual_at_site   = bytes(exe[PATCH_FO:PATCH_FO+5])
if actual_at_site == expected_at_site:
    print("Patch site FO 0x%05x: original (unpatched)" % PATCH_FO)
elif actual_at_site[0] == 0xe9:
    print("Patch site FO 0x%05x: already patched (JMP present)" % PATCH_FO)
    # verify it points to our cave
    rel = struct.unpack_from('<i', actual_at_site, 1)[0]
    tgt = FO2VA(PATCH_FO + 5) + rel
    print("  Existing JMP -> 0x%07x (cave=0x%07x)" % (tgt, patch_v_cave_va))
    if tgt == patch_v_cave_va:
        print("  -> points to our cave; verifying cave code...")
        cave_expected = bytes.fromhex('c746000100000' + '0b801000000' + '5e5bc20800')
        cave_actual   = bytes(exe[patch_v_cave_fo:patch_v_cave_fo+17])
        if cave_actual == cave_expected:
            print("  Cave code matches. Patch already applied.")
            sys.exit(0)
    print("  -> unexpected; re-applying")
else:
    print("ERROR: unexpected bytes at patch site: %s" % actual_at_site.hex(' '))
    sys.exit(1)

# ---- Build cave code ----
cave_code = bytearray()
cave_code += b'\xc7\x46\x00\x01\x00\x00\x00'   # MOV [ESI+0], 1  (7 bytes)
cave_code += b'\xb8\x01\x00\x00\x00'            # MOV EAX, 1      (5 bytes)
cave_code += b'\x5e\x5b\xc2\x08\x00'            # POP ESI; POP EBX; RET 8  (5 bytes)
assert len(cave_code) == 17, "cave size %d != 17" % len(cave_code)
print("Cave code (%d bytes): %s" % (len(cave_code), cave_code.hex(' ')))

# ---- Build JMP at patch site ----
next_instr_va = FO2VA(PATCH_FO + 5)
rel32 = (patch_v_cave_va - next_instr_va) & 0xFFFFFFFF
jmp_bytes = b'\xe9' + struct.pack('<I', rel32)
print("JMP bytes: %s  (rel=0x%08x -> cave 0x%07x)" % (
    jmp_bytes.hex(' '), rel32, patch_v_cave_va))
# Verify jump lands on cave
assert (next_instr_va + struct.unpack_from('<i', jmp_bytes, 1)[0]) & 0xFFFFFFFF == patch_v_cave_va, \
    "JMP target mismatch"
print("JMP target assertion: OK")

# ---- Extend VirtualSize to RawSize (makes caves executable) ----
new_vsiz = text_rawsiz
print(".text VirtualSize: 0x%x -> 0x%x" % (text_vsiz, new_vsiz))

# ---- Apply ----
exe[patch_v_cave_fo:patch_v_cave_fo+17] = cave_code
exe[PATCH_FO:PATCH_FO+5] = jmp_bytes
struct.pack_into('<I', exe, text_sh + 8, new_vsiz)

open(EXE, "wb").write(exe)
print("Patch V applied successfully.")
print("init_archive now sets vtable=1 at success return.")
print("0x443140 will skip cleanup for initialized archives, preserving field_0xC.")