#!/usr/bin/env python3
"""Second probe: dump the TS patch list (G[2]) with target slots resolved by
identity, then dump y again after the patch loop."""
import sys

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

ENV = 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
function __DUMPY(tag, y)
	local n = #y
	__realprint("PROTOS_BEGIN", tag, n)
	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)
			for o = 1, #keys do
				local k = keys[o]
				__realprint("SLOT", i, tostring(k), __hex(__encval(rawget(p, k))))
			end
		end
	end
	__realprint("PROTOS_END", tag)
end
function __DUMPPATCH(y, G)
	local n = #G[2]
	__realprint("PATCH_BEGIN", n)
	local id = {}
	for i = 1, #y do
		local p = y[i]
		if type(p) == "table" then
			for s = 1, 11 do
				local t = p[s]
				if type(t) == "table" then id[t] = tostring(i) .. "." .. tostring(s) end
			end
		end
	end
	local stats = {}
	local nres = 0
	for b = 1, n, 3 do
		local target, key, val = G[2][b], G[2][b + 1], G[2][b + 2]
		local name = id[target]
		if not name then nres = nres + 1; name = "?" .. type(target) end
		local k = tostring(key)
		local vt = type(y[val])
		local vv = vt == "number" and string.format("%.8g", y[val]) or vt == "string" and string.format("%q", y[val]:sub(1, 20)) or vt
		local sig = name .. "|" .. vt
		stats[sig] = (stats[sig] or 0) + 1
		if b <= 3000 then
			__realprint("PATCH", b, name, k, tostring(val), vv)
		end
	end
	__realprint("PATCH_UNRESOLVED", nres)
	local ks = {}
	for k in pairs(stats) do ks[#ks + 1] = k end
	table.sort(ks)
	for i = 1, #ks do __realprint("PSTAT", ks[i], stats[ks[i]]) end
	__realprint("PATCH_END")
end
'''

SHIM = open('/tmp/opencode/shim.luau').read() if False else r'''
local __realprint = print
tick = function() return 0 end
time = function() return 0 end
warn = function() end
typeof = function(v) return type(v) 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 + '__DUMPY("before",y);__DUMPPATCH(y,G);', 1)
    needle2 = 'for b=1.0,#G[2],3.0 do A=W:TS(b,A,y,G);end'
    assert stage2.count(needle2) == 1, stage2.count(needle2)
    stage2 = stage2.replace(needle2, needle2 + ';__DUMPY("after",y)', 1)
    out = SHIM + ENV + 'local __ok,__err=pcall(function()\n' + stage2 + '\nend)\n'
    out += '__realprint("EXEC", __ok, tostring(__err))\n'
    open(OUT, 'w').write(out)
    print("wrote", OUT, len(out))

if __name__ == '__main__':
    main()