// Capture rekordbox's .pdb writes with Frida (run on the WINDOWS host).
//
// Hooks at the ntdll layer (NtCreateFile / NtWriteFile / NtClose) so it catches
// file IO no matter whether the app calls kernel32, kernelbase, or the CRT.
// Logs EVERY file opened for writing (so you can see temp names / the real
// filenames even if the .pdb filter is wrong), reconstructs files whose path
// matches SAVE_RE byte-for-byte, and saves them to OUT_DIR via Frida's File API.
// Prints a backtrace the first time each saved file is written -> the
// module+offset of the serializer to open in Ghidra.
//
// Usage (Windows PowerShell, rekordbox running):
//     mkdir C:\Users\ntr\code\maple\maple-py\captures      # OUT_DIR must exist
//     frida-ps | Select-String rekordbox                   # find all PIDs
//     frida -n rekordbox.exe -l tools\rbx_pdb_capture.js   # or -p <pid>
// then trigger "Export to USB". If you see [w-open] lines but the export file
// isn't among them, widen SAVE_RE. If you see NOTHING, attach to another PID.
//
// Set VERBOSE = true to log every write-open path; false logs only SAVE_RE hits.

var OUT_DIR = "C:\\Users\\ntr\\code\\maple\\maple-py\\captures\\"; // must exist; Frida won't mkdir
var SAVE_RE = /export.*\.pdb|\.pdb$|PIONEER/i; // which opened files to reconstruct + save
var VERBOSE = true;                            // log ALL write-opens, not just SAVE_RE

var PID = Process.id;
var TAG = "[" + PID + "] ";
var DEBUG = true;                                   // log every interesting open + all mapping activity
var DEBUG_RE = /export|\.pdb|PIONEER|rekordbox|Contents|USBANLZ/i; // what DEBUG opens to print
var ntdll = Process.getModuleByName("ntdll.dll");
function nt(name) { return ntdll.getExportByName(name); }

// Resolve + confirm a hook target so we can see in the console what attached.
function resolve(name) {
    var a = ntdll.findExportByName(name);
    console.log(TAG + "hook " + name + (a ? " -> " + a : " MISSING"));
    return a;
}


// handle (string) -> { path, pos, end, chunks:[{pos,ab}], save, traced }
var files = {};
var counters = {};

function baseName(path) {
    var parts = path.replace(/\\/g, "/").replace(/\/+$/, "").split("/");
    return (parts[parts.length - 1] || "capture.bin").replace(/[^A-Za-z0-9._-]/g, "");
}

// Read the wide path out of an OBJECT_ATTRIBUTES* (x64 layout).
function readObjName(pObjAttr) {
    if (pObjAttr.isNull()) return null;
    var pUni = pObjAttr.add(16).readPointer();       // OBJECT_ATTRIBUTES.ObjectName
    if (pUni.isNull()) return null;
    var len = pUni.readU16();                          // UNICODE_STRING.Length (bytes)
    if (!len) return null;
    var buf = pUni.add(8).readPointer();              // UNICODE_STRING.Buffer
    if (buf.isNull()) return null;
    return buf.readUtf16String(len / 2);
}

// NtCreateFile / NtOpenFile: args[0]=PHANDLE out, [1]=ACCESS_MASK, [2]=OBJECT_ATTRIBUTES*
function hookOpen(name) {
    var addr = resolve(name);
    if (!addr) return;
    Interceptor.attach(addr, {
        onEnter: function (args) {
            this.pHandle = args[0];
            this.access = args[1].toUInt32();
            this.path = readObjName(args[2]);
        },
        onLeave: function (retval) {
            if (!this.path) return;
            var WRITE = 0x40000000 /*GENERIC_WRITE*/ | 0x10000000 /*GENERIC_ALL*/ |
                        0x2 /*FILE_WRITE_DATA*/ | 0x4 /*FILE_APPEND_DATA*/;
            var isWrite = (this.access & WRITE) !== 0;
            var save = SAVE_RE.test(this.path);
            // DEBUG: show every interesting open (incl. read-only) + its access mask,
            // so we can see if/how export.pdb is opened even when our filters miss it.
            if (DEBUG && DEBUG_RE.test(this.path))
                console.log(TAG + "[open] access=0x" + this.access.toString(16)
                    + (isWrite ? " W" : " r") + (save ? " *SAVE" : "") + "  " + this.path);
            if (!isWrite && !save) return;
            var handleStr = this.pHandle.readPointer().toString();
            files[handleStr] = { path: this.path, pos: 0, end: 0, chunks: [], save: save, traced: false };
            if (!DEBUG && (save || VERBOSE)) console.log(TAG + "[w-open] " + this.path);
        },
    });
}
hookOpen("NtCreateFile");
hookOpen("NtOpenFile");

// NtWriteFile: [0]=handle [5]=buffer [6]=length [7]=PLARGE_INTEGER ByteOffset
Interceptor.attach(resolve("NtWriteFile"), {
    onEnter: function (args) {
        var f = files[args[0].toString()];
        if (!f) return;
        var n = args[6].toUInt32();
        if (n <= 0) return;

        var pos = f.pos;
        var pOff = args[7];
        if (!pOff.isNull()) {
            var hi = pOff.add(4).readU32();
            if (hi !== 0xffffffff) pos = hi * 0x100000000 + pOff.readU32(); // skip sentinels
        }

        if (f.save) {
            f.chunks.push({ pos: pos, ab: args[5].readByteArray(n) });
            if (!f.traced) {
                f.traced = true;
                console.log(TAG + "[trace] writer for " + f.path + ":");
                Thread.backtrace(this.context, Backtracer.ACCURATE)
                    .slice(0, 14)
                    .forEach(function (a) {
                        // module+offset maps straight to Ghidra: static = 0x140000000 + offset
                        var m = Process.findModuleByAddress(a);
                        var loc = m ? (m.name + "+0x" + a.sub(m.base).toString(16)) : a.toString();
                        console.log("          " + loc + "   " + DebugSymbol.fromAddress(a));
                    });
            }
        }
        f.pos = pos + n;
        if (f.pos > f.end) f.end = f.pos;
    },
});

// Shared: write a captured buffer to OUT_DIR with a PID-tagged, de-duped name.
function saveBuffer(path, arrayBuffer, size) {
    var name = baseName(path);
    var c = counters[name] || 0;
    counters[name] = c + 1;
    var dest = OUT_DIR + "p" + PID + "_" + name + (c === 0 ? "" : "." + c);
    var fh = new File(dest, "wb");
    fh.write(arrayBuffer);
    fh.close();
    console.log(TAG + "[SAVED] " + dest + "  (" + size + " bytes)");
}

// NtClose: reconstruct + flush saved files written via NtWriteFile
function flush(handleStr) {
    var f = files[handleStr];
    if (!f) return;
    delete files[handleStr];
    if (!f.save || f.end === 0) return;
    var out = new Uint8Array(f.end);
    for (var i = 0; i < f.chunks.length; i++) out.set(new Uint8Array(f.chunks[i].ab), f.chunks[i].pos);
    saveBuffer(f.path, out.buffer, f.end);
}
Interceptor.attach(resolve("NtClose"), { onEnter: function (args) { flush(args[0].toString()); } });

// ---- memory-mapped writes (the .pdb is likely built via a file mapping) ----
// CreateFileMapping/MapViewOfFile funnel through these ntdll calls. We link
// section -> backing file (NtCreateSection takes the file handle we already
// track), record the mapped view, and snapshot its memory when it's unmapped.
var sections = {};       // sectionHandle -> { path, save, maxSize }
var views = {};          // baseAddr string -> { base, size, path, save, maxSize }

function hookSections() {
    var cs = resolve("NtCreateSection");
    if (cs) Interceptor.attach(cs, {
        onEnter: function (args) {
            this.pSection = args[0];
            this.pMax = args[3];                 // PLARGE_INTEGER MaximumSize (may be null)
            this.fileHandle = args[6].toString();// HANDLE FileHandle
        },
        onLeave: function (ret) {
            if (ret.toInt32() !== 0) return;     // not STATUS_SUCCESS
            var f = files[this.fileHandle];
            if (!f) return;                      // only file-backed sections we track
            var maxSize = (this.pMax && !this.pMax.isNull()) ? this.pMax.readU32() : 0;
            sections[this.pSection.readPointer().toString()] = { path: f.path, save: f.save, maxSize: maxSize };
        },
    });

    var mv = resolve("NtMapViewOfSection");
    if (mv) Interceptor.attach(mv, {
        onEnter: function (args) {
            this.section = args[0].toString();
            this.pBase = args[2];                // PVOID* BaseAddress
            this.pViewSize = args[6];            // PSIZE_T ViewSize
        },
        onLeave: function (ret) {
            if (ret.toInt32() < 0) return;       // negative NTSTATUS = error
            var s = sections[this.section];
            if (!s) return;
            var base = this.pBase.readPointer();
            var size = (this.pViewSize && !this.pViewSize.isNull()) ? this.pViewSize.readU32() : 0;
            views[base.toString()] = { base: base, size: size, path: s.path, save: s.save, maxSize: s.maxSize };
            if (s.save || VERBOSE) console.log(TAG + "[mmap] " + s.path + " @ " + base + " view=" + size + " max=" + s.maxSize);
        },
    });

    function captureView(base) {
        var v = views[base.toString()];
        if (!v) return;
        delete views[base.toString()];
        if (!v.save) return;
        var size = v.maxSize || v.size;
        if (!size) { var r = Process.findRangeByAddress(base); size = r ? r.size : 0; }
        if (!size) return;
        try { saveBuffer(v.path, Memory.readByteArray(base, size), size); }
        catch (e) { console.log(TAG + "[mmap] read failed @ " + base + ": " + e); }
    }

    var um = ntdll.findExportByName("NtUnmapViewOfSection");
    if (um) Interceptor.attach(um, { onEnter: function (args) { captureView(args[1]); } });
}
hookSections();

console.log(TAG + "ntdll + mmap hooks installed -- trigger 'Export to USB' in rekordbox now");