#!/usr/bin/env python3
"""Probe 6: identify which proto slot array the tagC==5 (import) path targets."""
STAGE2 = '/root/luad/samples/boulder/stage2.lua'
PAYLOAD = '/root/luad/samples/boulder/payload.bin'
OUT = '/tmp/opencode/dump_run6.luau'

SHIM = r'''
local __realprint = print
tick = function() return 0 end
time = function() return 0 end
warn = function() return end
typeof = function(v) return type(v) end
__PROTO_ID = 0
__CUR = nil
__CK = function(x)
	local cur = __CUR
	if type(cur) ~= "table" or __NLOGGED > 5 then return x end
	__NLOGGED = __NLOGGED + 1
	local hit = "?"
	for s = 1, 11 do
		if type(cur[s]) == "table" and cur[s] == x then hit = tostring(s) end
	end
	__realprint("CKT", __PROTO_ID, hit)
	return x
end
__NLOGGED = 0
'''

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)
    needle1 = 'W:n7(m,h,l,P);'
    assert stage2.count(needle1) == 1
    stage2 = stage2.replace(needle1, needle1 + '__CUR=m;', 1)
    needle2 = 'W:cS(j,l,e,t,d,A)'
    assert stage2.count(needle2) == 1, stage2.count(needle2)
    stage2 = stage2.replace(needle2, 'W:cS(j,__CK(l),e,t,d,A)', 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()