"""
Wow.exe binary patcher for MoonWell custom features.

Patches:
  1. Script_CreateCharacter: reads lua_tonumber(L, 2) as outfitId,
     stores to global, passes to CMSG_CHAR_CREATE packet instead of 0.
  2. GetCharacterInfo: adds 11th return value — charFlags as a Lua number,
     read from struct_offset +0x170 (same field as ghost flag 0x2000).
     Lua checks bit 0x40000000 (CHARACTER_FLAG_UNK31) for traitor status.

Usage:
    python make_modes_in_exe.py <path_to_Wow.exe>

A backup (Wow.exe.bak) is created before patching.
"""

import struct
import sys
import shutil
from pathlib import Path


# ---------------------------------------------------------------------------
# PE helpers
# ---------------------------------------------------------------------------

def _u16(data, off):
    return struct.unpack_from("<H", data, off)[0]

def _u32(data, off):
    return struct.unpack_from("<I", data, off)[0]


def parse_pe_sections(data):
    if data[0:2] != b"MZ":
        raise ValueError("Not a valid PE file (missing MZ)")
    pe_off = _u32(data, 0x3C)
    if data[pe_off:pe_off + 4] != b"PE\x00\x00":
        raise ValueError("Not a valid PE file (missing PE signature)")
    coff = pe_off + 4
    num_secs = _u16(data, coff + 2)
    opt_size = _u16(data, coff + 16)
    opt = coff + 20
    magic = _u16(data, opt)
    if magic == 0x10b:
        image_base = _u32(data, opt + 28)
    elif magic == 0x20b:
        image_base = struct.unpack_from("<Q", data, opt + 24)[0]
    else:
        raise ValueError(f"Unknown PE magic: 0x{magic:04x}")
    sec_table = opt + opt_size
    sections = []
    for i in range(num_secs):
        o = sec_table + i * 40
        name      = data[o:o + 8].rstrip(b"\x00").decode("ascii", errors="replace")
        virt_size = _u32(data, o + 8)
        virt_addr = _u32(data, o + 12)
        raw_size  = _u32(data, o + 16)
        raw_off   = _u32(data, o + 20)
        sections.append((name, virt_addr, virt_size, raw_off, raw_size))
    return image_base, sections


def va_to_offset(va, sections, image_base):
    rva = va - image_base
    for name, va0, vs, ro, _ in sections:
        if va0 <= rva < va0 + vs:
            return ro + (rva - va0)
    raise ValueError(f"VA 0x{va:08x} not in any section")


def offset_to_va(off, sections, image_base):
    for name, va0, _, ro, rs in sections:
        if ro <= off < ro + rs:
            return image_base + va0 + (off - ro)
    raise ValueError(f"File offset 0x{off:08x} not in any section")


def find_cave(data, sections, image_base, size, search_from_va):
    """Return VA of the first run of >= size consecutive CC bytes after search_from_va."""
    text = next((s for s in sections if s[0] == ".text"), None)
    if not text is None:
        _, va0, vs, ro, rs = text
    else:
        raise ValueError("No .text section")

    rva_start = max(search_from_va - image_base, va0)
    off_start = ro + (rva_start - va0)
    off_end   = ro + min(vs, rs)

    i = off_start
    while i <= off_end - size:
        if data[i] != 0xcc:
            i += 1
            continue
        j = i
        while j < off_end and data[j] == 0xcc:
            j += 1
        if j - i >= size:
            return offset_to_va(i, sections, image_base)
        i = j
    raise ValueError(f"No {size}-byte code cave found from VA 0x{search_from_va:08x}")


# ---------------------------------------------------------------------------
# Patch 1 & 2: outfitId in CMSG_CHAR_CREATE  (fixed addresses)
# ---------------------------------------------------------------------------

_FIXED_PATCHES = [
    (
        0x004e0c60,
        bytes.fromhex(
            "558bec568b75086a0156e8f1d23600"
            "83c40885c0740d6a006a0156e860d4"
            "360083c40c6a006a0156e853d43600"
            "50e8edf6ffff83c41033c05e5dc3cc"
            "cccccccc"
        ),
        bytes.fromhex(
            "558bec568b75086a0256e8c1d33600"
            "db1c245883c404a221 42ac006a006a"
            "0156e85bd4360050e8f5f6ffff83c4"
            "1033c05e5dc38855f08a1521 42ac00"
            "8855f4c3"
        ).replace(b" ", b""),  # spaces stripped for readability
        "Script_CreateCharacter: read outfitId from Lua arg 2 -> global",
    ),
    (
        0x004e042f,
        bytes.fromhex("8855f0c645f400"),
        bytes.fromhex("e85f08000090 90").replace(b" ", b""),
        "Packet builder: CALL helper -> load face + outfitId from global",
    ),
]

# ---------------------------------------------------------------------------
# Patch 3: GetCharacterInfo — add 11th return value (charFlags as number)
# ---------------------------------------------------------------------------
#
# Assembly context (FUN_004e3170 epilogue, valid branch):
#
#   004e31cd  56              PUSH ESI                   ← saves caller's ESI
#   004e31ce  8b f0           MOV ESI, EAX               ← ESI = char struct ptr
#   ...  [function body — ESI used throughout for char struct] ...
#   004e32e5  8b 8e 70010000  MOV ECX, [ESI+0x170]       ← charFlags (ghost flag check)
#   ...
#   004e332d  83 c4 2c        ADD ESP, 0x2c              ← clean accumulated args  ← PATCH HERE
#   004e3330  5e              POP ESI                    ← restores caller's ESI (ESI valid until here)
#   004e3331  b8 0a 00 00 00  MOV EAX, 0xa              (dead after patch)
#   004e3336  5f              POP EDI                    ← JMP back target
#   004e3337  8b e5           MOV ESP, EBP
#   004e3339  5d              POP EBP
#   004e333a  c3              RET
#
# KEY: ESI = char struct ptr up to and including ADD ESP,0x2c (004e332d).
#      After POP ESI (004e3330) it is caller's ESI — NOT char struct.
#      Previous approach: CALL cave at 004e3331 (after POP ESI) -> ESI already wrong.
#      Fix: JMP cave at 004e332d (before POP ESI) -> ESI still valid in cave.
#
# New approach: patch 5 bytes at 004e332d with JMP cave.
#   Old bytes: 83 c4 2c (ADD ESP,0x2c) + 5e (POP ESI) + b8 (MOV EAX first byte)
#   New bytes: e9 XX XX XX XX  JMP cave
#   The 4 dead bytes at 004e3332-004e3335 are unreachable.
#
# Cave (30 bytes, found automatically via CC scan over entire .text):
#   83 ec 08              SUB ESP, 8              ; make room for double
#   db 86 70 01 00 00     FILD [ESI+0x170]        ; load charFlags into FPU (ESI valid here!)
#   dd 1c 24              FSTP [ESP]              ; store as double
#   57                    PUSH EDI                ; lua_State* L
#   e8 XX XX XX XX        CALL FUN_0084e2a0       ; lua_pushnumber(L, charFlags)
#   83 c4 38              ADD ESP, 0x38           ; merged: own 0xc + stolen 0x2c
#   5e                    POP ESI                 ; stolen: restore caller's ESI
#   6a 0b                 PUSH 0xb               ; \  EAX = 11 in 3 bytes
#   58                    POP EAX                ; /  (vs MOV EAX,0xb = 5 bytes)
#   e9 XX XX XX XX        JMP 004e3336            ; continue epilogue (POP EDI ...)
#
# Lua code reads info[11] and checks bit 0x40000000 (CHARACTER_FLAG_UNK31 = traitor).
# charFlags confirmed at [ESI+0x170]: same field used for ghost flag (AND ECX,0x2000).

_CHARINFO_PATCH_VA      = 0x004e332d    # ADD ESP,0x2c (5 bytes including POP ESI + MOV EAX[0])
_CHARINFO_JMPBACK_VA    = 0x004e3336    # POP EDI — resume epilogue after cave
_CHARINFO_CAVE_SIZE     = 30
_CHARINFO_CAVE_SEARCH_FROM = 0x00401000  # search entire .text for first 35-byte CC block
_LUA_PUSHNUMBER_VA      = 0x0084e2a0    # lua_pushnumber


def _build_charinfo_patches(data, sections, image_base):
    """
    Build the dynamic GetCharacterInfo patches.
    Returns list of (va, old_bytes, new_bytes, description), [] if already applied,
    or None on unrecoverable error.
    """
    patch_off = va_to_offset(_CHARINFO_PATCH_VA, sections, image_base)
    current = bytes(data[patch_off:patch_off + 5])

    # Already patched with new JMP approach?
    if current[0:1] == b"\xe9":
        return []

    # Verify first 4 bytes look like the epilogue (ADD ESP,0x2c + POP ESI).
    # Accept both unpatched (b8 at [4]) and old-CALL-patched (e8 at [4]) states.
    if current[0:4] != bytes.fromhex("83c42c5e"):
        print(f"[WARN] 0x{_CHARINFO_PATCH_VA:08x}: unexpected bytes {current.hex(' ')}")
        print("       Expected ADD ESP,0x2c + POP ESI at this address.")
        print("       Restore Wow.exe.bak and re-run.")
        return None

    old_b = current  # exact 5 bytes currently in file (b8 or e8 at index 4)

    # Find a 35-byte code cave in .text
    cave_va = find_cave(data, sections, image_base,
                        _CHARINFO_CAVE_SIZE, _CHARINFO_CAVE_SEARCH_FROM)

    # Cave layout offsets (30 bytes total):
    #   0  SUB ESP,8           (3)
    #   3  FILD [ESI+0x170]    (6)
    #   9  FSTP [ESP]          (3)
    #  12  PUSH EDI            (1)
    #  13  CALL lua_pushnumber (5) — next instr at 18
    #  18  ADD ESP, 0x38       (3)  merged: own 0xc + stolen 0x2c
    #  21  POP ESI             (1)  stolen
    #  22  PUSH 0xb            (2)  \  shorter than MOV EAX,0xb (5 bytes)
    #  24  POP EAX             (1)  /
    #  25  JMP 004e3336        (5) — next instr at cave_va+30
    call_rel     = _LUA_PUSHNUMBER_VA - (cave_va + 18)
    jmp_back_rel = _CHARINFO_JMPBACK_VA - (cave_va + 30)

    cave_bytes = (
        bytes.fromhex("83ec08")                      # SUB ESP, 8               (3)
        + bytes.fromhex("db8670010000")               # FILD [ESI+0x170]         (6)
        + bytes.fromhex("dd1c24")                     # FSTP [ESP]               (3)
        + bytes.fromhex("57")                         # PUSH EDI                 (1)
        + b"\xe8" + struct.pack("<i", call_rel)        # CALL lua_pushnumber      (5)
        + bytes.fromhex("83c438")                     # ADD ESP, 0x38 (0xc+0x2c) (3)
        + bytes.fromhex("5e")                         # POP ESI [stolen]         (1)
        + bytes.fromhex("6a0b")                       # PUSH 0xb                 (2)
        + bytes.fromhex("58")                         # POP EAX  -> EAX=11       (1)
        + b"\xe9" + struct.pack("<i", jmp_back_rel)   # JMP 004e3336             (5)
    )
    assert len(cave_bytes) == _CHARINFO_CAVE_SIZE, \
        f"Cave size {len(cave_bytes)} != {_CHARINFO_CAVE_SIZE}"

    # Patch: JMP cave (5 bytes at 004e332d)
    jmp_rel = cave_va - (_CHARINFO_PATCH_VA + 5)
    patch_new = b"\xe9" + struct.pack("<i", jmp_rel)

    return [
        (
            cave_va,
            bytes([0xcc] * _CHARINFO_CAVE_SIZE),
            cave_bytes,
            f"GetCharacterInfo cave @ 0x{cave_va:08x}: push charFlags, stolen epilogue",
        ),
        (
            _CHARINFO_PATCH_VA,
            old_b,
            patch_new,
            "GetCharacterInfo: JMP to cave before POP ESI (ESI=char struct), return 11",
        ),
    ]


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    if len(sys.argv) != 2:
        print(f"Usage: python {Path(sys.argv[0]).name} <path_to_Wow.exe>")
        sys.exit(1)

    exe_path = Path(sys.argv[1])
    if not exe_path.exists():
        print(f"Error: file not found: {exe_path}")
        sys.exit(1)

    data = bytearray(exe_path.read_bytes())
    image_base, sections = parse_pe_sections(data)

    print(f"Image base: 0x{image_base:08x}")
    print(f"Sections: {len(sections)}")
    for n, va, vs, ro, rs in sections:
        print(f"  {n:8s}  VA=0x{va:08x}  VSize=0x{vs:08x}  Raw=0x{ro:08x}  RSize=0x{rs:08x}")
    print()

    charinfo_patches = _build_charinfo_patches(data, sections, image_base)
    if charinfo_patches is None:
        sys.exit(1)
    all_patches = list(_FIXED_PATCHES) + charinfo_patches

    # Verify all first
    ok = True
    for va, old_b, new_b, desc in all_patches:
        off = va_to_offset(va, sections, image_base)
        cur = bytes(data[off:off + len(old_b)])
        if cur == new_b:
            print(f"[SKIP] 0x{va:08x}: already patched — {desc}")
        elif cur == old_b:
            print(f"[ OK ] 0x{va:08x}: verified — {desc}")
        else:
            print(f"[FAIL] 0x{va:08x}: {desc}")
            print(f"  Expected: {old_b.hex(' ')}")
            print(f"  Found:    {cur.hex(' ')}")
            ok = False

    if not ok:
        print("\nByte mismatch — aborting. Is this the correct Wow.exe (3.3.5a build 12340)?")
        sys.exit(1)

    # Backup
    bak = exe_path.with_suffix(".exe.bak")
    if not bak.exists():
        shutil.copy2(exe_path, bak)
        print(f"\nBackup: {bak}")
    else:
        print(f"\nBackup already exists: {bak}")

    # Apply
    applied = 0
    for va, old_b, new_b, desc in all_patches:
        off = va_to_offset(va, sections, image_base)
        cur = bytes(data[off:off + len(old_b)])
        if cur == new_b:
            continue
        data[off:off + len(old_b)] = new_b
        applied += 1
        print(f"[PATCH] 0x{va:08x}: {desc}")

    if applied == 0:
        print("\nNothing to patch.")
    else:
        exe_path.write_bytes(data)
        print(f"\n{applied} patch(es) written to {exe_path}")

    print("Done!")


if __name__ == "__main__":
    main()