#!/usr/bin/env python3
"""Fourth probe: log every instruction word (all protos) and the reader's
resolved constant pool."""
STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
OUT = '/tmp/opencode/dump_run4.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
__PROTO_ID = 0
__HEX = {}
for __i = 0, 255 do __HEX[__i] = string.format("%02x", __i) end
local function __hex(s)
	local out = {}
	for j = 1, #s do out[#out + 1] = __HEX[string.byte(s, j)] end
	return table.concat(out)
end
__RAW = function(pid, ins, x0, x1, op, y1)
	__realprint("RAW", pid, ins, x0, x1, op, y1)
end
__DUMPCONST = function(rd)
	local t = rd[22]
	local n = 0
	if type(t) == "table" then for _ in pairs(t) do n = n + 1 end end
	__realprint("CONST_BEGIN", n)
	if type(t) == "table" then
		local keys = {}
		for k in pairs(t) do keys[#keys + 1] = k end
		table.sort(keys, function(a, b) return a < b end)
		for i = 1, #keys do
			local k = keys[i]
			local v = rawget(t, k)
			local tv = type(v)
			if tv == "string" then
				__realprint("CONST", k, "s", #v, __hex(v))
			elseif tv == "number" then
				__realprint("CONST", k, "n", string.format("%.17g", v))
			elseif tv == "table" then
				local a, b = v[1], v[2]
				local sa = type(a) == "string" and __hex(a) or tostring(a)
				local sb = type(b) == "string" and __hex(b) or tostring(b)
				__realprint("CONST", k, "t", tostring(a == nil and "nil" or type(a)), sa, tostring(type(b)), sb)
			else
				__realprint("CONST", k, tv, tostring(v))
			end
		end
	end
	__realprint("CONST_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)
    needle0 = '(l)[56]=(function()local T,O,m,A,e,P;'
    assert stage2.count(needle0) == 1
    stage2 = stage2.replace(needle0, '(l)[56]=(function()__PROTO_ID=__PROTO_ID+1;local T,O,m,A,e,P;', 1)
    needle = 'T,m,Q,S,A,I,s,O=W:QS(S,y,d,b,h,Q,P,M,l,m,s,A,I,B,e,q,t,o,p);'
    assert stage2.count(needle) == 1
    stage2 = stage2.replace(needle, needle + '__RAW(__PROTO_ID,e,y,q,M,B);', 1)
    needle2 = 'for W=1.0,m,1 do(j)[W]=l[56]();end;return j,G,nil,T,u;end'
    assert stage2.count(needle2) == 1, stage2.count(needle2)
    stage2 = stage2.replace(needle2, 'for W=1.0,m,1 do(j)[W]=l[56]();end;__DUMPCONST(l);return j,G,nil,T,u;end', 1)
    out = SHIM + '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()