#!/usr/bin/env python3
"""Probe 10: identify which state-table slots are buffer readers, and what
they return at a known offset. Also probe the closure factory with a
one-instruction synthetic prototype."""
STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
OUT = '/tmp/opencode/dump_run10.luau'

SHIM = r'''
local __realprint = print
tick = function() return 0 end
time = function() return 0 end
warn = function() end
typeof = function(v) return type(v) end
local __HOST = getfenv()
__ENV = setmetatable({}, {
	__index = function(t, k) return __HOST[k] end,
	__newindex = function(t, k, v) __HOST[k] = v end,
})
__STEPON = false
function __STEP(Q) return false end

local function V(v)
	local t = type(v)
	if t == "string" then return string.format("%q", v:sub(1, 24))
	elseif t == "number" then return string.format("%.10g", v)
	elseif t == "table" then return "tbl"
	elseif t == "boolean" then return tostring(v)
	else return t end
end

__PROBE_READERS = function(G, R)
	local slots = {5, 7, 8, 9, 11, 12, 13, 14, 15, 20, 26, 29, 42, 44, 45, 46, 47, 49, 50, 52, 53}
	for _, s in ipairs(slots) do
		local f = rawget(G, s)
		if type(f) == "function" then
			local results = {}
			for _, args in ipairs({ { R, 0 }, { R, 0, 4 }, { R, 0, 1 }, { R } }) do
				local ok, a, b = pcall(f, table.unpack(args))
				results[#results + 1] = (ok and "ok:" or "err:") .. "(" .. V(a) .. (b ~= nil and ("," .. V(b)) or "") .. ")"
			end
			__realprint("RD", s, table.concat(results, " "))
		end
	end
end

__PROBE_FACTORY = function(G)
	local function mk(op, a, b, c, frame)
		return { [1] = 0, [2] = {}, [3] = { op }, [4] = { b }, [5] = { c },
			[6] = frame or 16, [7] = {}, [8] = {}, [9] = {}, [10] = { a }, [11] = {} }
	end
	local factory = G[55]
	__realprint("FACT", type(factory))
	for _, spec in ipairs({ {173, 0, 0, 0}, {15, 0, 0, 0}, {180, 0, 0, 3}, {81, 0, 0, 3} }) do
		local op, a, b, c = spec[1], spec[2], spec[3], spec[4]
		local ok, err = pcall(function()
			local fn = factory(mk(op, a, b, c), {})
			local ok2, res = pcall(fn)
			__realprint("RUN", op, ok2, V(res), type(fn))
		end)
		if not ok then __realprint("RUNERR", op, tostring(err)) end
	end
end
'''

def main():
    stage2 = open(STAGE2, 'rb').read().decode('latin-1')
    payload = open(PAYLOAD, 'rb').read()
    rlit = '"' + ''.join('\\x%02x' % b for b in payload) + '"'
    stage2 = stage2.replace('local R=...;', 'local R=buffer.fromstring(' + rlit + ');', 1)
    needle = 'y,A,m,P,e=W:KS(e,P,A,G,y);'
    assert stage2.count(needle) == 1
    stage2 = stage2.replace(needle, needle + '__PROBE_READERS(G,R);__PROBE_FACTORY(G);', 1)
    out = SHIM + 'local __ok,__err=pcall(function()\nlocal __fn=function()\n' + stage2 + '\nend\nsetfenv(__fn,__ENV)\nreturn __fn()\nend)\n'
    out += '__realprint("EXEC", __ok, tostring(__err))\n'
    open(OUT, 'w').write(out)
    print("wrote", OUT, len(out))

if __name__ == '__main__':
    main()