import ctypes
import ctypes.wintypes as wt
import struct
import sys
from dataclasses import dataclass

kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
advapi32 = ctypes.WinDLL("advapi32", use_last_error=True)

# -----------------------------
# Windows constants
# -----------------------------
TH32CS_SNAPPROCESS = 0x00000002

PROCESS_QUERY_INFORMATION = 0x0400
PROCESS_VM_READ = 0x0010
PROCESS_QUERY_LIMITED_INFORMATION = 0x1000

MEM_COMMIT = 0x1000

MEM_IMAGE = 0x1000000
MEM_MAPPED = 0x40000
MEM_PRIVATE = 0x20000

PAGE_NOACCESS = 0x01
PAGE_READONLY = 0x02
PAGE_READWRITE = 0x04
PAGE_WRITECOPY = 0x08
PAGE_EXECUTE = 0x10
PAGE_EXECUTE_READ = 0x20
PAGE_EXECUTE_READWRITE = 0x40
PAGE_EXECUTE_WRITECOPY = 0x80
PAGE_GUARD = 0x100
PAGE_NOCACHE = 0x200
PAGE_WRITECOMBINE = 0x400

READABLE_MASK = (
    PAGE_READONLY | PAGE_READWRITE | PAGE_WRITECOPY |
    PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY
)

# SeDebugPrivilege (opcjonalnie)
SE_PRIVILEGE_ENABLED = 0x00000002
TOKEN_ADJUST_PRIVILEGES = 0x20
TOKEN_QUERY = 0x8

# -----------------------------
# Structures
# -----------------------------
class PROCESSENTRY32(ctypes.Structure):
    _fields_ = [
        ("dwSize", wt.DWORD),
        ("cntUsage", wt.DWORD),
        ("th32ProcessID", wt.DWORD),
        ("th32DefaultHeapID", ctypes.c_void_p),
        ("th32ModuleID", wt.DWORD),
        ("cntThreads", wt.DWORD),
        ("th32ParentProcessID", wt.DWORD),
        ("pcPriClassBase", wt.LONG),
        ("dwFlags", wt.DWORD),
        ("szExeFile", wt.WCHAR * 260),
    ]

# MEMORY_BASIC_INFORMATION for VirtualQueryEx:
# This layout is correct for 32-bit target from a 64-bit Python too, because VirtualQueryEx
# returns pointers sized to the caller, but BaseAddress/RegionSize are pointer-sized.
# We'll store them in ctypes.c_size_t and then cast.
class MEMORY_BASIC_INFORMATION(ctypes.Structure):
    _fields_ = [
        ("BaseAddress", ctypes.c_void_p),
        ("AllocationBase", ctypes.c_void_p),
        ("AllocationProtect", wt.DWORD),
        ("RegionSize", ctypes.c_size_t),
        ("State", wt.DWORD),
        ("Protect", wt.DWORD),
        ("Type", wt.DWORD),
    ]

class SYSTEM_INFO(ctypes.Structure):
    _fields_ = [
        ("wProcessorArchitecture", wt.WORD),
        ("wReserved", wt.WORD),
        ("dwPageSize", wt.DWORD),
        ("lpMinimumApplicationAddress", ctypes.c_void_p),
        ("lpMaximumApplicationAddress", ctypes.c_void_p),
        ("dwActiveProcessorMask", ctypes.c_void_p),
        ("dwNumberOfProcessors", wt.DWORD),
        ("dwProcessorType", wt.DWORD),
        ("dwAllocationGranularity", wt.DWORD),
        ("wProcessorLevel", wt.WORD),
        ("wProcessorRevision", wt.WORD),
    ]

class LUID(ctypes.Structure):
    _fields_ = [("LowPart", wt.DWORD), ("HighPart", wt.LONG)]

class LUID_AND_ATTRIBUTES(ctypes.Structure):
    _fields_ = [("Luid", LUID), ("Attributes", wt.DWORD)]

class TOKEN_PRIVILEGES(ctypes.Structure):
    _fields_ = [("PrivilegeCount", wt.DWORD),
                ("Privileges", LUID_AND_ATTRIBUTES * 1)]

# -----------------------------
# API prototypes
# -----------------------------
kernel32.CreateToolhelp32Snapshot.argtypes = [wt.DWORD, wt.DWORD]
kernel32.CreateToolhelp32Snapshot.restype = wt.HANDLE

kernel32.Process32FirstW.argtypes = [wt.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
kernel32.Process32FirstW.restype = wt.BOOL

kernel32.Process32NextW.argtypes = [wt.HANDLE, ctypes.POINTER(PROCESSENTRY32)]
kernel32.Process32NextW.restype = wt.BOOL

kernel32.OpenProcess.argtypes = [wt.DWORD, wt.BOOL, wt.DWORD]
kernel32.OpenProcess.restype = wt.HANDLE

kernel32.CloseHandle.argtypes = [wt.HANDLE]
kernel32.CloseHandle.restype = wt.BOOL

kernel32.VirtualQueryEx.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.POINTER(MEMORY_BASIC_INFORMATION), ctypes.c_size_t]
kernel32.VirtualQueryEx.restype = ctypes.c_size_t

kernel32.ReadProcessMemory.argtypes = [wt.HANDLE, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)]
kernel32.ReadProcessMemory.restype = wt.BOOL

kernel32.GetSystemInfo.argtypes = [ctypes.POINTER(SYSTEM_INFO)]
kernel32.GetSystemInfo.restype = None

advapi32.OpenProcessToken.argtypes = [wt.HANDLE, wt.DWORD, ctypes.POINTER(wt.HANDLE)]
advapi32.OpenProcessToken.restype = wt.BOOL

advapi32.LookupPrivilegeValueW.argtypes = [wt.LPCWSTR, wt.LPCWSTR, ctypes.POINTER(LUID)]
advapi32.LookupPrivilegeValueW.restype = wt.BOOL

advapi32.AdjustTokenPrivileges.argtypes = [wt.HANDLE, wt.BOOL, ctypes.POINTER(TOKEN_PRIVILEGES),
                                          wt.DWORD, ctypes.c_void_p, ctypes.c_void_p]
advapi32.AdjustTokenPrivileges.restype = wt.BOOL

kernel32.GetCurrentProcess.restype = wt.HANDLE

# -----------------------------
# Helpers
# -----------------------------
def winerr(msg: str):
    err = ctypes.get_last_error()
    raise OSError(err, f"{msg} (WinError={err})")

def enable_debug_privilege():
    hToken = wt.HANDLE()
    if not advapi32.OpenProcessToken(kernel32.GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, ctypes.byref(hToken)):
        return False
    try:
        luid = LUID()
        if not advapi32.LookupPrivilegeValueW(None, "SeDebugPrivilege", ctypes.byref(luid)):
            return False
        tp = TOKEN_PRIVILEGES(1, (LUID_AND_ATTRIBUTES(luid, SE_PRIVILEGE_ENABLED),))
        if not advapi32.AdjustTokenPrivileges(hToken, False, ctypes.byref(tp), 0, None, None):
            return False
        # AdjustTokenPrivileges can succeed but still not enable; GetLastError==ERROR_NOT_ALL_ASSIGNED
        if ctypes.get_last_error() == 1300:  # ERROR_NOT_ALL_ASSIGNED
            return False
        return True
    finally:
        kernel32.CloseHandle(hToken)

def find_pid_by_name(exe_name: str) -> int:
    snap = kernel32.CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
    if snap == wt.HANDLE(-1).value:
        winerr("CreateToolhelp32Snapshot failed")

    try:
        pe = PROCESSENTRY32()
        pe.dwSize = ctypes.sizeof(PROCESSENTRY32)
        if not kernel32.Process32FirstW(snap, ctypes.byref(pe)):
            winerr("Process32FirstW failed")

        exe_name_l = exe_name.lower()
        while True:
            if pe.szExeFile.lower() == exe_name_l:
                return int(pe.th32ProcessID)
            if not kernel32.Process32NextW(snap, ctypes.byref(pe)):
                break
        return 0
    finally:
        kernel32.CloseHandle(snap)

def is_readable_page(prot: int) -> bool:
    if prot & PAGE_GUARD:
        return False
    if prot & PAGE_NOACCESS:
        return False
    if prot & READABLE_MASK:
        return True
    return False

@dataclass
class RegionIdx:
    base: int
    size: int
    protect: int
    type: int
    fileOffset: int

def dump_process_memory(pid: int, mem_path: str, idx_path: str, chunk_size: int = 1 << 20):
    access = PROCESS_VM_READ | PROCESS_QUERY_INFORMATION
    hProc = kernel32.OpenProcess(access, False, pid)
    if not hProc:
        # fallback for some processes
        hProc = kernel32.OpenProcess(PROCESS_VM_READ | PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
    if not hProc:
        winerr("OpenProcess failed")

    index = []
    try:
        si = SYSTEM_INFO()
        kernel32.GetSystemInfo(ctypes.byref(si))
        p = ctypes.c_void_p(si.lpMinimumApplicationAddress).value
        max_addr = ctypes.c_void_p(si.lpMaximumApplicationAddress).value

        mbi = MEMORY_BASIC_INFORMATION()

        with open(mem_path, "wb") as fmem:
            while p < max_addr:
                got = kernel32.VirtualQueryEx(hProc, ctypes.c_void_p(p), ctypes.byref(mbi), ctypes.sizeof(mbi))
                if not got:
                    # jump by a page to avoid infinite loop on weird edges
                    p += 0x1000
                    continue

                base = ctypes.c_void_p(mbi.BaseAddress).value
                rsize = int(mbi.RegionSize)
                prot = int(mbi.Protect)
                state = int(mbi.State)
                rtype = int(mbi.Type)

                if state == MEM_COMMIT and is_readable_page(prot) and rsize > 0:
                    file_off = fmem.tell()
                    written = 0

                    cur = base
                    left = rsize

                    buf = (ctypes.c_ubyte * min(chunk_size, rsize))()
                    while left > 0:
                        to_read = min(left, chunk_size)
                        bytes_read = ctypes.c_size_t(0)
                        ok = kernel32.ReadProcessMemory(
                            hProc,
                            ctypes.c_void_p(cur),
                            ctypes.byref(buf),
                            to_read,
                            ctypes.byref(bytes_read),
                        )
                        n = int(bytes_read.value)
                        if n > 0:
                            fmem.write(bytes(buf[:n]))
                            written += n
                            cur += n
                            left -= n
                        # jeśli RPM fail albo short read — kończymy region (jak w Twoim C++)
                        if (not ok) or (n < to_read):
                            break

                    if written > 0:
                        # Twój format trzyma base jako u32
                        index.append(RegionIdx(
                            base=base & 0xFFFFFFFF,
                            size=written & 0xFFFFFFFF,
                            protect=prot & 0xFFFFFFFF,
                            type=rtype & 0xFFFFFFFF,
                            fileOffset=file_off & 0xFFFFFFFFFFFFFFFF,
                        ))

                # następny region
                p = base + rsize

        # zapis indexu
        with open(idx_path, "wb") as fidx:
            hdr = struct.pack("<IIII", 0x584D454D, 1, len(index), 0)  # 'MEMX'
            fidx.write(hdr)
            for ri in index:
                fidx.write(struct.pack("<IIIIQ", ri.base, ri.size, ri.protect, ri.type, ri.fileOffset))

        print(f"[+] DONE: wrote {len(index)} regions -> {mem_path} (+ {idx_path})")

    finally:
        kernel32.CloseHandle(hProc)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python dump_mem.py nemo.exe [dump.mem dump.idx]")
        sys.exit(1)

    exe = sys.argv[1]
    mem_path = sys.argv[2] if len(sys.argv) >= 3 else "dump.mem"
    idx_path = sys.argv[3] if len(sys.argv) >= 4 else "dump.idx"

    dbg = enable_debug_privilege()
    print(f"[*] SeDebugPrivilege: {'enabled' if dbg else 'not enabled (ok, maybe not needed)'}")

    pid = find_pid_by_name(exe)
    if not pid:
        print(f"[-] Process not found: {exe}")
        sys.exit(2)

    print(f"[*] PID for {exe}: {pid}")
    dump_process_memory(pid, mem_path, idx_path)