import struct, sys

# Set TEST_MODE=True to use dwExStyle=0 (original), False for WS_EX_APPWINDOW
TEST_MODE = True

exe = bytearray(open("HAVOC_NOCD.EXE", "rb").read())
BASE = 0x400000

# ---- 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

TEXT_SECHDOFF = None
for i in range(num_sects):
    sh = sects_off + i * 40
    name = exe[sh:sh+8].rstrip(b'\x00')
    if name == b'.text':
        TEXT_SECHDOFF = sh
        break

if TEXT_SECHDOFF is None:
    print("ERROR: .text section not found"); sys.exit(1)

sh = TEXT_SECHDOFF
TEXT_VRVA   = struct.unpack_from('<I', exe, sh+12)[0]
TEXT_VSIZ   = struct.unpack_from('<I', exe, sh+8)[0]
TEXT_RAWSIZ = struct.unpack_from('<I', exe, sh+16)[0]
TEXT_RAWOFF = struct.unpack_from('<I', exe, sh+20)[0]

print(f".text: VirtRVA=0x{TEXT_VRVA:x} VirtSize=0x{TEXT_VSIZ:x} "
      f"RawOff=0x{TEXT_RAWOFF:x} RawSize=0x{TEXT_RAWSIZ:x}")

def va2fo(va):
    rva = va - BASE
    return TEXT_RAWOFF + (rva - TEXT_VRVA)

# ---- Cave: placed in raw padding beyond ORIGINAL VirtualSize ----------------
# Hardcode original VirtualSize so this script is safe to re-run even after
# VirtualSize has already been extended to RawSize.
ORIG_TEXT_VSIZ = 0x66bbd
cave_fo = TEXT_RAWOFF + ORIG_TEXT_VSIZ   # FO 0x66fbd (raw padding start)
cave_va = BASE + TEXT_VRVA + ORIG_TEXT_VSIZ  # VA 0x467bbd
PATCH_VA = 0x4260c7
patch_fo = va2fo(PATCH_VA)
BACK_VA  = 0x4260d3

print(f"Cave FO=0x{cave_fo:05x}  VA=0x{cave_va:07x}")

# Verify cave area is zero/CC/or previously-patched bytes (safe to overwrite)
cave_bytes = exe[cave_fo:cave_fo+20]
all_safe = all(b in (0x00, 0xCC) for b in cave_bytes)
print(f"Cave current bytes: {cave_bytes.hex(' ')}")
if not all_safe:
    print("(overwriting previous cave code)")

# ---- Verify patch site (may already be patched; accept either) -------------
original_12 = bytes.fromhex("740051689cda47006a00ffd0")
actual_12   = bytes(exe[patch_fo:patch_fo+12])
if actual_12 == original_12:
    print("Patch site: original (unpatched)")
elif actual_12[0] == 0x51 and actual_12[1] == 0xe9:
    print("Patch site: already patched -- overwriting")
else:
    print(f"ERROR: unexpected bytes at patch site: {actual_12.hex(' ')}")
    sys.exit(1)

# ---- Cave code: explicit pushes (no LEA trick) to minimise risk -----------
# Push className as immediate, not computed.
# arg3 (title/lpWindowName) is pushed at patch site via 'push ecx'.
# Cave pushes arg2 (className) and arg1 (dwExStyle), then calls CreateWindowExA.
CLASSNAME_VA = 0x0047da9c
EX_STYLE = 0 if TEST_MODE else 0x00040000
print(f"dwExStyle = 0x{EX_STYLE:08x}  ({'TEST: original style' if TEST_MODE else 'WS_EX_APPWINDOW'})")

cave_code = bytearray()
cave_code += b'\x68' + struct.pack('<I', CLASSNAME_VA)   # push className (5b)
cave_code += b'\x68' + struct.pack('<I', EX_STYLE)       # push dwExStyle (5b)
cave_code += b'\xFF\xD0'                                  # call eax       (2b)
jmp_back = BACK_VA - (cave_va + len(cave_code) + 5)
cave_code += b'\xE9' + struct.pack('<i', jmp_back)        # jmp BACK_VA    (5b)

CAVE_SIZE = len(cave_code)  # 17
print(f"Cave ({CAVE_SIZE}b): {cave_code.hex(' ')}")
print(f"  jmp_back = {jmp_back:+d} -> IP after jmp = 0x{cave_va+CAVE_SIZE:07x}, "
      f"target = 0x{cave_va+CAVE_SIZE+jmp_back:07x}")

# ---- Patch (12 bytes at 0x4260c7) -----------------------------------------
jmp_cave = cave_va - (PATCH_VA + 1 + 5)
patch = bytearray()
patch += b'\x51'                                      # push ecx (title)  (1b)
patch += b'\xE9' + struct.pack('<i', jmp_cave)        # jmp cave           (5b)
patch += b'\x90' * 6                                  # 6x NOP             (6b)
print(f"Patch (12b): {patch.hex(' ')}")
print(f"  jmp_cave = {jmp_cave:+d} -> IP after jmp = 0x{PATCH_VA+6:07x}, "
      f"target = 0x{PATCH_VA+6+jmp_cave:07x}")

# Confirm jump targets
assert PATCH_VA + 6 + jmp_cave == cave_va, "jmp_cave target mismatch"
assert cave_va + CAVE_SIZE + jmp_back == BACK_VA, "jmp_back target mismatch"
print("Jump target assertions: OK")

# ---- Apply -----------------------------------------------------------------
new_vsiz = TEXT_RAWSIZ  # expose all raw data as executable
struct.pack_into('<I', exe, sh + 8, new_vsiz)
print(f".text VirtualSize: 0x{TEXT_VSIZ:x} -> 0x{new_vsiz:x}")

exe[cave_fo:cave_fo+CAVE_SIZE] = cave_code
exe[patch_fo:patch_fo+12] = patch

open("HAVOC_NOCD.EXE", "wb").write(exe)
print("Patch R applied.")