"""install_leakfix.py — Install leakfix.dll into an Asheron's Call install.

Patches acclient.exe to statically import leakfix.dll. Adds a new PE
section ".limport" containing a rebuilt import table that includes
leakfix.dll alongside the original imports.

Usage:
    python install_leakfix.py [ac_dir]

Defaults ac_dir to "C:\\Turbine\\Asheron's Call".

The script:
  1. Locates acclient.exe and leakfix.dll in ac_dir
  2. Checks whether acclient.exe already imports leakfix.dll
     (idempotent — exits cleanly if already patched)
  3. Saves a backup of the original acclient.exe (if no backup exists)
  4. Builds a new import table with leakfix.dll added at the end
  5. Appends a new PE section ".limport" containing the new import table
  6. Updates the PE Optional Header DataDirectory[1] (Import Table)
     to point at the new section
  7. Updates NumberOfSections in the COFF header

acclient.exe section table has 1 unused slot already (verified — it has
7 sections but room for 8 in headers), so we can append cleanly without
needing to move .text.
"""
import struct
import sys
import os
import shutil
import hashlib

DEFAULT_AC_DIR = r"C:\Turbine\Asheron's Call"

# leakfix.dll exports — leakfix.dll has only DllMain (no exported functions
# we need to call by name), so the import name table just needs ONE entry
# with a fake function name or an ordinal import. Easiest: import by ordinal
# 1 if the DLL exports anything, or use a forwarder. But leakfix.dll has no
# exports, so we need to add one — see leakfix DLL build for a stub export.
#
# For a DLL that has *zero* exports, an empty IAT entry (just the terminator)
# in the import descriptor would mean "load the DLL but don't bind anything".
# Windows still calls DllMain on load. That's what we want.
#
# Required structure per import descriptor:
#   OriginalFirstThunk (RVA to import lookup table) - terminator entry only
#   TimeDateStamp = 0
#   ForwarderChain = 0
#   Name (RVA to DLL name string)
#   FirstThunk (RVA to import address table) - terminator entry only

def hexdump_at(data, off, n=64):
    chunk = data[off:off+n]
    return ' '.join(f'{b:02x}' for b in chunk)

def find_section(data, pe_off, name):
    num_sections = struct.unpack_from('<H', data, pe_off + 6)[0]
    opt_header_size = struct.unpack_from('<H', data, pe_off + 20)[0]
    sect_off = pe_off + 24 + opt_header_size
    for i in range(num_sections):
        so = sect_off + i*40
        sname = data[so:so+8].rstrip(b'\0')
        if sname == name.encode():
            return so, i
    return None, -1

def get_sections(data, pe_off):
    num_sections = struct.unpack_from('<H', data, pe_off + 6)[0]
    opt_header_size = struct.unpack_from('<H', data, pe_off + 20)[0]
    sect_off = pe_off + 24 + opt_header_size
    out = []
    for i in range(num_sections):
        so = sect_off + i*40
        out.append({
            'name': data[so:so+8].rstrip(b'\0').decode(errors='replace'),
            'vsize': struct.unpack_from('<I', data, so+8)[0],
            'vaddr': struct.unpack_from('<I', data, so+12)[0],
            'rsize': struct.unpack_from('<I', data, so+16)[0],
            'rawoff': struct.unpack_from('<I', data, so+20)[0],
            'chars': struct.unpack_from('<I', data, so+36)[0],
            'tabentry_off': so,
        })
    return out

def rva_to_off(sections, rva):
    for s in sections:
        if s['vaddr'] <= rva < s['vaddr'] + s['vsize']:
            return s['rawoff'] + (rva - s['vaddr'])
    return None

def already_imports_leakfix(data, pe_off, sections):
    opt_off = pe_off + 24
    dd_off = opt_off + 96
    import_rva = struct.unpack_from('<I', data, dd_off + 8)[0]
    foff = rva_to_off(sections, import_rva)
    if foff is None: return False
    while True:
        name_rva = struct.unpack_from('<I', data, foff + 12)[0]
        if name_rva == 0: return False
        name_foff = rva_to_off(sections, name_rva)
        if name_foff is None: return False
        name_end = data.index(b'\0', name_foff)
        dll_name = data[name_foff:name_end].decode(errors='replace').lower()
        if 'leakfix' in dll_name:
            return True
        foff += 20

def align(n, a):
    return (n + a - 1) & ~(a - 1)

def patch(ac_dir):
    exe_path = os.path.join(ac_dir, 'acclient.exe')
    dll_path = os.path.join(ac_dir, 'leakfix.dll')
    if not os.path.exists(exe_path):
        print(f"ERROR: {exe_path} not found")
        return 1
    if not os.path.exists(dll_path):
        print(f"ERROR: {dll_path} not found")
        return 1

    with open(exe_path, 'rb') as f:
        data = bytearray(f.read())
    orig_sha = hashlib.sha256(data).hexdigest()
    print(f"acclient.exe: {len(data):,} bytes, SHA-256 {orig_sha}")

    pe_off = struct.unpack_from('<I', data, 0x3C)[0]
    sections = get_sections(data, pe_off)
    print(f"\nCurrent sections ({len(sections)}):")
    for s in sections:
        print(f"  {s['name']:10s} vaddr=0x{s['vaddr']:08x} vsize=0x{s['vsize']:08x} raw=0x{s['rawoff']:08x} rsize=0x{s['rsize']:08x}")

    if already_imports_leakfix(data, pe_off, sections):
        print("\n[already patched] acclient.exe already imports leakfix.dll — nothing to do.")
        return 0

    # Backup original (if no backup exists)
    backup = exe_path + '.bare_original'
    if not os.path.exists(backup):
        shutil.copy2(exe_path, backup)
        print(f"\nbackup saved: {backup}")
    else:
        print(f"\nbackup exists: {backup}")

    # Read PE optional header info we'll need
    opt_off = pe_off + 24
    file_align = struct.unpack_from('<I', data, opt_off + 36)[0]
    sect_align = struct.unpack_from('<I', data, opt_off + 32)[0]
    print(f"file_align=0x{file_align:x} sect_align=0x{sect_align:x}")

    # Verify we have a free slot in the section header table.
    num_sections = struct.unpack_from('<H', data, pe_off + 6)[0]
    opt_header_size = struct.unpack_from('<H', data, pe_off + 20)[0]
    sect_table_off = pe_off + 24 + opt_header_size
    next_sect_entry = sect_table_off + num_sections * 40
    first_section_raw = min(s['rawoff'] for s in sections if s['rawoff'] > 0)
    if next_sect_entry + 40 > first_section_raw:
        print(f"\nERROR: no room in section table at file 0x{next_sect_entry:x} "
              f"(first section raw 0x{first_section_raw:x}). Would need to expand headers.")
        return 1
    print(f"section table free slot at 0x{next_sect_entry:x}, first section raw 0x{first_section_raw:x}: OK")

    # Read original import table to copy as-is
    dd_off = opt_off + 96
    old_imp_rva  = struct.unpack_from('<I', data, dd_off + 8)[0]
    old_imp_size = struct.unpack_from('<I', data, dd_off + 12)[0]
    old_imp_foff = rva_to_off(sections, old_imp_rva)
    # Count old entries (excluding terminator)
    walk = old_imp_foff
    old_n = 0
    while struct.unpack_from('<I', data, walk + 12)[0] != 0:
        old_n += 1
        walk += 20
    old_iid_total = (old_n + 1) * 20   # +1 terminator
    print(f"old imports: {old_n} entries, IID block {old_iid_total} bytes at RVA 0x{old_imp_rva:x}")

    # === Build new section content ===
    #
    # Layout in the new section (.limport) at vaddr V (chosen as next aligned
    # virtual address past last section):
    #   +0:                Import Descriptors (rebuilt: old N + leakfix + terminator)
    #   +(N+2)*20:         IAT terminator for leakfix (one ULONG = 0)
    #   +(N+2)*20+4:       INT terminator for leakfix (one ULONG = 0)
    #   +(N+2)*20+8:       Hint/Name table for leakfix (empty — none needed)
    #   +(N+2)*20+8:       "leakfix.dll\0" name string
    #
    # Wait — actually for a DLL with no required imports we'd have an empty IAT/INT
    # consisting of just a 0 terminator. But Windows loaders are fine with that.

    last = max(sections, key=lambda s: s['vaddr'] + s['vsize'])
    new_vaddr = align(last['vaddr'] + last['vsize'], sect_align)
    new_rawoff = align(len(data), file_align)
    print(f"new section: vaddr=0x{new_vaddr:08x} rawoff=0x{new_rawoff:08x}")

    # Lay out contents
    blob = bytearray()
    # 1) IID block placeholder (size = (old_n + 2) * 20)
    iid_start = 0
    iid_size  = (old_n + 2) * 20
    blob += b'\0' * iid_size
    # 2) Empty IAT/INT terminator for leakfix (4 bytes each, RVA pointers stored
    #    in the leakfix descriptor)
    leakfix_int_off = len(blob)
    blob += b'\0\0\0\0'   # INT terminator
    leakfix_iat_off = len(blob)
    blob += b'\0\0\0\0'   # IAT terminator
    # 3) leakfix.dll name string
    leakfix_name_off = len(blob)
    blob += b'leakfix.dll\0'
    # Pad to file alignment
    while len(blob) % 4 != 0:
        blob += b'\0'

    new_section_size = len(blob)
    new_section_raw_size = align(new_section_size, file_align)
    blob += b'\0' * (new_section_raw_size - new_section_size)

    # Fill in IIDs: copy old N, then add leakfix at index N, then terminator
    # Each IID: OriginalFirstThunk, TimeDateStamp, ForwarderChain, Name, FirstThunk = 5 * 4B = 20B
    # Old IIDs we copy verbatim (their RVAs already point into the existing import section)
    blob[iid_start:iid_start + old_n*20] = data[old_imp_foff : old_imp_foff + old_n*20]
    # leakfix IID at index N
    leakfix_iid_off = iid_start + old_n * 20
    leakfix_iid = bytearray(20)
    struct.pack_into('<I', leakfix_iid, 0,  new_vaddr + leakfix_int_off)   # OriginalFirstThunk (INT)
    struct.pack_into('<I', leakfix_iid, 4,  0)                              # TimeDateStamp
    struct.pack_into('<I', leakfix_iid, 8,  0)                              # ForwarderChain
    struct.pack_into('<I', leakfix_iid, 12, new_vaddr + leakfix_name_off)   # Name
    struct.pack_into('<I', leakfix_iid, 16, new_vaddr + leakfix_iat_off)    # FirstThunk (IAT)
    blob[leakfix_iid_off:leakfix_iid_off + 20] = leakfix_iid
    # Terminator IID already zeroed

    # === Append new section ===
    if new_rawoff > len(data):
        data += b'\0' * (new_rawoff - len(data))   # pad to alignment
    data += blob

    # === Write new section header ===
    new_sect_hdr = bytearray(40)
    new_sect_hdr[0:8]   = b'.limport'                                          # name
    struct.pack_into('<I', new_sect_hdr, 8,  new_section_size)                 # VirtualSize
    struct.pack_into('<I', new_sect_hdr, 12, new_vaddr)                       # VirtualAddress
    struct.pack_into('<I', new_sect_hdr, 16, new_section_raw_size)            # SizeOfRawData
    struct.pack_into('<I', new_sect_hdr, 20, new_rawoff)                      # PointerToRawData
    struct.pack_into('<I', new_sect_hdr, 36, 0xC0000040)                      # READ|WRITE|INITIALIZED_DATA
    data[next_sect_entry:next_sect_entry + 40] = new_sect_hdr

    # === Update PE headers ===
    struct.pack_into('<H', data, pe_off + 6, num_sections + 1)                # NumberOfSections
    # DataDirectory[1] (Import Table) — point at new IID block
    struct.pack_into('<I', data, dd_off + 8,  new_vaddr + iid_start)
    struct.pack_into('<I', data, dd_off + 12, iid_size)
    # SizeOfImage = aligned end of new section
    new_size_of_image = align(new_vaddr + new_section_size, sect_align)
    struct.pack_into('<I', data, opt_off + 56, new_size_of_image)

    # === Save patched executable ===
    out_path = exe_path
    with open(out_path, 'wb') as f:
        f.write(data)
    new_sha = hashlib.sha256(data).hexdigest()
    print(f"\nPatched acclient.exe written: {len(data):,} bytes, SHA-256 {new_sha}")
    print(f"  delta vs original: +{len(data) - os.path.getsize(backup):,} bytes")
    print(f"  new section: .limport @ vaddr 0x{new_vaddr:08x} raw 0x{new_rawoff:08x}")
    print(f"\n[OK] acclient.exe now imports leakfix.dll. Restart any running clients.")
    return 0

def verify(ac_dir):
    """Read patched exe and confirm leakfix.dll is in imports."""
    exe_path = os.path.join(ac_dir, 'acclient.exe')
    with open(exe_path, 'rb') as f:
        data = f.read()
    pe_off = struct.unpack_from('<I', data, 0x3C)[0]
    sections = get_sections(data, pe_off)
    if already_imports_leakfix(data, pe_off, sections):
        print(f"[verified] {exe_path} imports leakfix.dll.")
        return 0
    print(f"[FAILED] {exe_path} does NOT import leakfix.dll.")
    return 1

def main():
    ac_dir = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_AC_DIR
    cmd = sys.argv[2] if len(sys.argv) > 2 else 'patch'
    print(f"AC directory: {ac_dir}\nCommand: {cmd}\n")
    if cmd == 'patch':
        return patch(ac_dir)
    elif cmd == 'verify':
        return verify(ac_dir)
    else:
        print(f"unknown command: {cmd}")
        return 1

if __name__ == '__main__':
    sys.exit(main())