#!/usr/bin/env python3
"""Generate a Luau runner that dumps the VM prototype array right after the
load pass (KS), before the TS patch loop.  Experiment harness."""
import sys, re

STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
OUT = '/tmp/opencode/dump_run.luau'

ENC = r'''
local __HEX = {}
for __i = 0, 255 do __HEX[__i] = string.format("%02x", __i) end
local function __hex(s)
	local out = {}
	local n = #s
	local i = 1
	while i <= n do
		local chunk = {}
		local lim = i + 4095
		if lim > n then lim = n end
		local j = i
		while j <= lim do
			chunk[#chunk + 1] = __HEX[string.byte(s, j)]
			j += 1
		end
		out[#out + 1] = table.concat(chunk)
		i = lim + 1
	end
	return table.concat(out)
end
local function __cmpk(a, b)
	local ta, tb = type(a), type(b)
	if ta ~= tb then return ta < tb end
	if ta == "number" or ta == "string" then return a < b end
	return tostring(a) < tostring(b)
end
local function __enc(v, out)
	local t = type(v)
	if t == "number" then
		out[#out + 1] = "\1"
		out[#out + 1] = string.pack("<d", v)
	elseif t == "string" then
		out[#out + 1] = "\2"
		out[#out + 1] = string.pack("<I4", #v)
		out[#out + 1] = v
	elseif t == "table" then
		local keys = {}
		for k in pairs(v) do keys[#keys + 1] = k end
		table.sort(keys, __cmpk)
		out[#out + 1] = "\4"
		out[#out + 1] = string.pack("<I4", #keys)
		for o = 1, #keys do
			local k = keys[o]
			__enc(k, out)
			__enc(rawget(v, k), out)
		end
	elseif t == "boolean" then
		out[#out + 1] = v and "\3\1" or "\3\0"
	elseif t == "function" then
		out[#out + 1] = "\6"
	elseif t == "nil" then
		out[#out + 1] = "\5"
	else
		out[#out + 1] = "\7"
		local s = tostring(v)
		out[#out + 1] = string.pack("<I4", #s)
		out[#out + 1] = s
	end
end
local function __encval(v)
	local out = {}
	__enc(v, out)
	return table.concat(out)
end
__DUMPED = false
function __DUMP(y, A, m, P, e, G)
	if __DUMPED then return end
	__DUMPED = true
	local n = #y
	__realprint("DUMP_BEGIN", n, type(A), tostring(m), type(P), type(e), type(G))
	for i = 1, n do
		local p = y[i]
		__realprint("PROTO", i, type(p))
		if type(p) == "table" then
			local keys = {}
			for k in pairs(p) do keys[#keys + 1] = k end
			table.sort(keys, __cmpk)
			local ks = {}
			for o = 1, #keys do ks[#ks + 1] = tostring(keys[o]) .. ":" .. type(p[keys[o]]) end
			__realprint("KEYS", i, table.concat(ks, ","))
			for o = 1, #keys do
				local k = keys[o]
				__realprint("SLOT", i, tostring(k), __hex(__encval(rawget(p, k))))
			end
		end
	end
	__realprint("DUMP_END", n)
end
'''

SHIM = r'''
local __realprint = print
local __MISS = {}
local __MISSOK = {}
local __type, __tostring, __rawget, __setmetatable, __getmetatable = type, tostring, rawget, setmetatable, getmetatable
local __env = getfenv()
local __mt = __getmetatable(__env)
local __prev = __type(__mt) == "table" and __mt.__index or nil
__setmetatable(__env, {
	__index = function(t, k)
		local v
		if __prev ~= nil then
			if __type(__prev) == "function" then
				v = __prev(t, k)
			elseif __type(__prev) == "table" then
				v = __prev[k]
			end
		end
		if v == nil and not __MISSOK[__tostring(k)] then
			__MISSOK[__tostring(k)] = true
			__MISS[#__MISS + 1] = __tostring(k)
		end
		return v
	end,
})
tick = function() return 0 end
time = function() return 0 end
warn = function() end
typeof = function(v) return type(v) end
__MISSING_REPORT = function()
	__realprint("MISSING", #__MISS, table.concat(__MISS, ","))
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) + '"'
    assert stage2.startswith('local R=...;')
    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 + '__DUMP(y,A,m,P,e,G);', 1)
    out = SHIM + ENC + 'local __ok,__err=pcall(function()\n' + stage2 + '\nend)\n'
    out += '__realprint("EXEC", __ok, tostring(__err))\n__MISSING_REPORT()\n'
    open(OUT, 'w').write(out)
    print("wrote", OUT, len(out), "payload", len(payload))

if __name__ == '__main__':
    main()