"""Frida JS agent injected into the target process.

setup() probes the target's pymalloc internals once (py_version, pool_size).
collect() scans all rw- ranges for pool_header signatures — no GIL, no
           knowledge of _obmalloc_state required, works on stripped builds.
readPool() returns raw block data for the block-level hex view.
reprBlock() acquires the GIL, bumps refcount, and calls PyObject_Repr safely.
"""

from __future__ import annotations

_AGENT_JS = r"""
'use strict';

// ── cached state (set during setup) ──────────────────────────────────────────
let _poolSize  = 16384;
let _arenaSize = 1048576;

// ── symbol / NativeFunction helpers (pattern from zap-allocator) ─────────────
function sym(name) {
    if (typeof Module.findGlobalExportByName === 'function') {
        const p = Module.findGlobalExportByName(name);
        if (p && !p.isNull()) return p;
    }
    for (const mod of Process.enumerateModules()) {
        try {
            const p = mod.getExportByName(name);
            if (p && !p.isNull()) return p;
        } catch (e) {}
    }
    return null;
}

function nfn(name, ret, args) {
    const p = sym(name);
    if (!p) throw new Error('symbol not found: ' + name);
    return new NativeFunction(p, ret, args);
}

function makeSetupRunner() {
    const ensure  = nfn('PyGILState_Ensure',  'int',  []);
    const release = nfn('PyGILState_Release', 'void', ['int']);
    const runStr  = nfn('PyRun_SimpleString', 'int',  ['pointer']);
    return function run(code) {
        const state = ensure();
        const ret   = runStr(Memory.allocUtf8String(code));
        release(state);
        return ret;
    };
}

// Read one Python integer attribute from __main__ via PyLong_AsVoidPtr.
// Returns a NativePointer whose integer value equals the attribute,
// or null pointer if missing / not an int.
function readIntAttr(run, attrName) {
    const ensure    = nfn('PyGILState_Ensure',       'int',     []);
    const release   = nfn('PyGILState_Release',      'void',    ['int']);
    const addMod    = nfn('PyImport_AddModule',      'pointer', ['pointer']);
    const getAttr   = nfn('PyObject_GetAttrString',  'pointer', ['pointer', 'pointer']);
    const asVoidPtr = nfn('PyLong_AsVoidPtr',        'pointer', ['pointer']);
    const decRef    = nfn('Py_DecRef',               'void',    ['pointer']);

    const cMain = Memory.allocUtf8String('__main__');
    const cAttr = Memory.allocUtf8String(attrName);

    const state   = ensure();
    const mainMod = addMod(cMain);
    const obj     = getAttr(mainMod, cAttr);
    let result    = ptr('0x0');
    if (!obj.isNull()) {
        const p = asVoidPtr(obj);
        if (!p.isNull()) result = p;
        decRef(obj);
    }
    release(state);
    return result;
}

// ── free-list walker (no GIL needed) ─────────────────────────────────────────
function walkFreeList(poolPtr, maxBlocks) {
    const addrs = [];
    let cur;
    try { cur = poolPtr.add(8).readPointer(); } catch (e) { return addrs; }
    let limit = 0;
    while (cur && !cur.isNull() && limit < maxBlocks) {
        addrs.push(cur.toString());          // hex string preserves 64-bit precision
        try { cur = cur.readPointer(); } catch (e) { break; }
        limit++;
    }
    return addrs;
}

// ── pool-header scan over ALL rw- ranges ─────────────────────────────────────
// Validation uses the exact maxnextoffset = POOL_SIZE - block_size invariant.
// Scans in 4 MiB chunks to bound peak memory usage from readByteArray.
function scanAllPools() {
    const POOL_SIZE     = _poolSize;
    const POOL_MASK     = POOL_SIZE - 1;
    const POOL_OVERHEAD = 48;
    const CHUNK_SIZE    = 4 * 1024 * 1024;

    const pools = [];

    for (const range of Process.enumerateRanges('rw-')) {
        if (range.size < POOL_SIZE) continue;

        const rangeBase = parseInt(range.base.toString());
        const rangeEnd  = rangeBase + range.size;

        for (let chunkStart = rangeBase; chunkStart < rangeEnd; chunkStart += CHUNK_SIZE) {
            const chunkSize = Math.min(CHUNK_SIZE, rangeEnd - chunkStart);
            let buf;
            try {
                buf = ptr('0x' + chunkStart.toString(16)).readByteArray(chunkSize);
            } catch (e) { continue; }

            const dv = new DataView(buf);
            // First pool-aligned offset within this chunk
            const firstOff = (POOL_SIZE - (chunkStart & POOL_MASK)) & POOL_MASK;

            for (let off = firstOff; off + POOL_SIZE <= chunkSize; off += POOL_SIZE) {
                const szidx = dv.getUint32(off + 36, true);
                if (szidx >= 32) continue;

                const blockSize     = (szidx + 1) * 16;
                const maxnextoffset = dv.getUint32(off + 44, true);
                // Exact invariant from CPython obmalloc.c: POOL_SIZE - block_size
                if (maxnextoffset !== POOL_SIZE - blockSize) continue;

                const refCount   = dv.getUint32(off + 0,  true);
                const nextoffset = dv.getUint32(off + 40, true);
                const totalBlocks = (POOL_SIZE - POOL_OVERHEAD) / blockSize;

                if (refCount > totalBlocks) continue;
                if (nextoffset > POOL_SIZE)  continue;

                const paddr  = chunkStart + off;

                const arenaIndex = dv.getUint32(off + 32, true);

                pools.push({
                    address:      paddr.toString(16),
                    arenaIndex,
                    szidx,
                    refCount,
                    nextoffset,
                    maxnextoffset,
                    freeAddrs:    [],   // deferred to readPool() to keep collect() payload small
                });
            }
        }
    }
    return pools;
}

// ── single-arena pool scan ────────────────────────────────────────────────────
// Scans exactly one arena_size window starting at arenaBase (must be
// arena_size-aligned).  Returns the same pool-record shape as scanAllPools()
// but without freeAddrs.  Errors on individual pool reads are silently skipped.
function scanOneArena(arenaBase) {
    const POOL_SIZE     = _poolSize;
    const POOL_MASK     = POOL_SIZE - 1;
    const POOL_OVERHEAD = 48;

    let buf;
    try {
        buf = ptr('0x' + arenaBase.toString(16)).readByteArray(_arenaSize);
    } catch (e) { return []; }

    const dv       = new DataView(buf);
    const firstOff = (POOL_SIZE - (arenaBase & POOL_MASK)) & POOL_MASK;
    const pools    = [];

    for (let off = firstOff; off + POOL_SIZE <= _arenaSize; off += POOL_SIZE) {
        const szidx = dv.getUint32(off + 36, true);
        if (szidx >= 32) continue;

        const blockSize     = (szidx + 1) * 16;
        const maxnextoffset = dv.getUint32(off + 44, true);
        if (maxnextoffset !== POOL_SIZE - blockSize) continue;

        const refCount    = dv.getUint32(off +  0, true);
        const nextoffset  = dv.getUint32(off + 40, true);
        const totalBlocks = (POOL_SIZE - POOL_OVERHEAD) / blockSize;
        if (refCount > totalBlocks || nextoffset > POOL_SIZE) continue;

        const arenaIndex = dv.getUint32(off + 32, true);
        pools.push({
            address: (arenaBase + off).toString(16),
            arenaIndex,
            szidx,
            refCount,
            nextoffset,
            maxnextoffset,
        });
    }
    return pools;
}

// ── RPC exports ───────────────────────────────────────────────────────────────
rpc.exports = {

    setup: function () {
        let run;
        try { run = makeSetupRunner(); }
        catch (e) { return {ok: false, error: e.message}; }

        // Retrieve Python version — no struct offset guessing needed.
        const ret = run(`
import sys as _sys, __main__ as _m
_m._mw_pymaj = _sys.version_info.major
_m._mw_pymin = _sys.version_info.minor
`);
        if (ret !== 0) return {ok: false, error: 'setup Python snippet failed'};

        const pyMaj = parseInt(readIntAttr(run, '_mw_pymaj').toString());
        const pyMin = parseInt(readIntAttr(run, '_mw_pymin').toString());

        return {ok: true, poolSize: _poolSize, arenaSize: _arenaSize,
                pyMajor: pyMaj, pyMinor: pyMin};
    },

    collect: function () {
        try {
            return {ok: true, pools: scanAllPools()};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // collectFiltered: like collect() but restricted to ranges whose size is
    // exactly a small multiple of ARENA_SIZE (the heuristic that identifies
    // CPython mmap arenas on 64-bit systems).  One RPC call, so no per-arena
    // round-trip overhead.  freeAddrs are omitted (deferred to readPool).
    collectFiltered: function () {
        const ARENA_SIZE = _arenaSize;
        const pools = [];
        try {
            for (const range of Process.enumerateRanges('rw-')) {
                if (range.size < ARENA_SIZE) continue;
                // Use floor division so that a run of N arenas that got merged
                // with a small adjacent allocation (e.g. 3 MiB + 4 KiB) is still
                // fully scanned.  Any sub-ARENA_SIZE residual at the end is
                // ignored — it can't contain a complete pool.
                const rangeBase = parseInt(range.base.toString());
                const count     = Math.floor(range.size / ARENA_SIZE);
                for (let i = 0; i < count; i++) {
                    for (const p of scanOneArena(rangeBase + i * ARENA_SIZE))
                        pools.push(p);
                }
            }
            return {ok: true, pools};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // Refresh pool header + free list + raw bytes for the block view.
    // NOTE: must be named readPool (camelCase) to match Frida's Python→JS mapping.
    // raw bytes are returned as a base64 string because ArrayBuffer inside an RPC
    // dict doesn't serialize reliably across the Frida bridge.
    readPool: function (poolAddrHex) {
        const p = ptr('0x' + poolAddrHex);
        try {
            const szidx         = p.add(36).readU32();
            const refCount      = p.add(0).readU32();
            const nextoffset    = p.add(40).readU32();
            const maxnextoffset = p.add(44).readU32();
            const blockSize     = (szidx + 1) * 16;
            const totalBlocks   = (_poolSize - 48) / blockSize;
            const maxBl         = Math.ceil(totalBlocks) + 2;
            const freeAddrs     = walkFreeList(p, maxBl);
            const rawBuf        = p.readByteArray(_poolSize);
            // hex-encode so it survives the JSON RPC bridge
            const rawHex        = Array.from(new Uint8Array(rawBuf),
                                      b => b.toString(16).padStart(2, '0')).join('');
            return {ok: true, szidx, refCount, nextoffset, maxnextoffset, freeAddrs, rawHex};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // Safely repr a live block.
    // 1) Read ob_type + tp_name without GIL (memory reads only).
    // 2) Sanity-check refcount.
    // 3) Acquire GIL, Py_IncRef, PyObject_Repr, Py_DecRef, release GIL.
    // NOTE: must be named reprBlock (camelCase) to match Frida's Python→JS mapping.
    reprBlock: function (blockAddrHex) {
        const blockPtr = ptr('0x' + blockAddrHex);

        // For non-GC objects ob_type is at block+8; for GC-tracked objects (dict, list,
        // frame, …) a 16-byte PyGC_Head precedes the PyObject so ob_type is at block+24.
        let objPtr = blockPtr;
        let typeName = '?';
        for (const obTypeOff of [8, 24]) {
            try {
                const obType = blockPtr.add(obTypeOff).readPointer();
                if (obType.isNull()) continue;
                const tpNamePtr = obType.add(24).readPointer();
                const name = tpNamePtr.readUtf8String();
                if (name && name.length > 0 && !name.startsWith('\x00')) {
                    typeName = name;
                    if (obTypeOff === 24) objPtr = blockPtr.add(16);
                    break;
                }
            } catch (_) {}
        }

        let refcnt = 0;
        try { refcnt = objPtr.add(0).readU32(); } catch (_) {}
        if (refcnt === 0 || refcnt > 0x7fffffff) {
            return {ok: false, typeName, error: 'implausible refcount — block may be free'};
        }

        const ensure   = nfn('PyGILState_Ensure',  'int',     []);
        const release  = nfn('PyGILState_Release', 'void',    ['int']);
        const incRef   = nfn('Py_IncRef',           'void',    ['pointer']);
        const decRef   = nfn('Py_DecRef',           'void',    ['pointer']);
        const reprFn   = nfn('PyObject_Repr',       'pointer', ['pointer']);
        const asUtf8   = nfn('PyUnicode_AsUTF8',    'pointer', ['pointer']);
        const clearErr = nfn('PyErr_Clear',         'void',    []);

        const state = ensure();
        incRef(objPtr);
        let reprStr = null;
        let errMsg  = null;
        try {
            const reprObj = reprFn(objPtr);
            if (!reprObj.isNull()) {
                const cstr = asUtf8(reprObj);
                reprStr = cstr.isNull() ? null : cstr.readUtf8String();
                decRef(reprObj);
            } else {
                clearErr();
                errMsg = 'PyObject_Repr returned NULL';
            }
        } catch (e) {
            errMsg = e.message;
        } finally {
            decRef(objPtr);
            release(state);
        }

        return reprStr !== null
            ? {ok: true, typeName, reprStr}
            : {ok: false, typeName, error: errMsg || 'repr failed'};
    },

    // Bulk-resolve ob_type pointers → tp_name strings (no GIL needed).
    // Input: array of hex strings (ob_type pointer values from raw block bytes).
    // Returns: object mapping each input hex string to its tp_name (or '?').
    // NOTE: must be named resolveTypeNames (camelCase) to match Frida's mapping.
    resolveTypeNames: function (typeAddrHexList) {
        const result = {};
        for (const addrHex of typeAddrHexList) {
            try {
                const typePtr = ptr('0x' + addrHex);
                const tpNamePtr = typePtr.add(24).readPointer();
                result[addrHex] = tpNamePtr.readUtf8String();
            } catch (e) {
                result[addrHex] = '?';
            }
        }
        return result;
    },

    // Return the memory protection string (e.g. 'rw-', 'rwx') for the range
    // containing the given address, or null if the address is unmapped.
    // NOTE: must be named getRangeProtection (camelCase) for Frida's Python→JS mapping.
    getRangeProtection: function (addrHex) {
        try {
            const range = Process.findRangeByAddress(ptr('0x' + addrHex));
            return range ? range.protection : null;
        } catch (_) {
            return null;
        }
    },

    // ── streaming collect API ─────────────────────────────────────────────────
    // listArenas: enumerate rw- ranges of exactly ARENA_SIZE — no memory reads.
    // CPython allocates each arena with a single mmap(NULL, ARENA_SIZE, …) call,
    // so those regions appear as distinct ARENA_SIZE-sized rw- ranges.
    // Adjacent arenas coalesced by the OS (size == N*ARENA_SIZE) are split.
    // arenaIndex here is a synthetic sequence number; actual pool_header.arenaIndex
    // values come from listPoolsInArena().
    // NOTE: must be named listArenas (camelCase) for Frida's Python→JS mapping.
    listArenas: function () {
        try {
            const ARENA_SIZE = _arenaSize;
            const arenas = [];
            let idx = 0;
            for (const range of Process.enumerateRanges('rw-')) {
                if (range.size < ARENA_SIZE) continue;
                const base  = parseInt(range.base.toString());
                const count = Math.floor(range.size / ARENA_SIZE);
                for (let i = 0; i < count; i++) {
                    arenas.push({
                        arenaIndex: idx++,
                        base: (base + i * ARENA_SIZE).toString(16),
                        poolCount: -1,   // filled in by listPoolsInArena
                    });
                }
            }
            return {ok: true, arenas};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // listPoolsInArena: scan one arena_size window, return pool headers only.
    // arenaBaseHex must be the arena_size-aligned base returned by listArenas.
    // NOTE: must be named listPoolsInArena (camelCase) for Frida's mapping.
    listPoolsInArena: function (arenaBaseHex) {
        try {
            const arenaBase = parseInt('0x' + arenaBaseHex, 16);
            return {ok: true, pools: scanOneArena(arenaBase)};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // listPoolSubset: read raw bytes for [firstAddr, halfOpenLastAddr) and return
    // as a single hex string.  Python splits by blockSize*2 chars per block.
    // Half-open range; size must be <= 65536 bytes.
    // NOTE: must be named listPoolSubset (camelCase) for Frida's mapping.
    listPoolSubset: function (firstAddrHex, halfOpenLastAddrHex) {
        try {
            const first = parseInt('0x' + firstAddrHex,        16);
            const last  = parseInt('0x' + halfOpenLastAddrHex, 16);
            const size  = last - first;
            if (size <= 0 || size > 65536)
                return {ok: false, error: 'range out of bounds: ' + size};
            const buf = ptr('0x' + firstAddrHex).readByteArray(size);
            const hex = Array.from(new Uint8Array(buf),
                            b => b.toString(16).padStart(2, '0')).join('');
            return {ok: true, hex};
        } catch (e) {
            return {ok: false, error: e.message};
        }
    },

    // Bulk-resolve addresses → {module, offset, symbol} or null.
    // Uses Process.findModuleByAddress (no GIL, cross-platform) for module lookup
    // and DebugSymbol.fromAddress (best-effort) for symbol names.
    // NOTE: must be named symbolizeAddresses (camelCase) for Frida's Python→JS mapping.
    symbolizeAddresses: function(addrHexList) {
        const result = {};
        for (const addrHex of addrHexList) {
            try {
                const p = ptr('0x' + addrHex);
                // DebugSymbol.fromAddress uses dladdr() which covers __DATA segments
                // correctly on macOS (unlike Module.size which only covers __TEXT).
                const sym = DebugSymbol.fromAddress(p);
                const modName = sym && sym.moduleName && sym.moduleName.length > 0
                    ? sym.moduleName : null;
                if (!modName) { result[addrHex] = null; continue; }

                // Compute byte offset from the module's load address.
                let offset = 0;
                const mod = Process.findModuleByName(modName);
                if (mod) offset = parseInt(p.sub(mod.base).toString());

                // sym.name may be an address string ("0x...") when no symbol is found.
                const rawName = sym.name || '';
                const symbolName = (rawName.length > 0 && !rawName.startsWith('0x'))
                    ? rawName : null;

                result[addrHex] = { module: modName, path: mod ? mod.path : '', offset: offset, symbol: symbolName };
            } catch (_) {
                result[addrHex] = null;
            }
        }
        return result;
    },
};
"""
