"""Capture every `buffer` library call made by the inner VM.

The VM resolves buffer methods dynamically from the `buffer` table it captures
at startup, so wrapping the table before stage2 runs records all reads,
including the string-extraction path that pulls user constants out of the
unread payload tail.
"""

import re
from pathlib import Path

STAGE2 = Path("/root/luad/samples/boulder/stage2.lua")
PAYLOAD = Path("/root/luad/samples/boulder/payload.bin")
OUT = Path("/tmp/opencode/bufprobe.luau")


def main() -> None:
    stage2 = STAGE2.read_text()
    payload = PAYLOAD.read_bytes()

    harness = r'''
local __print = print
local __n = 0
local __log = {}
local function __hex(s)
  return (s:gsub(".", function(c) return string.format("%02x", string.byte(c)) end))
end
local __real = {}
local __copy = {}
for k, v in pairs(buffer) do __copy[k] = v end
for k, v in pairs(buffer) do
  if type(v) == "function" and k ~= "fromstring" and k ~= "tostring" then
    __real[k] = v
    __copy[k] = function(...)
      local a = table.pack(...)
      local r = table.pack(v(...))
      __n = __n + 1
      if __n <= 400 then
        local parts = {}
        for i = 1, math.min(a.n, 4) do
          local x = a[i]
          if type(x) == "number" then parts[#parts+1] = string.format("%.6g", x)
          elseif type(x) == "buffer" then parts[#parts+1] = "buf"
          elseif type(x) == "string" then parts[#parts+1] = "s:" .. __hex(x:sub(1, 24))
          else parts[#parts+1] = type(x) end
        end
        local rets = {}
        for i = 1, math.min(r.n, 3) do
          local x = r[i]
          if type(x) == "string" then rets[#rets+1] = "s:" .. __hex(x:sub(1, 60))
          elseif type(x) == "number" then rets[#rets+1] = string.format("%.6g", x)
          else rets[#rets+1] = type(x) end
        end
        __print("BUF\t" .. tostring(k) .. "\t" .. table.concat(parts, ",") .. "\t->\t" .. table.concat(rets, ","))
      end
      return table.unpack(r, 1, r.n)
    end
  end
end
buffer = __copy
print("BUFWRAP", #__log)
''' 

    # permissive environment (same shape as the working runner)
    shim = r'''
local __print2 = print
local function __proxy(name)
  local t = {}
  return setmetatable(t, {
    __index = function(self, k)
      local v = __proxy(name .. "." .. tostring(k)); rawset(self, k, v); return v
    end,
    __call = function(_, ...) return __proxy(name .. "()") end,
    __tostring = function() return "<" .. name .. ">" end,
    __len = function() return 0 end,
    __add = function() return 0 end, __sub = function() return 0 end,
    __mul = function() return 0 end, __div = function() return 0 end,
    __mod = function() return 0 end, __pow = function() return 0 end,
    __unm = function() return 0 end,
    __eq = function() return false end, __lt = function() return false end, __le = function() return false end,
    __concat = function() return "" end,
  })
end
game = __proxy("game")
workspace = __proxy("workspace")
Instance = {new = function(c) return __proxy("Instance:" .. tostring(c)) end}
Enum = setmetatable({}, {__index = function(t, k) local v = __proxy("E." .. k); rawset(t, k, v); return v end})
task = {wait = function() end, spawn = function() end}
typeof = function(v) return type(v) end
tick = function() return 0 end
'''

    lit = '"' + "".join("\\x%02x" % b for b in payload) + '"'
    assert stage2.startswith("local R=...;")
    stage2 = "local R=buffer.fromstring(" + lit + ");" + stage2[len("local R=...;"):]

    body = (
        shim
        + harness
        + "\nlocal __ok, __f = pcall(function()\n"
        + stage2
        + "\nend)\n"
        + '__print("LOAD", __ok, type(__f))\n'
        + 'if __ok then local a, b = pcall(__f) __print("RUN", a, tostring(b)) end\n'
    )
    OUT.write_text(body)
    print("wrote", OUT, len(body))


if __name__ == "__main__":
    main()