import ctypes
import os
import subprocess
import sys
from ctypes import wintypes

from capstone import CS_ARCH_X86, CS_MODE_64, Cs

PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_VM_READ = 0x0010
PAGE = 0x1000

k32 = ctypes.WinDLL("kernel32", use_last_error=True)
psapi = ctypes.WinDLL("psapi", use_last_error=True)


class ModuleInfo(ctypes.Structure):
    _fields_ = [
        ("base", ctypes.c_void_p),
        ("size", wintypes.DWORD),
        ("entry", ctypes.c_void_p),
    ]


k32.OpenProcess.restype = wintypes.HANDLE
k32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
k32.ReadProcessMemory.argtypes = [wintypes.HANDLE, ctypes.c_void_p, ctypes.c_void_p,
                                  ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)]
psapi.EnumProcessModules.argtypes = [wintypes.HANDLE, ctypes.POINTER(wintypes.HMODULE),
                                     wintypes.DWORD, ctypes.POINTER(wintypes.DWORD)]
psapi.GetModuleInformation.argtypes = [wintypes.HANDLE, wintypes.HMODULE,
                                       ctypes.c_void_p, wintypes.DWORD]


def main_module(pid):
    handle = k32.OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, False, pid)
    if not handle:
        raise OSError(f"OpenProcess failed: {ctypes.get_last_error()}")
    mods = (wintypes.HMODULE * 256)()
    needed = wintypes.DWORD()
    psapi.EnumProcessModules(handle, mods, ctypes.sizeof(mods), ctypes.byref(needed))
    info = ModuleInfo()
    psapi.GetModuleInformation(handle, mods[0], ctypes.byref(info), ctypes.sizeof(info))
    return handle, int(mods[0]), info.size


def read_image(handle, base, size):
    image = bytearray(size)
    got = ctypes.c_size_t(0)
    pages = 0
    for off in range(0, size, PAGE):
        buf = (ctypes.c_char * PAGE)()
        if k32.ReadProcessMemory(handle, base + off, buf, PAGE, ctypes.byref(got)):
            image[off:off + got.value] = buf.raw[:got.value]
            pages += 1
    return image, pages


def main():
    exe = os.path.abspath(sys.argv[1] if len(sys.argv) > 1 else "hello2_protected.exe")
    rva = int(sys.argv[2], 16) if len(sys.argv) > 2 else 0x10E0

    proc = subprocess.Popen([exe], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                            bufsize=1, universal_newlines=True)
    print(f"[+] {exe} pid={proc.pid}")
    for line in proc.stdout:
        print("   child:", line.strip())
        if line.startswith("RESULT"):
            break

    handle, base, size = main_module(proc.pid)
    print(f"[+] module base={base:#x} size={size:#x}")
    image, pages = read_image(handle, base, size)
    with open("live_dump.bin", "wb") as fh:
        fh.write(image)
    print(f"[+] pages read {pages}, live_dump.bin ({len(image)} bytes)")

    md = Cs(CS_ARCH_X86, CS_MODE_64)
    va = base + rva
    print(f"\n== transform @ {va:#x} (live) ==")
    for i, insn in enumerate(md.disasm(bytes(image[rva:rva + 0x120]), va)):
        raw = " ".join(f"{b:02x}" for b in insn.bytes)
        print(f"  {insn.address:016x}  {raw:<24}  {insn.mnemonic} {insn.op_str}".rstrip())
        if insn.mnemonic == "ret" or i + 1 >= 24:
            break

    try:
        proc.stdin.write("\n")
        proc.stdin.flush()
    except OSError:
        pass
    proc.wait(timeout=5)


if __name__ == "__main__":
    main()