Xentax forum attachments archive/xentax.7z/14707/Spellforce050818_2.7z/luadec.py

Python Source | 2018-08-01 | 5.3 KB | 177 lines

"""

Primitive Lua disassembler.

The Lua scripts are interesting for their data as well as their logic.  For
example, sql_building.lua maps building ids to models.  Therefore, it would
be useful to write a basic decompiler for scripts that do nothing more than
returning a table.

The disassembler can be improved easily by adding line numbers, labels for
jump instructions, string tables, and so on.

"""

import io
import struct


class LuaOpType:
    NONE = 0
    U = 1
    S = 2
    AB = 3
    J = 4
    K = 5
    L = 6
    N = 7


OPS = (
    ('END', LuaOpType.NONE),
    ('RETURN', LuaOpType.U),
    ('CALL', LuaOpType.AB),
    ('TAILCALL', LuaOpType.AB),
    ('PUSHNIL', LuaOpType.U),
    ('POP', LuaOpType.U),
    ('PUSHINT', LuaOpType.S),
    ('PUSHSTRING', LuaOpType.K),
    ('PUSHNUM', LuaOpType.N),
    ('PUSHNEGNUM', LuaOpType.N),
    ('PUSHUPVALUE', LuaOpType.U),
    ('GETLOCAL', LuaOpType.L),
    ('GETGLOBAL', LuaOpType.K),
    ('GETTABLE', LuaOpType.NONE),
    ('GETDOTTED', LuaOpType.K),
    ('GETINDEXED', LuaOpType.L),
    ('PUSHSELF', LuaOpType.K),
    ('CREATETABLE', LuaOpType.U),
    ('SETLOCAL', LuaOpType.L),
    ('SETGLOBAL', LuaOpType.K),
    ('SETTABLE', LuaOpType.AB),
    ('SETLIST', LuaOpType.AB),
    ('SETMAP', LuaOpType.U),
    ('ADD', LuaOpType.NONE),
    ('ADDI', LuaOpType.S),
    ('SUB', LuaOpType.NONE),
    ('MULT', LuaOpType.NONE),
    ('DIV', LuaOpType.NONE),
    ('POW', LuaOpType.NONE),
    ('CONCAT', LuaOpType.U),
    ('MINUS', LuaOpType.NONE),
    ('NOT', LuaOpType.NONE),
    ('JMPNE', LuaOpType.J),
    ('JMPEQ', LuaOpType.J),
    ('JMPLT', LuaOpType.J),
    ('JMPLE', LuaOpType.J),
    ('JMPGT', LuaOpType.J),
    ('JMPGE', LuaOpType.J),
    ('JMPT', LuaOpType.J),
    ('JMPF', LuaOpType.J),
    ('JMPONT', LuaOpType.J),
    ('JMPONF', LuaOpType.J),
    ('JMP', LuaOpType.J),
    ('PUSHNILJMP', LuaOpType.NONE),
    ('FORPREP', LuaOpType.J),
    ('FORLOOP', LuaOpType.J),
    ('LFORPREP', LuaOpType.J),
    ('LFORLOOP', LuaOpType.J),
    ('CLOSURE', LuaOpType.AB)
)


def read_string(ins):
    length, = struct.unpack('<L', ins.read(4))
    return ins.read(length).decode('latin-1')


def read_lua_number(ins):
    return struct.unpack('<d', ins.read(8))[0]


def read_uint32(ins):
    return struct.unpack('<L', ins.read(4))[0]


def read_array(ins, elemtype):
    count, = struct.unpack('<L', ins.read(4))
    return [elemtype(ins) for _ in range(count)]


class LuaFileHeader:
    def __init__(self, ins):
        self.magic, self.version, self.endianness, self.size_int, \
            self.size_size_t, self.size_instruction,\
            self.test_size_instruction, self.test_size_op, self.test_size_b, \
            self.size_number, self.test_fp_number = \
            struct.unpack('<4s9Bd', ins.read(21))


class LuaLocalVariable:
    def __init__(self, ins):
        self.name = read_string(ins)
        self.start_pc, self.end_pc = struct.unpack('<2L', ins.read(8))


class LuaCodeChunk:
    def __init__(self, ins):
        self.source = read_string(ins)
        self.line_number, self.num_params, isvarargs, self.max_stack_size = \
            struct.unpack('<2LbL', ins.read(13))
        self.isvarargs = isvarargs == 1
        self.local_vars = read_array(ins, LuaLocalVariable)
        self.line_info = read_array(ins, read_uint32)
        self.strings = read_array(ins, read_string)
        self.numbers = read_array(ins, read_lua_number)
        self.functions = read_array(ins, LuaCodeChunk)
        self.instructions = read_array(ins, read_uint32)


class LuaFile:
    def __init__(self, data):
        ins = io.BytesIO(data)
        self.header = LuaFileHeader(ins)
        self.root_code_chunk = LuaCodeChunk(ins)


def disassemble_function(c, lines, depth=0):
    """Disassemble LuaCodeChunk c into lines."""
    spacer = '    ' * depth
    lines.append('{}Function {}'.format(spacer, c.source))
    fmt = '{}locals: {} strings: {} numbers: {} functions: {}'
    lines.append(fmt.format(spacer, len(c.local_vars), len(c.strings),
                            len(c.numbers), len(c.functions)))
    for instruction in c.instructions:
        opcode = instruction & 0x3f
        raw_args = instruction >> 6
        opname, optype = OPS[opcode]
        if optype == LuaOpType.NONE:
            args = ''
        elif optype == LuaOpType.U:
            args = '{}'.format(raw_args)
        elif optype == LuaOpType.S:
            args = '{}'.format(raw_args - (0x7fffffff >> 6))
        elif optype == LuaOpType.AB:
            args = '{} {}'.format(raw_args >> 9, raw_args & 0x1ff)
        elif optype == LuaOpType.K:
            args = c.strings[raw_args]
        elif optype == LuaOpType.J:
            args = '@{}'.format(raw_args - (0x7fffffff >> 6))
        elif optype == LuaOpType.L:
            args = '{}'.format(raw_args)
        elif optype == LuaOpType.N:
            args = '{}'.format(c.numbers[raw_args])
        else:
            raise ValueError('bad lua optype')
        lines.append('{}{} {}'.format(spacer, opname, args))

    for f in c.functions:
        disassemble_function(f, lines, depth + 1)


def disassemble(data):
    f = LuaFile(data)
    lines = []
    disassemble_function(f.root_code_chunk, lines)
    return '\n'.join(lines)