leakhunt/tools/patch_v7_test.py

"""patch_v7_test.py <pid> [--revert]

EXPERIMENTAL v7: NOP BOTH guards in UIElement_ItemList::UpdateEmptySlots
to force the trim logic to run regardless of:
  1. visibility    (jz at 0x004e439d — same NOP as v6)
  2. auto-fit mode (jne at 0x004e43bd — NEW for v7)

Function prologue (EoR):
  0x004e4390:  sub esp, 0x10
  0x004e4393:  push esi
  0x004e4394:  mov esi, ecx               ; this = ecx
  0x004e4396:  call IsVisible
  0x004e439b:  test al, al
  0x004e439d:  jz  +0x1f3                 ; <-- NOP (v6/v7 site 1)
  0x004e43a3:  push ebx
  0x004e43a4:  push ebp
  0x004e43a5:  push edi
  0x004e43a6:  lea eax, [esp+0x14]
  0x004e43aa:  push eax
  0x004e43ab:  push 0x10000015
  0x004e43b0:  mov ecx, esi
  0x004e43b2:  call GetAttribute_Int
  0x004e43b7:  mov eax, [esp+0x14]
  0x004e43bb:  or  edi, -1                ; edi = 0xFFFFFFFF (= -1)
  0x004e43be:  cmp eax, edi
  0x004e43bd:  jne +0x1cd                 ; <-- NOP (v7 site 2)
  ...

After v7 both guards are eliminated. The function runs unconditionally.
The trim loop inside then runs and deletes WAITING-state items at the
end of the array.
"""
import argparse
import ctypes
import ctypes.wintypes as wt
import sys


SITE_VIS_VA   = 0x004e439d
SITE_VIS_ORIG = bytes([0x0f, 0x84, 0xf3, 0x01, 0x00, 0x00])   # jz +0x1f3
SITE_VIS_NOP  = bytes([0x90] * 6)

SITE_AF_VA    = 0x004e43c0
SITE_AF_ORIG  = bytes([0x0f, 0x85, 0xcd, 0x01, 0x00, 0x00])   # jne +0x1cd
SITE_AF_NOP   = bytes([0x90] * 6)


PROCESS_VM_READ           = 0x0010
PROCESS_VM_WRITE          = 0x0020
PROCESS_VM_OPERATION      = 0x0008
PROCESS_QUERY_INFORMATION = 0x0400
PAGE_EXECUTE_READWRITE    = 0x40


k32 = ctypes.windll.kernel32
OpenProcess        = k32.OpenProcess
OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]; OpenProcess.restype = wt.HANDLE
CloseHandle        = k32.CloseHandle
CloseHandle.argtypes = [wt.HANDLE]; CloseHandle.restype = wt.BOOL
WriteProcessMemory = k32.WriteProcessMemory
WriteProcessMemory.argtypes = [wt.HANDLE, wt.LPVOID, wt.LPCVOID, ctypes.c_size_t,
                               ctypes.POINTER(ctypes.c_size_t)]
WriteProcessMemory.restype = wt.BOOL
ReadProcessMemory  = k32.ReadProcessMemory
ReadProcessMemory.argtypes = [wt.HANDLE, wt.LPCVOID, wt.LPVOID, ctypes.c_size_t,
                              ctypes.POINTER(ctypes.c_size_t)]
ReadProcessMemory.restype = wt.BOOL
VirtualProtectEx   = k32.VirtualProtectEx
VirtualProtectEx.argtypes = [wt.HANDLE, wt.LPVOID, ctypes.c_size_t, wt.DWORD,
                             ctypes.POINTER(wt.DWORD)]
VirtualProtectEx.restype = wt.BOOL


def read_bytes(h, addr, n):
    buf = (ctypes.c_ubyte * n)()
    sz = ctypes.c_size_t(0)
    if not ReadProcessMemory(h, addr, buf, n, ctypes.byref(sz)):
        raise OSError(f"read 0x{addr:08x} err={ctypes.get_last_error()}")
    return bytes(buf[:sz.value])


def write_bytes(h, addr, data):
    old_prot = wt.DWORD(0)
    if not VirtualProtectEx(h, addr, len(data), PAGE_EXECUTE_READWRITE, ctypes.byref(old_prot)):
        raise OSError(f"VirtualProtectEx 0x{addr:08x} err={ctypes.get_last_error()}")
    sz = ctypes.c_size_t(0)
    ok = WriteProcessMemory(h, addr, data, len(data), ctypes.byref(sz))
    err = ctypes.get_last_error() if not ok else 0
    restored = wt.DWORD(0)
    VirtualProtectEx(h, addr, len(data), old_prot.value, ctypes.byref(restored))
    if not ok:
        raise OSError(f"write 0x{addr:08x} err={err}")


def apply_or_revert(h, site_va, orig, patched, label, revert):
    cur = read_bytes(h, site_va, 6)
    print(f"  {label} @ 0x{site_va:08x}: current {cur.hex()}")
    if revert:
        if cur == orig:
            print(f"    already original")
            return
        if cur != patched:
            print(f"    UNEXPECTED — neither orig nor patched. refusing.")
            sys.exit(3)
        write_bytes(h, site_va, orig)
        after = read_bytes(h, site_va, 6)
        print(f"    reverted; bytes now: {after.hex()}")
        return

    if cur == patched:
        print(f"    already patched")
        return
    if cur != orig:
        print(f"    UNEXPECTED — bytes don't match expected original. refusing.")
        sys.exit(4)
    write_bytes(h, site_va, patched)
    after = read_bytes(h, site_va, 6)
    print(f"    patched; bytes now: {after.hex()}")
    if after != patched:
        print(f"    MISMATCH — write didn't take")
        sys.exit(5)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("pid", type=int)
    ap.add_argument("--revert", action="store_true",
                    help="restore original instructions at both sites")
    args = ap.parse_args()

    h = OpenProcess(
        PROCESS_VM_READ | PROCESS_VM_WRITE | PROCESS_VM_OPERATION | PROCESS_QUERY_INFORMATION,
        False, args.pid,
    )
    if not h:
        print(f"OpenProcess({args.pid}) err={ctypes.get_last_error()}")
        sys.exit(2)

    print(f"PID {args.pid}")
    apply_or_revert(h, SITE_VIS_VA, SITE_VIS_ORIG, SITE_VIS_NOP,
                    "visibility guard (jz)", args.revert)
    apply_or_revert(h, SITE_AF_VA, SITE_AF_ORIG, SITE_AF_NOP,
                    "auto-fit guard (jne) ", args.revert)
    print("  OK")
    CloseHandle(h)


if __name__ == "__main__":
    main()