#!/usr/bin/env python3
import re, struct, sys, json

def b_xs(state):
    state &= 0xffffffff
    state ^= (state << 13) & 0xffffffff
    state ^= (state >> 17)
    state ^= (state << 5) & 0xffffffff
    return state & 0xffffffff

def decode_graph(vmsrc):
    m = re.search(r'local __G="([0-9a-fA-F]+)"', vmsrc)
    if not m:
        raise SystemExit('no __G')
    raw = bytes.fromhex(m.group(1))
    assert raw[:4] == b'SGB2', raw[:4]
    version, key_raw, regbase, flag = raw[4], raw[5], raw[6], raw[7]
    assert version == 1
    key = key_raw ^ ((regbase * 37 + 0xA5) & 0xFF)
    state = (key * 257 + regbase * 17 + 0x5EC20) & 0xffffffff
    plain = bytearray()
    for b in raw[8:]:
        state = b_xs(state)
        plain.append(b ^ (state & 0xFF))
    d = bytes(plain)
    p = 0
    def u8():
        nonlocal p
        v = d[p]; p += 1; return v
    def u32():
        nonlocal p
        v = struct.unpack_from('<I', d, p)[0]; p += 4; return v
    def f64():
        nonlocal p
        v = struct.unpack_from('<d', d, p)[0]; p += 8; return v
    def bstr():
        nonlocal p
        n = u32(); v = d[p:p+n]; p += n; return v
    def bval():
        nonlocal p
        t = u8()
        if t == 0: return None
        if t == 1: return f64()
        if t == 2: return bstr()
        if t == 3:
            n = u32(); return [bval() for _ in range(n)]
        if t == 4: return u8() != 0
        raise ValueError('bval tag %d at %d' % (t, p))
    opkey = u32(); root = u32(); nproto = u32()
    protos = {}
    for _ in range(nproto):
        pid = u32(); params = u32(); vararg = u8(); stack = u32(); nbox = u32()
        boxed = [u32() for _ in range(nbox)]
        maxpc = u32()
        code = []
        for _pc in range(maxpc):
            opcode = u32(); nf = u8(); ins = {'op': opcode}
            names = {0:'A',1:'B',2:'C',3:'K1',4:'K2',5:'K3'}
            for _ in range(nf):
                fid = u8(); val = bval()
                ins[names.get(fid, 'f%d' % fid)] = val
            code.append(ins)
        protos[pid] = {'params':params,'vararg':vararg,'stack':stack,'boxed':boxed,'code':code}
    nop = u32(); opmap = {}
    for _ in range(nop):
        e = u32(); o = u32(); opmap[e] = o
    hs = u8(); sbox_inv = None
    if hs == 1:
        sbox = [u8() for _ in range(256)]
        sbox_inv = [u8() for _ in range(256)]
    want = u32()
    h = 0x13579BDF
    for i in range(len(d) - 4):
        h = ((h * 33) & 0xffffffff) ^ d[i]
    print('hash check:', h == want, 'graph bytes', len(d))
    return {'key':key,'regbase':regbase,'flag':flag,'opkey':opkey,'root':root,
            'protos':protos,'opmap':opmap,'sbox_inv':sbox_inv}

def decode_string(s, key, sbox_inv):
    if not s or s[0] != 1:
        return s
    seed = (57 + key * 17) % 256
    out = bytearray()
    for b in s[1:]:
        if sbox_inv is not None:
            b = sbox_inv[b]
        out.append(b ^ seed ^ 92)
        seed = (141 * seed + 29) % 256
    return bytes(out)

def main():
    vmsrc = open(sys.argv[1], 'rb').read().decode('latin-1')
    g = decode_graph(vmsrc)
    print('key', g['key'], 'regbase', g['regbase'], 'flag', g['flag'], 'opkey', g['opkey'], 'root', g['root'])
    print('protos', len(g['protos']), 'total ins', sum(len(p['code']) for p in g['protos'].values()))
    strings = {}
    for pid, p in g['protos'].items():
        for ins in p['code']:
            for k in ('K1','K2','K3'):
                v = ins.get(k)
                if isinstance(v, bytes):
                    dec = decode_string(v, g['key'], g['sbox_inv'])
                    try:
                        s = dec.decode('utf-8')
                    except Exception:
                        s = repr(dec)
                    strings[s] = strings.get(s, 0) + 1
    print('unique strings', len(strings))
    out = sorted(strings.items(), key=lambda kv: (-kv[1], kv[0]))
    with open(sys.argv[2], 'w') as f:
        for s, c in out:
            f.write('%5d %r\n' % (c, s))
    for s, c in out[:120]:
        print('%5d %r' % (c, s[:160]))

if __name__ == "__main__":
    main()
