#!/usr/bin/env python3
"""vault_disasm.py — disassemble MuseBank vault bytecode, dump per-selector bodies.
Stdlib only. Focus: reward accounting semantics (order-takes-all vs pro-rata).
"""
import json, os, sys, urllib.request

RPCS = ["https://rpc.ordofi.network", "https://rpc.mainnet.chain.robinhood.com",
        "https://robinhood-rpc.publicnode.com", "https://robinhood.drpc.org"]
H = {"Content-Type": "application/json", "User-Agent": "ARION-vaultdisasm/1.0"}

OPS = {
 0x00:"STOP",0x01:"ADD",0x02:"MUL",0x03:"SUB",0x04:"DIV",0x05:"SDIV",0x06:"MOD",
 0x07:"SMOD",0x08:"ADDMOD",0x09:"MULMOD",0x0a:"EXP",0x0b:"SIGNEXTEND",
 0x10:"LT",0x11:"GT",0x12:"SLT",0x13:"SGT",0x14:"EQ",0x15:"ISZERO",0x16:"AND",
 0x17:"OR",0x18:"XOR",0x19:"NOT",0x1a:"BYTE",0x1b:"SHL",0x1c:"SHR",0x1d:"SAR",
 0x20:"SHA3",
 0x30:"ADDRESS",0x31:"BALANCE",0x32:"ORIGIN",0x33:"CALLER",0x34:"CALLVALUE",
 0x35:"CALLDATALOAD",0x36:"CALLDATASIZE",0x37:"CALLDATACOPY",0x38:"CODESIZE",
 0x39:"CODECOPY",0x3a:"GASPRICE",0x3b:"EXTCODESIZE",0x3c:"EXTCODECOPY",
 0x3d:"RETURNDATASIZE",0x3e:"RETURNDATACOPY",0x3f:"EXTCODEHASH",
 0x40:"BLOCKHASH",0x41:"COINBASE",0x42:"TIMESTAMP",0x43:"NUMBER",0x44:"PREVRANDAO",
 0x45:"GASLIMIT",0x46:"CHAINID",0x47:"SELFBALANCE",0x48:"BASEFEE",
 0x50:"POP",0x51:"MLOAD",0x52:"MSTORE",0x53:"MSTORE8",0x54:"SLOAD",0x55:"SSTORE",
 0x56:"JUMP",0x57:"JUMPI",0x58:"PC",0x59:"MSIZE",0x5a:"GAS",0x5b:"JUMPDEST",
 0x5c:"TLOAD",0x5d:"TSTORE",0x5e:"MCOPY",0x5f:"PUSH0",
 0xf0:"CREATE",0xf1:"CALL",0xf2:"CALLCODE",0xf3:"RETURN",0xf4:"DELEGATECALL",
 0xf5:"CREATE2",0xfa:"STATICCALL",0xfd:"REVERT",0xfe:"INVALID",0xff:"SELFDESTRUCT",
}
for i in range(32): OPS[0x60+i] = f"PUSH{i+1}"
for i in range(16): OPS[0x80+i] = f"DUP{i+1}"
for i in range(16): OPS[0x90+i] = f"SWAP{i+1}"
for i in range(5):  OPS[0xa0+i] = f"LOG{i}"

_i = 0
def rpc(method, params, budget=30):
    global _i
    body = json.dumps({"jsonrpc":"2.0","id":1,"method":method,"params":params}).encode()
    import time
    t0, delay, last = time.time(), 0.4, None
    while True:
        url = RPCS[_i % len(RPCS)]; _i += 1
        try:
            d = json.load(urllib.request.urlopen(
                urllib.request.Request(url, data=body, headers=H), timeout=30))
            if "result" in d: return d["result"]
            last = d.get("error")
        except Exception as e:
            last = e
        if time.time() - t0 > budget: raise RuntimeError(f"rpc: {last}")
        time.sleep(delay); delay = min(delay*1.7, 2.5)

def getcode(addr):
    return bytes.fromhex(rpc("eth_getCode", [addr, "latest"])[2:])

def disasm(code):
    """Return list of (offset, op, imm_hex_or_None)."""
    out = []
    i = 0
    while i < len(code):
        op = code[i]
        imm = None
        if 0x60 <= op <= 0x7f:
            n = op - 0x5f
            imm = code[i+1:i+1+n].hex()
            i += 1 + n
        else:
            i += 1
        out.append((i - (0 if imm is None else 0), op, imm)) if False else out.append((None, op, imm))
    # recompute offsets
    out = []
    i = 0
    while i < len(code):
        op = code[i]
        imm = None
        off = i
        if 0x60 <= op <= 0x7f:
            n = op - 0x5f
            imm = code[i+1:i+1+n].hex()
            i += 1 + n
        else:
            i += 1
        out.append((off, op, imm))
    return out

def fmt(code):
    ins = disasm(code)
    jumpdests = {o for o, op, _ in ins if op == 0x5b}
    # find PUSH4 constants that look like selectors + JUMPI dispatch
    print(f"code len {len(code)}B, {len(jumpdests)} jumpdests")
    # dump all push4 constants (likely selectors)
    sels = [imm for _, op, imm in ins if op == 0x63 and imm and len(imm) <= 8]
    print("PUSH4 consts:", sorted(set(sels)))
    # dispatch targets: PUSH2/PUSH3 followed by JUMPI
    for idx, (o, op, imm) in enumerate(ins):
        if op in (0x61, 0x62) and idx+1 < len(ins) and ins[idx+1][1] == 0x57:
            # look back for push4
            back = [x for x in ins[max(0,idx-4):idx] if x[1] == 0x63]
            tag = back[-1][2] if back else ""
            print(f"dispatch sel 0x{tag or '?'} -> dest 0x{int(imm,16):04x}")
    return ins, jumpdests

def dump_region(ins, start, end=None):
    for o, op, imm in ins:
        if o < start: continue
        if end is not None and o >= end: break
        name = OPS.get(op, f"OP{op:02x}")
        extra = ""
        if imm: extra = " 0x" + imm
        print(f"{o:05x}: {name}{extra}")

if __name__ == "__main__":
    addr = sys.argv[1] if len(sys.argv) > 1 else "0xB5FBDAE914C696467D6D9E6de0c476d8d8402F45"
    code = getcode(addr)
    ins, jd = fmt(code)
    if len(sys.argv) > 2:
        start = int(sys.argv[2], 16)
        end = int(sys.argv[3], 16) if len(sys.argv) > 3 else None
        dump_region(ins, start, end)
    else:
        # dump everything compactly to file
        out = []
        for o, op, imm in ins:
            out.append(f"{o:05x}: {OPS.get(op, f'OP{op:02x}')}" + (f" 0x{imm}" if imm else ""))
        open("/tmp/vault_disasm.txt", "w").write("\n".join(out))
        print("wrote /tmp/vault_disasm.txt")