OpenMFOR/extra/decomp_logic.py

import xdis
import ast

from collections import defaultdict

def main():
    pyc_filename = 'Randomizer.pyc'

    (
        version_tuple,
        timestamp,
        magic_int,
        topco,
        is_pypy,
        source_size,
        sip_hash,
    ) = xdis.load_module(pyc_filename)

    opc = xdis.get_opcode(version_tuple, is_pypy)

    for co in topco.co_consts:
        if getattr(co, 'co_name', '') == 'update_requirements':
            bc = xdis.Bytecode(co, opc)

            ops = set()
            for i in bc:
                ops.add(i.opname)
                #if i.opname == 'EXTENDED_ARG' or 'JUMP' in i.opname:
                #    print(i)
            #print('\n'.join(sorted(ops)))

            dump(bc)

def mkbool(op, a1, a2):
    if op == 'and':
        expr = []
        for x in a1, a2:
            if isinstance(x, ast.BoolOp) and isinstance(x.op, ast.And):
                expr.extend(x.values)
            else:
                expr.append(x)
        return ast.BoolOp(ast.And(), expr)
    elif op == 'or':
        expr = []
        for x in a1, a2:
            if isinstance(x, ast.BoolOp) and isinstance(x.op, ast.Or):
                expr.extend(x.values)
            else:
                expr.append(x)
        return ast.BoolOp(ast.Or(), expr)
    else:
        return (op, a1, a2)

def dump(bytecode):
    op_stack = []
    expr_stack = []
    value_stack = []
    reduce = defaultdict(int)
    push = defaultdict(int)
    rindex = dict()
    retired = set()
    i = 0
    for inst in bytecode:
        addr = inst.offset
        op = inst.opname
        arg = inst.argval
        #print(addr, op, reduce[addr], push[addr], rindex.get(addr, ""))
        if reduce[addr]:
            if push[addr]:
                if reduce[addr] != push[addr]:
                    print("stack mismatch at %s: %d != %d", addr, reduce[addr], push[addr])
            stash = []
            assert len(expr_stack) == len(op_stack), (expr_stack, op_stack)
            if push[addr]:
                expr = value_stack.pop()
                j,k = i,i
                i += 1
            else:
                stash = [op_stack.pop()]
                j,k,expr = expr_stack.pop()
            r = rindex[addr]
            while not j <= r <= k:
                j2,k2,a2 = j,k,expr
                j1,k1,a1 = expr_stack.pop()
                xop = op_stack.pop()
                expr = mkbool(xop, a1, a2)
                j, k = j1, k2 # min(j1,j2), max(k1, k2)
            expr_stack.append((j,k,expr))
            op_stack.extend(stash)
        if push[addr]:
            _,_,top = expr_stack.pop()
            #print(top)
            #print(ast.dump(top, indent=4))
            #print(ast.unparse(top))
            value_stack.append(top)
        if op == 'LOAD_NAME' or op == 'LOAD_GLOBAL' or op == 'LOAD_FAST':
            value_stack.append(ast.Name(arg))
        elif op == 'LOAD_CONST':
            value_stack.append(ast.Constant(arg))
        elif op == 'EXTENDED_ARG':
            pass
        elif op == 'STORE_NAME':
            sys.stdout.flush()
            assert not value_stack, value_stack
            assert not expr_stack, expr_stack
        elif op == 'POP_JUMP_IF_TRUE':
            dest = arg
            val = value_stack.pop()
            expr_stack.append((i,i,val))
            op_stack.append('or')
            reduce[dest] += 1
            rindex.setdefault(dest, i)
            i += 1
        elif op == 'POP_JUMP_IF_FALSE':
            dest = arg
            val = value_stack.pop()
            expr_stack.append((i,i,val))
            op_stack.append('and')
            reduce[dest] += 1
            rindex.setdefault(dest, i)
            i += 1
        elif op == 'JUMP_IF_TRUE_OR_POP':
            dest = arg
            val = value_stack.pop()
            expr_stack.append((i,i,val))
            op_stack.append('or')
            reduce[dest] += 1
            push[dest] += 1
            rindex.setdefault(dest, i)
            i += 1
        elif op == 'JUMP_IF_FALSE_OR_POP':
            dest = arg
            val = value_stack.pop()
            expr_stack.append((i,i,val))
            op_stack.append('and')
            reduce[dest] += 1
            push[dest] += 1
            rindex.setdefault(dest, i)
            i += 1
        elif op == 'COMPARE_OP':
            a2 = value_stack.pop()
            a1 = value_stack.pop()
            if arg == '>':
                cop = ast.Gt()
            elif arg == '>=':
                cop = ast.GtE()
            else:
                print('unknown comparator', arg)
                cop = ''
            if cop:
                value_stack.append(ast.Compare(a1, ops=[cop], comparators=[a2]))
        elif op == 'BUILD_CONST_KEY_MAP':
            key_expr = value_stack.pop()
            if len(value_stack) < arg:
                print("error: BUILD_CONST_KEY_MAP: not enough expressions on stack (%d < %d)" % (len(value_stack), arg))
            vals = value_stack[-arg:]
            value_stack[-arg:] = []
            assert isinstance(key_expr, ast.Constant)
            assert len(key_expr.value) == arg
            l = max(len(str(k)) for k in key_expr.value)
            for k, v in zip(key_expr.value, vals):
                print("    %s: %s%s," % (k, " "*(l-len(str(k))), ast.unparse(v)))
            m = ast.Dict(map(ast.Constant, key_expr.value), vals)
            value_stack.append(m)
        elif op == 'STORE_ATTR':
            obj = value_stack.pop()
            val = value_stack.pop()
            name = arg
            #print(obj, name, val)
            x = ast.Assign([ast.Attribute(val, name)], val, lineno=0)
            #print(ast.unparse(x))
        else:
            print('unknown op', op, arg)

    print(expr_stack)
    print(value_stack)


main()