#!/usr/bin/env python3
"""Dump a DRM-DECRYPTED game image from a running (Proton) process.

Why: the on-disk exe is SteamStub-encrypted (.text entropy 8.0). The real UE4
code only exists in memory after the stub decrypts it at launch. This grabs that
memory so we can verify/derive the exact SUWSF byte patterns offline.

Usage:
    # 1. Launch Gundam Breaker 4 (get to the main menu / lobby).
    # 2. Run:
    python3 tools/dump_decrypted.py <ExeName.exe> [pid]
    #    -> writes scratch dump + region metadata, then prints AOB match report.

If you get "Operation not permitted", either run with sudo, or temporarily:
    echo 0 | sudo tee /proc/sys/kernel/yama/ptrace_scope     (reset to 1 after)
"""
import ctypes, ctypes.util, os, re, sys, json, struct, glob

MODULE_HINT = next((a for a in sys.argv[1:] if a.lower().endswith(".exe")), "GB4-Win64-Shipping.exe")
OUTDIR = os.environ.get("UWP_DUMP_DIR", os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "dumps"))

libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)


class iovec(ctypes.Structure):
    _fields_ = [("iov_base", ctypes.c_void_p), ("iov_len", ctypes.c_size_t)]


libc.process_vm_readv.restype = ctypes.c_ssize_t
libc.process_vm_readv.argtypes = [ctypes.c_int, ctypes.POINTER(iovec), ctypes.c_ulong,
                                  ctypes.POINTER(iovec), ctypes.c_ulong, ctypes.c_ulong]


def find_pid():
    if len(sys.argv) > 1 and sys.argv[1].isdigit():
        return int(sys.argv[1])
    for p in glob.glob("/proc/[0-9]*"):
        try:
            with open(f"{p}/cmdline", "rb") as f:
                cl = f.read().replace(b"\0", b" ").decode("utf-8", "replace")
            if MODULE_HINT in cl:
                return int(os.path.basename(p))
        except OSError:
            continue
    return None


def read_mem(pid, addr, size):
    buf = ctypes.create_string_buffer(size)
    local = iovec(ctypes.cast(buf, ctypes.c_void_p), size)
    remote = iovec(ctypes.c_void_p(addr), size)
    n = libc.process_vm_readv(pid, ctypes.byref(local), 1, ctypes.byref(remote), 1, 0)
    if n < 0:
        # fallback: /proc/pid/mem pread
        try:
            with open(f"/proc/{pid}/mem", "rb", 0) as m:
                m.seek(addr)
                return m.read(size)
        except OSError as e:
            raise OSError(f"read fail @0x{addr:x}: {os.strerror(ctypes.get_errno())} / {e}")
    return buf.raw[:n]


def module_regions(pid):
    """Return list of (start,end,perms,path) VMAs backed by the module exe."""
    regs = []
    modpath = None
    with open(f"/proc/{pid}/maps") as f:
        for line in f:
            m = re.match(r"([0-9a-f]+)-([0-9a-f]+) (\S{4}) \S+ \S+ \S+ *(.*)", line)
            if not m:
                continue
            start, end, perms, path = int(m[1], 16), int(m[2], 16), m[3], m[4]
            if MODULE_HINT in path:
                if modpath is None:
                    modpath = path
                regs.append((start, end, perms, path))
    return modpath, regs


def main():
    pid = find_pid()
    if not pid:
        print(f"[!] No running process matching '{MODULE_HINT}'. Launch the game first.")
        sys.exit(1)
    print(f"[+] pid = {pid}")
    modpath, regs = module_regions(pid)
    if not regs:
        print("[!] Module not mapped yet. Reach the main menu, then retry.")
        sys.exit(1)
    base = min(r[0] for r in regs)
    print(f"[+] module base = 0x{base:x}  ({modpath})")
    os.makedirs(OUTDIR, exist_ok=True)
    meta = {"pid": pid, "module": modpath, "base": base, "regions": []}
    blob = bytearray()
    for start, end, perms, path in regs:
        size = end - start
        try:
            data = read_mem(pid, start, size)
        except OSError as e:
            print(f"    skip 0x{start:x}-0x{end:x} {perms}: {e}")
            continue
        rva = start - base
        # pad blob so file offset == rva (sparse-ish, capped)
        if rva >= 0 and rva < 0x20000000:
            if len(blob) < rva:
                blob.extend(b"\x00" * (rva - len(blob)))
            blob[rva:rva + len(data)] = data
        meta["regions"].append({"start": start, "end": end, "rva": rva,
                                 "perms": perms, "size": size, "got": len(data)})
        print(f"    dumped 0x{start:x}-0x{end:x} {perms} rva=0x{rva:x} ({len(data)} bytes)")
    stem = os.path.splitext(os.path.basename(MODULE_HINT))[0]
    dump_path = os.path.join(OUTDIR, f"{stem}_decrypted.bin")
    with open(dump_path, "wb") as f:
        f.write(blob)
    with open(os.path.join(OUTDIR, f"{stem}_meta.json"), "w") as f:
        json.dump(meta, f, indent=2)
    print(f"[+] wrote {dump_path} ({len(blob)} bytes, rva-aligned)")
    print(f"[+] wrote {os.path.join(OUTDIR, stem + '_meta.json')}")

    # quick sanity: is it decrypted? look for the UE4 version string / a UTF-16 hint
    if b"++UE4+Release-4.27" in blob or b"CalculateProjectionMatrix" in blob:
        print("[+] decrypted UE4 code confirmed in dump.")
    print(f"\nNext: python3 tools/aob.py {dump_path} games/<Game>/SUWSF.ini")


if __name__ == "__main__":
    main()