Extractor for "actionscripts" used in NES and SNES games developed by HAL Laboratory
#!/usr/bin/python3

# Some of the code you'll see here is bad, but it works.

import argparse
import os
import string
import sys

LABEL_CHARSET = string.ascii_letters + string.digits + '_'  # Valid characters for labels

LAST_EB_SCRIPT = 895  # Amount of scripts in Earthbound
LAST_M2_SCRIPT = 890  # Amount of scripts in MOTHER 2

# In the format (START, END)
EB_SCRIPT_BLOCKS = (
    (0x00AD8A, 0x00AD9F),
    (0x02FFB7, 0x02FFCC),
    (0x030195, 0x039FF2),
    (0x03A043, 0x03DFE8),
    (0x040E24, 0x040E31),
    (0x040F18, 0x041A9E),
    (0x042172, 0x0423DC),
    (0x04279F, 0x04283F)
)

# In the format (START, END)
M2_SCRIPT_BLOCKS = (
    # TODO
)

DATA_TYPE_SIZES = {
    'label_16': 2,  # 16-bit address (prefferably a label, when available)
    'label_24': 3,  # 24-bit address (prefferably a label, when available)
    'imm_8':    1,  # 8-bit immediate (hex representation)
    'imm_s8':   1,  # 8-bit signed immediate (decimal representation)
    'imm_u8':   1,  # 8-bit unsigned immediate (decimal representation)
    'imm_16':   2,  # 16-bit immediate (hex representation)
    'imm_s16':  2,  # 16-bit signed immediate (decimal representation)
    'imm_u16':  2,  # 16-bit unsigned immediate (decimal representation)
    'addr_16':  2,  # 16-bit address
    'addr_24':  3,  # 24-bit address
    'addr_32':  4,  # 32-bit address (high byte ignored)
    'obj_var':  1,  # Object variable (i.e. VAR0, VAR1, ..., VAR7)
    'reg':      0,  # The work register
    'nop':      0,  # Special, used only by opcode 0x0F (ONTICK NOP)
}

# (mnemonic, data types separated by spaces)
OPCODES = (
    ('END',         ''),                       # 00
    ('STARTLOOP',   'imm_u8'),                 # 01
    ('ENDLOOP',     ''),                       # 02
    ('JML',         'label_24'),               # 03
    ('JSL',         'label_24'),               # 04
    ('RTL',         ''),                       # 05
    ('WAIT',        'imm_u8'),                 # 06
    ('TASK',        'label_16'),               # 07
    ('ONTICK',      'addr_24'),                # 08
    ('HALT',        ''),                       # 09
    ('JEQ',         'label_16'),               # 0A
    ('JNE',         'label_16'),               # 0B
    ('ENDTASK',     ''),                       # 0C
    ('BINOP.w',     'addr_16 imm_u8 imm_16'),  # 0D - BINOP with 16-bit memory value
    ('MOV',         'obj_var imm_16'),         # 0E
    ('ONTICK',      'nop'),                    # 0F
    ('MULTIJMP',    'imm_u8'),                 # 10
    ('MULTIJSR',    'imm_u8'),                 # 11
    ('MOV.b',       'addr_16 imm_8'),          # 12
    ('ENDLASTTASK', ''),                       # 13
    ('BINOP',       'obj_var imm_u8 imm_16'),  # 14 - BINOP with object variable
    ('MOV.w',       'addr_16 imm_16'),         # 15
    ('BREAKEQ',     'label_16'),               # 16
    ('BREAKNE',     'label_16'),               # 17
    ('BINOP.b',     'addr_16 imm_u8 imm_8'),   # 18 - BINOP with 8-bit memory value
    ('JMP',         'label_16'),               # 19
    ('JSR',         'label_16'),               # 1A
    ('RTS',         ''),                       # 1B
    ('SETANIMPTR',  'addr_24'),                # 1C
    ('MOV',         'reg imm_16'),             # 1D
    ('MOV',         'reg addr_16'),            # 1E
    ('MOV',         'obj_var reg'),            # 1F
    ('MOV',         'reg obj_var'),            # 20
    ('WAIT',        'obj_var'),                # 21
    ('ONDRAW',      'addr_16'),                # 22
    ('ONPOSITION',  'addr_16'),                # 23
    ('STARTLOOP',   'reg'),                    # 24
    ('ONMOVE',      'addr_16'),                # 25
    ('SETANIM',     'obj_var'),                # 26
    ('BINOP',       'reg imm_u8 imm_16'),      # 27 - BINOP with work register
    ('SETXPOS',     'imm_16'),                 # 28
    ('SETYPOS',     'imm_16'),                 # 29
    ('SETZPOS',     'imm_16'),                 # 2A
    ('ADDXPOS',     'imm_s16'),                # 2B
    ('ADDYPOS',     'imm_s16'),                # 2C
    ('ADDZPOS',     'imm_s16'),                # 2D
    ('ADDXVEL',     'imm_16'),                 # 2E
    ('ADDYVEL',     'imm_16'),                 # 2F
    ('ADDZVEL',     'imm_16'),                 # 30
    ('UNK31',       'imm_8 imm_8'),            # 31 - Mystery instruction, never used
    ('UNK32',       'imm_8 imm_8'),            # 32 - Mystery instruction, never used
    ('UNK33',       'imm_8 imm_8'),            # 33 - Mystery instruction, never used
    ('UNK34',       'imm_8 imm_8'),            # 34 - Mystery instruction, never used
    ('UNK35',       'imm_8 imm_8'),            # 35 - Mystery instruction, never used
    ('UNK36',       'imm_8 imm_8'),            # 36 - Mystery instruction, never used
    ('UNK37',       'imm_8 imm_8'),            # 37 - Mystery instruction, never used
    ('UNK38',       'imm_8 imm_8'),            # 38 - Mystery instruction, never used
    ('ZEROVEL',     ''),                       # 39
    ('UNK3A',       'imm_8'),                  # 3A - Mystery instruction, never used
    ('SETANIM',     'imm_s8'),                 # 3B
    ('INCANIM',     ''),                       # 3C
    ('DECANIM',     ''),                       # 3D
    ('ADDANIM',     'imm_s8'),                 # 3E
    ('SETXVEL',     'imm_16'),                 # 3F
    ('SETYVEL',     'imm_16'),                 # 40
    ('SETZVEL',     'imm_16'),                 # 41
    ('ASMCALL',     'addr_24'),                # 42
    ('PRIORITY',    'imm_u8'),                 # 43
    ('WAIT',        'reg')                     # 44
)

BINOPS = ('AND', 'OR', 'ADD', 'XOR')

class Disassembler(object):
    def __init__(self, rom_file, out_file, sym_file, asm_functions_file, header_offset, script_count, indent=4):
        self._indent = indent
        self.rom_file = rom_file
        self.out_file = out_file
        self.sym_file = sym_file
        self.asm_functions_file = asm_functions_file
        self.header_offset = header_offset
        self.script_count = script_count
        self.indentation = self._indent
        self.pc = 0
        self.was_linebreak = False
        self.force_label = False
        self.symbols = dict()
        self.asm_functions = dict()
        self.init_asm_functions()
        self.init_symbols()

    @property
    def snes_pc(self):
        return self.pc | 0xC00000

    def read_rom(self, amount, signed=False):
        b = self.rom_file.read(amount)
        return int.from_bytes(b, signed=signed, byteorder='little')

    def datatype_to_str(self, data_type, bytes_):
        value = int.from_bytes(bytes_, byteorder='little')
        if data_type == 'label_16':
            return self.symbols.get((self.snes_pc & 0xFF0000) | value, '${:04X}'.format(value))
        elif data_type == 'label_24':
            return self.symbols.get(value, '${:06X}'.format(value))
        elif data_type == 'imm_8':
            return '#${:02X}'.format(value)
        elif data_type == 'imm_16':
            return '#${:04X}'.format(value)
        elif data_type in ('imm_s8', 'imm_s16'):
            value = int.from_bytes(bytes_, signed=True, byteorder='little')
            return '#{}'.format(value)
        elif data_type in ('imm_u8', 'imm_u16'):
            return '#{}'.format(value)
        elif data_type == 'addr_16':
            return '${:04X}'.format(value)
        elif data_type == 'addr_24':
            return '${:06X}'.format(value)
        elif data_type == 'addr_32':
            value = ((value & 0xFFFF0000) >> 16) | ((value & 0xFFFF) << 16)  # addr_32 has a weird-ass format
            return '${:06X}'.format(value)
        elif data_type == 'obj_var':
            if value > 7:
                print('WARNING: INVALID OBJ_VAR OPERAND:', value)

            return 'VAR{}'.format(value)
        elif data_type in ('reg', 'nop'):
            return data_type.upper()
        else:
            raise TypeError('Unknown data type! ({})'.format(data_type))

    def try_add_label(self):
        label = self.symbols.get(self.snes_pc)

        if self.force_label:
            if not label:
                label = 'L_{:06X}'.format(self.snes_pc)
                self.symbols[self.snes_pc] = label

            label = '\n' + label

        if label:
            self.out_file.write(label + ':\n')

        return bool(label)

    def disassemble_all(self):
        script_blocks = EB_SCRIPT_BLOCKS if self.script_count == LAST_EB_SCRIPT else M2_SCRIPT_BLOCKS

        size = 0
        for start, end in script_blocks:
            size += end - start
            self.rom_file.seek(start)
            self.pc = self.rom_file.tell()

            while self.pc < end:
                was_label = self.try_add_label()
                self.was_linebreak = was_label or self.was_linebreak
                self.force_label = False

                opcode = self.read_rom(1)
                self.disasm_opcode(opcode)
                self.force_label = opcode in (0x00, 0x03, 0x05, 0x09, 0x0C, 0x19, 0x1B)
                self.pc = self.rom_file.tell()
        print(f'SIZE IS {size}')

    def disasm_opcode(self, op_byte):
        indentation = ' ' * self.indentation

        waited = None
        opcode = op_byte
        if opcode >= 0x70:
            waited = opcode & 0x0F
            opcode = ((opcode & 0x70) >> 4) + 0x3B  # Get the actual opcode from the "waited opcode" code

        if opcode >= len(OPCODES):
            to_write = (indentation + '.byte'.ljust(12) + '${:02X}'.format(opcode)).ljust(40 + len(indentation))
            to_write += '; {:06X}/{:02X}\n'.format(self.snes_pc, opcode)
            self.out_file.write(to_write)
            return

        op_addr = self.snes_pc

        mnemonic, types = OPCODES[opcode]
        types = types.split()

        bytes_ = bytearray([op_byte])
        operands = []
        for data_type in types:
            b = self.rom_file.read(DATA_TYPE_SIZES[data_type])
            bytes_ += b
            operands.append(self.datatype_to_str(data_type, b))

        if waited is not None:
            operands.append('WAIT #{}'.format(waited))

        if not self.was_linebreak:
            mnemonic = '\n' + mnemonic  # Add leading newline if previous line didn't have a line break
            self.was_linebreak = True

        comment = ''
        extra = ''

        if opcode == 0x42:  # ASMCALL
            address = int.from_bytes(bytes_[1:], byteorder='little')
            asm_func = self.asm_functions.get(address, None)
            if asm_func:
                DIRECTIVES = ('.byte', '.word', '.long', '.dword')
                comment = ' // ' + asm_func['comment'] if asm_func['comment'] else ''
                
                # Fucking hell this is horrible
                for p in asm_func['params']:
                    if p == 'varargs':
                        count = self.read_rom(1)
                        
                        self.pc = self.rom_file.tell()
                        extra += (indentation + '.byte'.ljust(12) + str(count)).ljust(40 + len(indentation))
                        extra += '; {:06X}/{:02X}\n'.format(self.snes_pc, count)
                        
                        for i in range(count):
                            self.pc = self.rom_file.tell()
                            b = self.rom_file.read(2)
                            addr = int.from_bytes(b, byteorder='little')
                            extra += (indentation + '.word'.ljust(12) + str(addr)).ljust(40 + len(indentation))
                            extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
                    else:
                        self.pc = self.rom_file.tell()
                        data_size = DATA_TYPE_SIZES[p]
                        b = self.rom_file.read(data_size)
                        arg = self.datatype_to_str(p, b)
                        arg = arg.replace('#', '')
                        extra += (indentation + DIRECTIVES[data_size - 1].ljust(12) + arg).ljust(40 + len(indentation))
                        extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
                        
            else:
                print("WARNING: I DON'T KNOW ANYTHING ABOUT ASM FUNCTION {:06X}".format(address))
        elif opcode in (0x0D, 0x14, 0x18, 0x27):  # BINOP
            # Now THIS is what I call hacky shit!
            if opcode == 0x27:
                op = bytes_[1]
            elif opcode == 0x14:
                op = bytes_[2]
            else:
                op = bytes_[3]

            if op in (0, 1, 2, 3):
                del operands[1]
                mnemonic = mnemonic.replace('BINOP', BINOPS[op])
                
                if op == 2:  # ADD, change operand type from HEXADECIMAL immediate to DECIMAL immediate
                    data_type = 'imm_s8' if opcode == 0x18 else 'imm_s16'                    
                    size = DATA_TYPE_SIZES[data_type]
                    self.rom_file.seek(-size, os.SEEK_CUR)
                    operands[-1] = self.datatype_to_str(data_type, self.rom_file.read(size))
        elif opcode in (0x10, 0x11):  # MULTIJMP, MULTIJSR
            count = bytes_[1]

            for i in range(count):
                self.pc = self.rom_file.tell()
                b = self.rom_file.read(2)
                addr = self.datatype_to_str('label_16', b)
                extra += (indentation + '.word'.ljust(12) + addr).ljust(40 + len(indentation))
                extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
        elif opcode in (0x00, 0x05, 0x09, 0x0C, 0x1B):  # END, RTL, HALT, ENDTASK, RTS
            self.indentation = self._indent
            indentation = ' ' * self.indentation
            self.was_linebreak = True
        elif opcode in (0x01, 0x24):  # STARTLOOP
            self.indentation += self._indent
        elif opcode == 0x02:  # ENDLOOP
            self.indentation = max(self._indent, self.indentation - self._indent)
            indentation = ' ' * self.indentation
        elif opcode in (0x03, 0x19):  # JML, JMP
            self.was_linebreak = True
            self.indentation = max(self._indent, self.indentation - self._indent)

        to_write = (indentation + mnemonic.ljust(12) + ','.join(operands)).ljust(40 + len(indentation))
        to_write += '; {:06X}/{}{}\n'.format(op_addr, bytes_.hex().upper(), comment)
        self.out_file.write(to_write + extra)

    def init_symbols(self):
        if self.sym_file:
            self.parse_sym_file()

        # Seek to the start of the "script entry points" pointer table
        if self.script_count == LAST_EB_SCRIPT:
            self.rom_file.seek(0x0400D4 + self.header_offset)  # Earthbound
        else:
            self.rom_file.seek(0x04002F + self.header_offset)  # MOTHER 2

        # Add default symbols for each script entry point, if they've not been defined
        for i in range(script_count):
            address = self.read_rom(3)
            if address not in self.symbols:
                self.symbols[address] = 'Script{:04d}'.format(i)

        script_blocks = EB_SCRIPT_BLOCKS if self.script_count == LAST_EB_SCRIPT else M2_SCRIPT_BLOCKS

        # BAD STUFF AHEAD, LOTS OF DUPLICATE CODE FROM "disasm_opcode"
        for start, end in script_blocks:
            self.rom_file.seek(start)
            self.pc = self.rom_file.tell()

            while self.pc < end:
                opcode = self.read_rom(1)

                if opcode >= 0x70:
                    opcode = ((opcode & 0x70) >> 4) + 0x3B  # Get the actual opcode from the "waited opcode" code

                if opcode >= len(OPCODES):
                    continue

                if opcode == 0x07:  # TASK
                    bank_mask = self.snes_pc & 0xFF0000
                    addr = bank_mask | self.read_rom(2)
                    if addr not in self.symbols:
                        self.symbols[addr] = f'TASK_{addr:06X}'
                elif opcode == 0x42:  # ASMCALL
                    addr = self.read_rom(3)
                    asm_func = self.asm_functions.get(addr, None)
                    if asm_func:
                        for p in asm_func['params']:
                            if p == 'varargs':
                                count = self.read_rom(1)
                                self.rom_file.read(count*2)  # Read and discard bytes
                            else:
                                data_size = DATA_TYPE_SIZES[p]
                                b = self.rom_file.read(data_size)  # Read and discard bytes
                elif opcode in (0x10, 0x11):  # MULTIJMP/MULTIJSR
                    count = self.read_rom(1)
                    self.rom_file.read(count*2)  # Read and discard bytes
                else:
                    types = OPCODES[opcode][1].split()
                    for data_type in types:
                        if data_type == 'label_16':
                            addr = (self.snes_pc & 0xFF0000) | self.read_rom(2)
                            if addr not in self.symbols:
                                self.symbols[addr] = f'L_{addr:06X}'
                        elif data_type == 'label_24':
                            addr = self.read_rom(3)
                            if addr not in self.symbols:
                                self.symbols[addr] = f'L_{addr:06X}'
                        else:
                            self.rom_file.read(DATA_TYPE_SIZES[data_type])  # Read and discard bytes

                self.pc = self.rom_file.tell()

    def parse_sym_file(self):
        for i, line in enumerate(self.sym_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                label, str_address = [s.strip() for s in pre_comment.split('=', maxsplit=1)]

                if not all(c in LABEL_CHARSET for c in label):
                    print('Ignoring line {} from {}: Invalid label name ({})'.format(i, self.sym_file.name, label), file=sys.stderr)
                    continue

                try:
                    address = int(str_address, 16)
                    self.symbols[address] = label
                except ValueError:
                    print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.sym_file.name, str_address), file=sys.stderr)
            elif pre_comment:
                print('Ignoring line {} from {}: Invalid line'.format(i, self.sym_file.name), file=sys.stderr)

    def init_asm_functions(self):
        if not self.asm_functions_file:
            return

        address = None
        asm_function = None
        for i, line in enumerate(self.asm_functions_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                if address is None:
                    print('Ignoring line {} from {}: Attribute assignment without ASM_FUNCTION directive'.format(i, self.asm_functions_file.name), file=sys.stderr)
                    continue

                key, value = [s.strip() for s in pre_comment.split('=', maxsplit=1)]                
                if key == 'COMMENT':
                    self.asm_functions[address]['comment'] = value
                elif key == 'PARAMS':
                    self.asm_functions[address]['params'] = [s.strip() for s in value.split(',')]
                else:
                    print('Ignoring line {} from {}: Invalid attribute ({})'.format(i, self.asm_functions_file.name, key), file=sys.stderr)
            elif pre_comment:
                directive, arg = [s.strip() for s in pre_comment.split(maxsplit=1)]
                if directive == 'ASM_FUNCTION':
                    try:
                        address = int(arg, 16)
                        self.asm_functions[address] = {'comment': '', 'params': []}
                    except ValueError:
                        print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.asm_functions_file.name, arg), file=sys.stderr)
                else:
                    print('Ignoring line {} from {}: Unknown directive ({})'.format(i, self.asm_functions_file.name, directive), file=sys.stderr)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--symfile', help='the symbols definition file')
    parser.add_argument('-a', '--asmfuncfile', help='the asm functions definition file')
    parser.add_argument('romfile', help='the Earthbound or MOTHER 2 ROM file')
    parser.add_argument('outfile', help='the output file')
    args = parser.parse_args()

    rom_size = os.path.getsize(args.romfile)
    header_offset = rom_size % 0x010000
    if rom_size < 0x300000 or (header_offset != 0 and header_offset != 512):
        print("The file {} doesn't look like a valid Earthbound/MOTHER 2 ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    rom_file = open(args.romfile, 'rb')
    out_file = open(args.outfile, 'w')
    sym_file = open(args.symfile, 'r') if args.symfile else None
    asm_funcs_file = open(args.asmfuncfile, 'r') if args.asmfuncfile else None

    rom_file.seek(0xFFC0 + header_offset)
    rom_name = rom_file.read(21)
    if rom_name == b'EARTH BOUND'.ljust(21):
        script_count = LAST_EB_SCRIPT
    elif rom_name == b'MOTHER 2'.ljust(21):
        script_count = LAST_M2_SCRIPT
    else:
        print("The file {} doesn't look like a valid Earthbound/MOTHER 2 ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    disassembler = Disassembler(rom_file, out_file, sym_file, asm_funcs_file, header_offset, script_count, indent=4)
    disassembler.disassemble_all()

    rom_file.close()
    out_file.close()
    if sym_file:
        sym_file.close()

    with open('symbols_out.txt', 'w') as f:
        for k, v in disassembler.symbols.items():
            f.write('{:06X}: {}\n'.format(k, v))

Earthbound/actionscript_dumper.py
CataLatas/hal-actionscript-extractor
#!/usr/bin/python3

# Some of the code you'll see here is bad, but it works.

import argparse
import os
import string
import sys

LABEL_CHARSET = string.ascii_letters + string.digits + '_'  # Valid characters for labels

LAST_SCRIPT = 134  # Amount of scripts

SCRIPT_BLOCKS = (
    (0x00772B, 0x0077BB),
    (0x00C609, 0x00C737),
    (0x00C74A, 0x00C770),
    (0x00C783, 0x00C825)
)

DATA_TYPE_SIZES = {
    'label_16': 2,  # 16-bit address (prefferably a label, when available)
    'label_24': 3,  # 24-bit address (prefferably a label, when available)
    'imm_8':    1,  # 8-bit immediate (hex representation)
    'imm_s8':   1,  # 8-bit signed immediate (decimal representation)
    'imm_u8':   1,  # 8-bit unsigned immediate (decimal representation)
    'imm_16':   2,  # 16-bit immediate (hex representation)
    'imm_s16':  2,  # 16-bit signed immediate (decimal representation)
    'imm_u16':  2,  # 16-bit unsigned immediate (decimal representation)
    'addr_16':  2,  # 16-bit address
    'addr_24':  3,  # 24-bit address
    'addr_32':  4,  # 32-bit address (high byte ignored)
    'obj_var':  1,  # Object variable (i.e. VAR0, VAR1, ..., VAR7)
    'reg':      0,  # The work register
    'nop':      0,  # Special, used only by opcode 0x0F (ONTICK NOP)
}

# (mnemonic, data types separated by spaces)
OPCODES = (
    ('END',         ''),
    ('STARTLOOP',   'imm_u8'),
    ('ENDLOOP',     ''),
    ('JML',         'label_24'),
    ('JSL',         'label_24'),
    ('RTL',         ''),
    ('WAIT',        'imm_u8'),
    ('ASMCALL',     'addr_24'),
    ('TASK',        'label_16'),
    ('ONTICK',      'addr_24'),
    ('HALT',        ''),
    ('JEQ',         'label_16'),
    ('JNE',         'label_16'),
    ('ENDTASK',     ''),
    ('BINOP.w',     'addr_16 imm_u8 imm_16'),  # BINOP with 16-bit memory value
    ('MOV',         'obj_var imm_16'),
    ('ONTICK',      'nop'),
    ('MULTIJMP',    'imm_u8'),
    ('MULTIJSR',    'imm_u8'),
    ('MOV.b',       'addr_16 imm_8'),
    ('UNK_TASK',    'imm_s8'),
    ('BINOP',       'obj_var imm_u8 imm_16'),  # BINOP with object variable
    ('MOV.w',       'addr_16 imm_16'),
    ('BREAKEQ',     'label_16'),
    ('BREAKNE',     'label_16'),
    ('BINOP.b',     'addr_16 imm_u8 imm_8'),   # BINOP with 8-bit memory value
    ('JMP',         'label_16'),
    ('JSR',         'label_16'),
    ('RTS',         ''),
    ('SETANIMPTR',  'addr_24'),
    ('MOV',         'reg imm_16'),
    ('MOV',         'reg addr_16'),
    ('WEIRD_1'      'imm_8'),        # Nonexistant instruction in Earthbound
    ('WEIRD_2',     'imm_8'),        # Nonexistant instruction in Earthbound
    ('WEIRD_3',     'imm_8'),        # Nonexistant instruction in Earthbound
    ('MOV',         'obj_var reg'),
    ('MOV',         'reg obj_var'),
    ('WAIT',        'obj_var'),
    # SPECIAL INSTRUCTIONS BEGIN
    ('SETANIM',     'imm_s8'),
    ('SETXPOS',     'imm_16'),
    ('SETYPOS',     'imm_16'),
    ('ADDXPOS',     'imm_s16'),
    ('ADDYPOS',     'imm_s16'),
    ('SETXVEL',     'imm_16'),
    ('SETYVEL',     'imm_16'),
    ('ADDXVEL',     'imm_16'),
    ('ADDYVEL',     'imm_16'),
    ('BGHDISP',     'imm_u8 imm_16'),  # Set background horizontal displacement          (Earthbound: UNK31)
    ('BGVDISP',     'imm_u8 imm_16'),  # Set background vertical displacement            (Earthbound: UNK32)
    ('BGHDISPVEL',  'imm_u8 imm_16'),  # Set background horizontal displacement velocity (Earthbound: UNK33)
    ('BGVDISPVEL',  'imm_u8 imm_16'),  # Set background vertical displacement velocity   (Earthbound: UNK34)
    ('UNK_35',      'imm_u8 imm_16'),  # UNK35 in Earthbound
    ('UNK_36',      'imm_u8 imm_16'),  # UNK36 in Earthbound
    ('INCANIM',     ''),
    ('DECANIM',     ''),
    ('ADDANIM',     'imm_s8'),
    ('UNK_37',      'imm_8 imm_16'),   # UNK37 in Earthbound
    ('UNK_38',      'imm_8 imm_16'),   # UNK38 in Earthbound
    ('ZEROVEL',     ''),
    ('ZRBGDISPVEL', 'imm_8'),          # Zero background displacement velocity (Earthbound: UNK3A)
    ('SETZPOS',     'imm_16'),
    ('ADDZPOS',     'imm_s16'),
    ('SETZVEL',     'imm_16'),
    ('ADDZVEL',     'imm_16'),
)

BINOPS = ('AND', 'OR', 'ADD')

class Disassembler(object):
    def __init__(self, rom_file, out_file, sym_file, asm_functions_file, header_offset, script_count, indent=4):
        self._indent = indent
        self.rom_file = rom_file
        self.out_file = out_file
        self.sym_file = sym_file
        self.asm_functions_file = asm_functions_file
        self.header_offset = header_offset
        self.script_count = script_count
        self.indentation = self._indent
        self.pc = 0
        self.was_linebreak = False
        self.force_label = False
        self.symbols = dict()
        self.asm_functions = dict()
        self.init_asm_functions()
        self.init_symbols()

    @property
    def snes_pc(self):
        addr = self.pc
        bank = addr >> 15
        addr = (addr & 0x7FFF) + ((bank & 0x2F) * 0x010000) | 0x8000

        return addr

    def read_rom(self, amount, signed=False):
        b = self.rom_file.read(amount)
        return int.from_bytes(b, signed=signed, byteorder='little')

    def datatype_to_str(self, data_type, bytes_):
        value = int.from_bytes(bytes_, byteorder='little')
        if data_type == 'label_16':
            bank_mask = self.snes_pc & 0xFF0000
            return self.symbols.get(bank_mask | value, '${:04X}'.format(value))
        elif data_type == 'label_24':
            return self.symbols.get(value, '${:06X}'.format(value))
        elif data_type == 'imm_8':
            return '#${:02X}'.format(value)
        elif data_type == 'imm_16':
            return '#${:04X}'.format(value)
        elif data_type in ('imm_s8', 'imm_s16'):
            value = int.from_bytes(bytes_, signed=True, byteorder='little')
            return '#{}'.format(value)
        elif data_type in ('imm_u8', 'imm_u16'):
            return '#{}'.format(value)
        elif data_type == 'addr_16':
            return '${:04X}'.format(value)
        elif data_type == 'addr_24':
            return '${:06X}'.format(value)
        elif data_type == 'addr_32':
            value = ((value & 0xFFFF0000) >> 16) | ((value & 0xFFFF) << 16)  # addr_32 has a weird-ass format
            return '${:06X}'.format(value)
        elif data_type == 'obj_var':
            if value > 3:
                print('WARNING: INVALID OBJ_VAR OPERAND:', value)

            return 'VAR{}'.format(value)
        elif data_type in ('reg', 'nop'):
            return data_type.upper()
        else:
            raise TypeError('Unknown data type! ({})'.format(data_type))

    def try_add_label(self):
        label = self.symbols.get(self.snes_pc)

        if self.force_label:
            if not label:
                label = 'L_{:06X}'.format(self.snes_pc)
                self.symbols[self.snes_pc] = label

            label = '\n' + label

        if label:
            self.out_file.write(label + ':\n')

        return bool(label)

    def disassemble_all(self):
        self.force_label = False  # For first
        for start, end in SCRIPT_BLOCKS:
            self.rom_file.seek(start)
            self.pc = self.rom_file.tell()

            while self.pc < end:
                was_label = self.try_add_label()
                self.was_linebreak = was_label or self.was_linebreak
                self.force_label = False

                opcode = self.read_rom(1)
                self.disasm_opcode(opcode)
                self.force_label = opcode in (0x00, 0x03, 0x05, 0x0A, 0x0E, 0x1A, 0x1C)
                self.pc = self.rom_file.tell()

    def disasm_opcode(self, op_byte):
        indentation = ' ' * self.indentation

        waited = 0
        opcode = op_byte
        if opcode >= 0x30:
            waited = opcode & 0x07
            opcode = (((opcode - 0x30) & 0xF8) >> 3) + 0x26

        if opcode >= len(OPCODES):
            to_write = (indentation + '.byte'.ljust(12) + '${:02X}'.format(opcode)).ljust(40 + len(indentation))
            to_write += '; {:06X}/{:02X}\n'.format(self.snes_pc, opcode)
            self.out_file.write(to_write)
            return

        op_addr = self.snes_pc

        mnemonic, types = OPCODES[opcode]
        types = types.split()

        bytes_ = bytearray([op_byte])
        operands = []
        for data_type in types:
            b = self.rom_file.read(DATA_TYPE_SIZES[data_type])
            bytes_ += b
            operands.append(self.datatype_to_str(data_type, b))

        if waited > 0:
            operands.append('WAIT #{}'.format(waited))

        if not self.was_linebreak:
            mnemonic = '\n' + mnemonic  # Add leading newline if previous line didn't have a line break
            self.was_linebreak = True

        comment = ''
        extra = ''

        if opcode == 0x07:  # ASMCALL
            address = int.from_bytes(bytes_[1:], byteorder='little')
            asm_func = self.asm_functions.get(address, None)
            if asm_func:
                DIRECTIVES = ('.byte', '.word', '.long', '.dword')
                comment = ' // ' + asm_func['comment'] if asm_func['comment'] else ''
                
                # Fucking hell this is horrible
                for p in asm_func['params']:
                    if p == 'varargs':
                        count = self.read_rom(1)
                        
                        self.pc = self.rom_file.tell()
                        extra += (indentation + '.byte'.ljust(12) + str(count)).ljust(40 + len(indentation))
                        extra += '; {:06X}/{:02X}\n'.format(self.snes_pc, count)
                        
                        for i in range(count):
                            self.pc = self.rom_file.tell()
                            b = self.rom_file.read(2)
                            addr = int.from_bytes(b, byteorder='little')
                            extra += (indentation + '.word'.ljust(12) + str(addr)).ljust(40 + len(indentation))
                            extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
                    else:
                        self.pc = self.rom_file.tell()
                        data_size = DATA_TYPE_SIZES[p]
                        b = self.rom_file.read(data_size)
                        arg = self.datatype_to_str(p, b)
                        arg = arg.replace('#', '')
                        extra += (indentation + DIRECTIVES[data_size - 1].ljust(12) + arg).ljust(40 + len(indentation))
                        extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
                        
            else:
                print("WARNING: I DON'T KNOW ANYTHING ABOUT ASM FUNCTION {:06X}".format(address))
        elif opcode in (0x0E, 0x15, 0x19):  # BINOP
            # Now THIS is what I call hacky shit!
            if opcode == 0x15:
                op = bytes_[2]
            else:
                op = bytes_[3]

            if op in (0, 1, 2):
                del operands[1]
                mnemonic = mnemonic.replace('BINOP', BINOPS[op])
                
                if op == 2:  # ADD, change operand type from HEXADECIMAL immediate to DECIMAL immediate
                    data_type = 'imm_s8' if opcode == 0x19 else 'imm_s16'
                    size = DATA_TYPE_SIZES[data_type]
                    self.rom_file.seek(-size, os.SEEK_CUR)
                    operands[-1] = self.datatype_to_str(data_type, self.rom_file.read(size))
        elif opcode in (0x11, 0x12):  # MULTIJMP, MULTIJSR
            count = bytes_[1]

            for i in range(count):
                self.pc = self.rom_file.tell()
                b = self.rom_file.read(2)
                addr = self.datatype_to_str('label_16', b)
                extra += (indentation + '.word'.ljust(12) + addr).ljust(40 + len(indentation))
                extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
        elif opcode in (0x00, 0x05, 0x0A, 0x0D, 0x1C):  # END, RTL, HALT, ENDTASK, RTS
            self.indentation = self._indent
            indentation = ' ' * self.indentation
            self.was_linebreak = True
        elif opcode == 0x01:  # STARTLOOP
            self.indentation += self._indent
        elif opcode == 0x02:  # ENDLOOP
            self.indentation = max(self._indent, self.indentation - self._indent)
            indentation = ' ' * self.indentation
        elif opcode in (0x03, 0x1A):  # JML, JMP
            self.was_linebreak = True
            self.indentation = max(self._indent, self.indentation - self._indent)

        to_write = (indentation + mnemonic.ljust(12) + ','.join(operands)).ljust(40 + len(indentation))
        to_write += '; {:06X}/{}{}\n'.format(op_addr, bytes_.hex().upper(), comment)
        self.out_file.write(to_write + extra)

    def init_symbols(self):
        if self.sym_file:
            self.parse_sym_file()

        # Seek to the start of the "script entry points" pointer table
        self.rom_file.seek(0x002785 + self.header_offset)

        # Add default symbols for each script entry point, if they've not been defined
        for i in range(self.script_count):
            address = self.read_rom(3)
            if address not in self.symbols:
                self.symbols[address] = 'Script{:04d}'.format(i)

        # BAD STUFF AHEAD, LOTS OF DUPLICATE CODE FROM "disasm_opcode"
        for start, end in SCRIPT_BLOCKS:
            self.rom_file.seek(start)
            self.pc = self.rom_file.tell()

            while self.pc < end:
                opcode = self.read_rom(1)

                if opcode >= 0x30:
                    opcode = (((opcode - 0x30) & 0xF8) >> 3) + 0x26  # Get the actual opcode from the "waited opcode" code

                if opcode >= len(OPCODES):
                    continue

                if opcode == 0x08:  # TASK
                    bank_mask = self.snes_pc & 0xFF0000
                    addr = bank_mask | self.read_rom(2)
                    if addr not in self.symbols:
                        self.symbols[addr] = f'TASK_{addr:06X}'
                elif opcode == 0x07:  # ASMCALL
                    addr = self.read_rom(3)
                    asm_func = self.asm_functions.get(addr, None)
                    if asm_func:
                        for p in asm_func['params']:
                            if p == 'varargs':
                                count = self.read_rom(1)
                                self.rom_file.read(count*2)  # Read and discard bytes
                            else:
                                data_size = DATA_TYPE_SIZES[p]
                                b = self.rom_file.read(data_size)  # Read and discard bytes
                elif opcode in (0x11, 0x12):  # MULTIJMP/MULTIJSR
                    count = self.read_rom(1)
                    self.rom_file.read(count*2)  # Read and discard bytes
                else:
                    types = OPCODES[opcode][1].split()
                    for data_type in types:
                        if data_type == 'label_16':
                            bank_mask = self.snes_pc & 0xFF0000
                            addr = bank_mask | self.read_rom(2)
                            if addr not in self.symbols:
                                self.symbols[addr] = f'L_{addr:06X}'
                        elif data_type == 'label_24':
                            addr = self.read_rom(3)
                            if addr not in self.symbols:
                                self.symbols[addr] = f'L_{addr:06X}'
                        else:
                            self.rom_file.read(DATA_TYPE_SIZES[data_type])  # Read and discard bytes

                self.pc = self.rom_file.tell()

    def parse_sym_file(self):
        for i, line in enumerate(self.sym_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                label, str_address = [s.strip() for s in pre_comment.split('=', maxsplit=1)]

                if not all(c in LABEL_CHARSET for c in label):
                    print('Ignoring line {} from {}: Invalid label name ({})'.format(i, self.sym_file.name, label), file=sys.stderr)
                    continue

                try:
                    address = int(str_address, 16)
                    self.symbols[address] = label
                except ValueError:
                    print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.sym_file.name, str_address), file=sys.stderr)
            elif pre_comment:
                print('Ignoring line {} from {}: Invalid line'.format(i, self.sym_file.name), file=sys.stderr)

    def init_asm_functions(self):
        if not self.asm_functions_file:
            return

        address = None
        asm_function = None
        for i, line in enumerate(self.asm_functions_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                if address is None:
                    print('Ignoring line {} from {}: Attribute assignment without ASM_FUNCTION directive'.format(i, self.asm_functions_file.name), file=sys.stderr)
                    continue

                key, value = [s.strip() for s in pre_comment.split('=', maxsplit=1)]                
                if key == 'COMMENT':
                    self.asm_functions[address]['comment'] = value
                elif key == 'PARAMS':
                    self.asm_functions[address]['params'] = [s.strip() for s in value.split(',')]
                else:
                    print('Ignoring line {} from {}: Invalid attribute ({})'.format(i, self.asm_functions_file.name, key), file=sys.stderr)
            elif pre_comment:
                directive, arg = [s.strip() for s in pre_comment.split(maxsplit=1)]
                if directive == 'ASM_FUNCTION':
                    try:
                        address = int(arg, 16)
                        self.asm_functions[address] = {'comment': '', 'params': []}
                    except ValueError:
                        print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.asm_functions_file.name, arg), file=sys.stderr)
                else:
                    print('Ignoring line {} from {}: Unknown directive ({})'.format(i, self.asm_functions_file.name, directive), file=sys.stderr)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--symfile', help='the symbols definition file')
    parser.add_argument('-a', '--asmfuncfile', help='the asm functions definition file')
    parser.add_argument('romfile', help="the HyperZone ROM file")
    parser.add_argument('outfile', help='the output file')
    args = parser.parse_args()

    rom_size = os.path.getsize(args.romfile)
    header_offset = rom_size % 0x010000
    if rom_size < 0x080000 or (header_offset != 0 and header_offset != 512):
        print("The file {} doesn't look like a valid HyperZone ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    rom_file = open(args.romfile, 'rb')
    out_file = open(args.outfile, 'w')
    sym_file = open(args.symfile, 'r') if args.symfile else None
    asm_funcs_file = open(args.asmfuncfile, 'r') if args.asmfuncfile else None

    rom_file.seek(0x7FC0 + header_offset)
    rom_name = rom_file.read(21)

    if rom_name != b'HYPER ZONE'.ljust(21):
        print("The file {} doesn't look like a valid HyperZone ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    disassembler = Disassembler(rom_file, out_file, sym_file, asm_funcs_file, header_offset, LAST_SCRIPT, indent=4)
    disassembler.disassemble_all()

    rom_file.close()
    out_file.close()
    if sym_file:
        sym_file.close()

    with open('symbols_out.txt', 'w') as f:
        for k, v in disassembler.symbols.items():
            f.write('{:06X}: {}\n'.format(k, v))

HyperZone/actionscript_dumper.py
CataLatas/hal-actionscript-extractor
#!/usr/bin/python3

# Some of the code you'll see here is bad, but it works.

import argparse
import os
import string
import sys

LABEL_CHARSET = string.ascii_letters + string.digits + '_'  # Valid characters for labels

LAST_SCRIPT = 229  # Amount of scripts

DATA_TYPE_SIZES = {
    'label_16': 2,  # 16-bit address (prefferably a label, when available)
    'label_24': 3,  # 24-bit address (prefferably a label, when available)
    'imm_8':    1,  # 8-bit immediate (hex representation)
    'imm_s8':   1,  # 8-bit signed immediate (decimal representation)
    'imm_u8':   1,  # 8-bit unsigned immediate (decimal representation)
    'imm_16':   2,  # 16-bit immediate (hex representation)
    'imm_s16':  2,  # 16-bit signed immediate (decimal representation)
    'imm_u16':  2,  # 16-bit unsigned immediate (decimal representation)
    'addr_16':  2,  # 16-bit address
    'addr_24':  3,  # 24-bit address
    'addr_32':  4,  # 32-bit address (high byte ignored)
    'obj_var':  1,  # Object variable (i.e. VAR0, VAR1, ..., VAR7)
    'reg':      0,  # The work register
    'nop':      0,  # Special, used only by opcode 0x0F (ONTICK NOP)
}

# (mnemonic, data types separated by spaces)
OPCODES = (
    ('END',         ''),                       # 00
    ('HALT',        ''),                       # 01
    ('WAIT',        'imm_u8'),                 # 02
    ('ONTICK',      'addr_24'),                # 03
    ('ENDTICK',     ''),                       # 04
    ('LOOP',        'imm_u8'),                 # 05
    ('LOOP',        'reg'),                    # 06
    ('ENDLOOP',     ''),                       # 07
    ('JSL',         'label_24'),               # 08
    ('RTL',         ''),                       # 09
    ('JSR',         'label_16'),               # 0A
    ('MULTIJSR',    'imm_u8'),                 # 0B
    ('RTS',         ''),                       # 0C
    ('JML',         'label_24'),               # 0D
    ('JMP',         'label_16'),               # 0E
    ('MULTIJMP',    'imm_u8'),                 # 0F
    ('JEQ',         'label_16'),               # 10
    ('JNE',         'label_16'),               # 11
    ('BREAKEQ',     'label_16'),               # 12
    ('BREAKNE',     'label_16'),               # 13
    ('SPRITEMAP',   'addr_24'),                # 14
    ('ONDRAW',      'addr_16'),                # 15
    ('UNK16',       'addr_16'),                # 16 - TODO: ONMOVE?
    ('MOV',         'reg imm_16'),             # 17
    ('MOV',         'reg addr_16'),            # 18
    ('MOV',         'obj_var imm_16'),         # 19
    ('MOV.b',       'addr_16 imm_8'),          # 1A
    ('MOV.w',       'addr_16 imm_16'),         # 1B
    ('BINOP',       'obj_var imm_u8 imm_16'),  # 1C
    ('BINOP.b',     'addr_16 imm_u8 imm_8'),   # 1D
    ('BINOP.w',     'addr_16 imm_u8 imm_16'),  # 1E
    ('BINOP',       'reg imm_u8 imm_16'),      # 1F
    ('MOV',         'obj_var reg'),            # 20
    ('MOV',         'reg obj_var'),            # 21
    ('WAIT',        'obj_var'),                # 22
    ('SETPOSE',     'obj_var'),                # 23
    ('UNK24',       ''),                       # 24 - TODO: ZERO $6B6E, $6BE8, $6C62, $6CDC
    ('UNK25',       ''),                       # 25 - TODO: ZERO $6C62, $6CDC
    ('SETXPOS',     'imm_16'),                 # 26
    ('SETYPOS',     'imm_16'),                 # 27
    ('ADDXPOS',     'imm_16'),                 # 28
    ('ADDYPOS',     'imm_16'),                 # 29
    ('UNK2A',       'imm_8 imm_16'),           # 2A
    ('UNK2B',       'imm_8 imm_16'),           # 2B
    ('UNK2C',       'imm_8 imm_16'),           # 2C
    ('UNK2D',       'imm_8 imm_16'),           # 2D
    ('UNK2E',       'imm_8'),                  # 2E
    ('UNK2F',       'imm_8 imm_16'),           # 2F
    ('UNK30',       'imm_8 imm_16'),           # 30
    ('UNK31',       'imm_8 imm_16'),           # 31
    ('UNK32',       'imm_8 imm_16'),           # 32
    ('UNK33',       'imm_16'),                 # 33 - TODO: SET $64B6
    ('UNK34',       'imm_8'),                  # 34 - TODO: CALL $00921B
    ('SETPOSE.w',   'imm_16'),                 # 35
)

OPCODES_WAITED = (
    ('SETPOSE.b',   'imm_s8'),                 # 4x
    ('ADDPOSE',     'imm_s8'),                 # 5x
    ('INCPOSE',     ''),                       # 6x
    ('DECPOSE',     ''),                       # 7x
    ('UNK8x',       'imm_16'),                 # 8x - TODO: SET $6B6E
    ('UNK9x',       'imm_16'),                 # 9x - TODO: SET $6BE8
    ('UNKAx',       'imm_16'),                 # Ax - TODO: SET $6C62
    ('UNKBx',       'imm_16'),                 # Bx - TODO: SET $6CDC
    ('ASMCALL',     'addr_24'),                # Cx
    ('PRESETCALL',  'imm_8'),                  # Dx
    ('UNKEx',       'imm_8'),                  # Ex - TODO: CALL $00D12D / SOUND SOMETHING
    ('UNKFx',       'imm_8'),                  # Fx - TODO: CALL $00D003 / SOUND SOMETHING
)

# Kirby Super Star introduced an alternate, mysterious interpreter!
# Opcodes are 00, 02, 04, 06, ...
# But if the lsb of an opcode is set, the next byte acts as "WAIT #x"
# TODO: When are these used????????????????
OPCODES_ALT = (
    OPCODES[0x00],                              # 00
    OPCODES[0x01],                              # 02
    OPCODES[0x02],                              # 04
    OPCODES[0x05],                              # 06
    OPCODES[0x06],                              # 08
    OPCODES[0x07],                              # 0A
    OPCODES[0x0B],                              # 0C
    OPCODES[0x0C],                              # 0E
    OPCODES[0x0E],                              # 10
    OPCODES[0x0F],                              # 12
    OPCODES[0x10],                              # 14
    OPCODES[0x11],                              # 16
    OPCODES[0x14],                              # 18
    OPCODES[0x15],                              # 1A
    OPCODES[0x16],                              # 1C
    OPCODES[0x17],                              # 1E
    OPCODES[0x18],                              # 20
    OPCODES[0x19],                              # 22
    OPCODES[0x1A],                              # 24
    OPCODES[0x1B],                              # 26
    ('MOVINDEXED.w', 'addr_16 imm_16'),         # 28 addr_16[this.index] = imm_16
    ('MOVINDEXED.w', 'reg addr_16'),            # 2A reg = addr_16[this.index]
    OPCODES[0x20],                              # 2C
    OPCODES[0x21],                              # 2E
    OPCODES[0x33],                              # 30
    OPCODES[0x34],                              # 32
    OPCODES[0x24],                              # 34
    OPCODES_WAITED[0x08-0x04],                  # 36
    OPCODES_WAITED[0x09-0x04],                  # 38
    OPCODES_WAITED[0x0A-0x04],                  # 3A
    OPCODES_WAITED[0x0B-0x04],                  # 3C
    OPCODES_WAITED[0x04-0x04],                  # 3E
    OPCODES_WAITED[0x05-0x04],                  # 40
    OPCODES_WAITED[0x06-0x04],                  # 42
    OPCODES_WAITED[0x07-0x04],                  # 44
    OPCODES[0x35],                              # 46
    ('UNK48',        'imm_16 imm_16'),          # 48
    ('UNK4A',        'addr_24'),                # 4A
    ('UNK4C',        'addr_24'),                # 4C
    OPCODES_WAITED[0x0C-0x04],                  # 4E
    OPCODES_WAITED[0x0D-0x04],                  # 50
    OPCODES_WAITED[0x0E-0x04],                  # 52
    OPCODES_WAITED[0x0F-0x04],                  # 54
    ('UNK56',        ''),                       # 56
    ('UNK58',        'imm_16'),                 # 58
    ('UNK5A',        'imm_16'),                 # 5A
    ('UNK5C',        'imm_16'),                 # 5C
)

BINOPS = ('AND', 'OR', 'ADD', 'XOR')

def snes2pc(addr):
    if addr >= 0xC00000:
        return addr - 0xC00000

    return ((addr & 0x3F0000) >> 1) | (addr & 0x7FFF)

def pc2snes(addr, super_mmc=False):
    if super_mmc:
        return addr + 0xC00000

    return ((addr & 0x3F8000) << 1) | (addr & 0x7FFF) | 0x8000

# Some epic tests because I was STUCK figuring out lorom conversion for 2 days straight like a dummy
'''
# SNES, PC
TESTS = (
    (0x008000, 0x000000),
    (0x00C000, 0x004000),
    (0x018000, 0x008000),
    (0x01C000, 0x00C000),
    (0x028000, 0x010000),
    (0x02C000, 0x014000),
    (0x038000, 0x018000),
    (0x03C000, 0x01C000)
)

for snes, pc in TESTS:
    print(f'snes2pc({snes:06X}) -> {snes2pc(snes):06X} (expect {pc:06X})')

print('-------------------------')
for snes, pc in TESTS:
    print(f'pc2snes({pc:06X}) -> {pc2snes(pc):06X} (expect {snes:06X})')
exit()
'''

class Disassembler(object):
    def __init__(self, rom_file, out_file, sym_file, asm_functions_file, header_offset, script_count, indent=4):
        self._indent = indent
        self.rom_file = rom_file
        self.out_file = out_file
        self.sym_file = sym_file
        self.asm_functions_file = asm_functions_file
        self.header_offset = header_offset
        self.script_count = script_count
        self.indentation = self._indent
        self.pc = 0
        self.was_linebreak = False
        self.force_label = False
        self.super_mmc = False # HACK: SA-1 stuff
        self.preset_asmcalls = []
        self.symbols = dict()
        self.asm_functions = dict()
        self.init_asm_functions()
        self.init_symbols()

    @property
    def snes_pc(self):
        return pc2snes(self.pc, self.super_mmc)

    def read_rom(self, amount, signed=False):
        b = self.rom_file.read(amount)
        return int.from_bytes(b, signed=signed, byteorder='little')

    def datatype_to_str(self, data_type, bytes_):
        value = int.from_bytes(bytes_, byteorder='little')
        if data_type == 'label_16':
            return self.symbols.get((self.snes_pc & 0xFF0000) | value, '${:04X}'.format(value))
        elif data_type == 'label_24':
            return self.symbols.get(value, '${:06X}'.format(value))
        elif data_type == 'imm_8':
            return '#${:02X}'.format(value)
        elif data_type == 'imm_16':
            return '#${:04X}'.format(value)
        elif data_type in ('imm_s8', 'imm_s16'):
            value = int.from_bytes(bytes_, signed=True, byteorder='little')
            return '#{}'.format(value)
        elif data_type in ('imm_u8', 'imm_u16'):
            return '#{}'.format(value)
        elif data_type == 'addr_16':
            return '${:04X}'.format(value)
        elif data_type == 'addr_24':
            return '${:06X}'.format(value)
        elif data_type == 'addr_32':
            value = ((value & 0xFFFF0000) >> 16) | ((value & 0xFFFF) << 16)  # addr_32 has a weird-ass format
            return '${:06X}'.format(value)
        elif data_type == 'obj_var':
            if value > 9:
                print(f'WARNING: INVALID OBJ_VAR OPERAND {value} @ {self.snes_pc:06X}')

            return 'VAR{}'.format(value)
        elif data_type in ('reg', 'nop'):
            return data_type.upper()
        else:
            raise TypeError('Unknown data type! ({})'.format(data_type))

    def try_add_label(self):
        label = self.symbols.get(self.snes_pc)

        if self.force_label:
            if not label:
                label = 'L_{:06X}'.format(self.snes_pc)
                self.symbols[self.snes_pc] = label

            label = '\n' + label

        if label:
            self.out_file.write(label + ':\n')

        return bool(label)

    def disassemble_all(self):
        l = list(self.symbols.keys())
        l.sort()

        self.bad_asmcall = set()
        for start in l:
            pcstart = snes2pc(start)
            self.rom_file.seek(pcstart + self.header_offset)
            self.pc = self.rom_file.tell()

            self.super_mmc = start >= 0xC00000

            while self.snes_pc in self.traversed:
                self.traversed.remove(self.snes_pc)

                was_label = self.try_add_label()
                self.was_linebreak = was_label or self.was_linebreak
                self.force_label = False

                opcode = self.read_rom(1)
                self.disasm_opcode(opcode)
                self.force_label = opcode in (0x00, 0x01, 0x09, 0x0C, 0x0D, 0x0E)
                self.pc = self.rom_file.tell()
                if self.force_label:
                    break

        l = list(self.bad_asmcall)
        l.sort()
        for address in l:
            print(f"WARNING: I DON'T KNOW ANYTHING ABOUT ASMCALL ${address:06X}!")

    def disasm_opcode(self, op_byte):
        indentation = ' ' * self.indentation

        waited = None
        opcode = op_byte
        if opcode >= 0x40:
            waited = opcode & 0x0F
            opcode = ((opcode & 0xF0) >> 4) - 0x04 # Get the actual opcode from the "waited opcode" code
            valid = True
        else:
            valid = (opcode < len(OPCODES))

        if not valid:
            to_write = (indentation + '.byte'.ljust(12) + '${:02X}'.format(opcode)).ljust(40 + len(indentation))
            to_write += '; {:06X}/{:02X}\n'.format(self.snes_pc, opcode)
            self.out_file.write(to_write)
            return

        op_addr = self.snes_pc

        if waited is None:
            mnemonic, types = OPCODES[opcode]
        else:
            mnemonic, types = OPCODES_WAITED[opcode]

        types = types.split()

        bytes_ = bytearray([op_byte])
        operands = []
        for data_type in types:
            b = self.rom_file.read(DATA_TYPE_SIZES[data_type])
            bytes_ += b
            operands.append(self.datatype_to_str(data_type, b))

        if waited is not None:
            operands.append(('' if len(operands) == 0 else ' ') + f'WAIT #{waited}')

        if not self.was_linebreak:
            mnemonic = '\n' + mnemonic  # Add leading newline if previous line didn't have a line break
            self.was_linebreak = True

        comment = ''
        extra = ''

        if mnemonic in ('ASMCALL', 'PRESETCALL'):
            if mnemonic == 'PRESETCALL':
                preset_id = bytes_[1]
                if preset_id >= len(self.preset_asmcalls):
                    address = -1 # Horrible hacks incoming!!!!!!
                    comment = ' // BAD PRESETCALL'
                    print(f'WARNING: BAD PRESETCALL: {preset_id:02X} @ {self.snes_pc:06X}')
                else:
                    address = self.preset_asmcalls[preset_id]
                    comment = f' // ASMCALL ${address:06X}'
            else:
                comment = ''
                address = int.from_bytes(bytes_[1:], byteorder='little')

            asm_func = self.asm_functions.get(address, None)
            if asm_func:
                DIRECTIVES = ('.byte', '.word', '.long', '.dword')
                if asm_func['comment']:
                    comment += ' // ' + asm_func['comment']

                # Fucking hell this is horrible
                for p in asm_func['params']:
                    if p == 'varargs':
                        count = self.read_rom(1)

                        self.pc = self.rom_file.tell()
                        extra += (indentation + '.byte'.ljust(12) + str(count)).ljust(40 + len(indentation))
                        extra += '; {:06X}/{:02X}\n'.format(self.snes_pc, count)

                        for i in range(count):
                            self.pc = self.rom_file.tell()
                            b = self.rom_file.read(2)
                            addr = int.from_bytes(b, byteorder='little')
                            extra += (indentation + '.word'.ljust(12) + str(addr)).ljust(40 + len(indentation))
                            extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
                    else:
                        self.pc = self.rom_file.tell()
                        data_size = DATA_TYPE_SIZES[p]
                        b = self.rom_file.read(data_size)
                        arg = self.datatype_to_str(p, b)
                        arg = arg.replace('#', '')
                        extra += (indentation + DIRECTIVES[data_size - 1].ljust(12) + arg).ljust(40 + len(indentation))
                        extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())

            elif address not in self.bad_asmcall:
                self.bad_asmcall.add(address)
        elif mnemonic == 'BINOP':
            op = bytes_[1 + DATA_TYPE_SIZES[types[0]]] # Ugh.

            if op in (0, 1, 2, 3):
                del operands[1]
                mnemonic = mnemonic.replace('BINOP', BINOPS[op])

                if op == 2:  # ADD, change operand type from HEXADECIMAL immediate to DECIMAL immediate
                    data_type = 'imm_s8' if opcode == 0x1D else 'imm_s16'
                    size = DATA_TYPE_SIZES[data_type]
                    self.rom_file.seek(-size, os.SEEK_CUR)
                    operands[-1] = self.datatype_to_str(data_type, self.rom_file.read(size))
        elif mnemonic in ('MULTIJMP', 'MULTIJSR'):
            count = bytes_[1]

            for i in range(count):
                self.pc = self.rom_file.tell()
                b = self.rom_file.read(2)
                addr = self.datatype_to_str('label_16', b)
                extra += (indentation + '.word'.ljust(12) + addr).ljust(40 + len(indentation))
                extra += '; {:06X}/{}\n'.format(self.snes_pc, b.hex().upper())
        elif mnemonic in ('END', 'RTL', 'HALT', 'ENDTASK', 'RTS'):
            self.indentation = self._indent
            indentation = ' ' * self.indentation
            self.was_linebreak = True
        elif mnemonic == 'LOOP':
            self.indentation += self._indent
        elif mnemonic == 'ENDLOOP':
            self.indentation = max(self._indent, self.indentation - self._indent)
            indentation = ' ' * self.indentation
        elif mnemonic in ('JML', 'JMP'):
            self.was_linebreak = True
            self.indentation = max(self._indent, self.indentation - self._indent)

        to_write = (indentation + mnemonic.ljust(12) + ','.join(operands)).ljust(40 + len(indentation))
        to_write += '; {:06X}/{}{}\n'.format(op_addr, bytes_.hex().upper(), comment)
        self.out_file.write(to_write + extra)

    def init_symbols(self):
        self.traversed = set()
        self.traverse_queue = [] # wtf...

        if self.sym_file:
            self.parse_sym_file()

        # Read preset ASMCALLs...
        self.rom_file.seek(0x384D + self.header_offset)
        for i in range(24):
            self.preset_asmcalls.append(self.read_rom(3))

        # Now onto the script themselves
        self.rom_file.seek(0x3895 + self.header_offset)

        # Add default symbols for each script entry point, if they've not been defined
        for i in range(script_count):
            address = self.read_rom(3)
            flags = self.read_rom(1)
            processor = 'SA-1' if flags & 0x20 == 0 else 'S-CPU'
            if flags & 0x40 == 0x40:
                print(f'"Script" {i:04X}: ASM @ ${address:06X} ({processor})')
            elif address not in self.symbols:
                print(f'Script {i:04X}: actionscript @ ${address:06X} ({processor})')
                label = f'Script{i:04X}'
                self.add_label(address, label)

        while self.traverse_queue:
            address = self.traverse_queue.pop()
            self.traverse(address)

    def add_label(self, address, label, traverse=False):
        if (address >= 0x400000 and address < 0x808000) or (address < 0xC00000 and (address & 0xFFFF) < 0x8000):
            return

        if address not in self.symbols:
            self.symbols[address] = label
            self.traverse_queue.append(address)

    def traverse(self, start):
        old_pc = self.pc
        old_super_mmc = self.super_mmc

        if start < 0x8000: return

        address = snes2pc(start)
        self.rom_file.seek(address + self.header_offset)
        self.pc = self.rom_file.tell()

        self.super_mmc = start >= 0xC00000

        while self.snes_pc not in self.traversed:
            op_byte = self.read_rom(1)

            opcode = op_byte
            if opcode >= 0x40:
                waited = True
                opcode = ((opcode & 0xF0) >> 4) - 0x04  # Get the actual opcode from the "waited opcode" code
                valid = True
            else:
                waited = False
                valid = (opcode < len(OPCODES))

            if not valid:
                break

            self.traversed.add(self.snes_pc)

            if op_byte & 0xF0 in (0xC0, 0xD0): # ASMCALL/PRESETCALL
                if op_byte & 0xF0 == 0xD0: # PRESETCALL
                    preset_id = self.read_rom(1)
                    if preset_id < len(self.preset_asmcalls):
                        addr = self.preset_asmcalls[preset_id]
                    else:
                        addr = -1 # HACK: Dummy
                else: # ASMCALL
                    addr = self.read_rom(3)

                asm_func = self.asm_functions.get(addr, None)
                if asm_func:
                    for p in asm_func['params']:
                        if p == 'varargs':
                            count = self.read_rom(1)
                            self.rom_file.read(count*2)  # Read and discard bytes
                        else:
                            data_size = DATA_TYPE_SIZES[p]
                            b = self.read_rom(data_size)  # Read and discard bytes
                            if p == 'label_24':
                                self.add_label(b, f'L_{b:06X}')
            elif op_byte in (0x0B, 0x0F):  # MULTIJSR/MULTIJMP
                count = self.read_rom(1)
                for i in range(count):
                    addr = (self.snes_pc & 0xFF0000) | self.read_rom(2)
                    self.add_label(addr, f'L_{addr:06X}')
            else:
                if not waited:
                    types = OPCODES[opcode][1].split()
                else:
                    types = OPCODES_WAITED[opcode][1].split()

                for data_type in types:
                    if data_type == 'label_16':
                        addr = (self.snes_pc & 0xFF0000) | self.read_rom(2)
                        self.add_label(addr, f'L_{addr:06X}')
                    elif data_type == 'label_24':
                        addr = self.read_rom(3)
                        self.add_label(addr, f'L_{addr:06X}')
                    else:
                        self.rom_file.read(DATA_TYPE_SIZES[data_type])  # Read and discard bytes

            self.pc = self.rom_file.tell()
            if op_byte in (0x00, 0x01, 0x09, 0x0C, 0x0D, 0x0E):
                break

        self.rom_file.seek(old_pc)
        self.pc = self.rom_file.tell()
        self.super_mmc = old_super_mmc

    def parse_sym_file(self):
        for i, line in enumerate(self.sym_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                label, str_address = [s.strip() for s in pre_comment.split('=', maxsplit=1)]

                if not all(c in LABEL_CHARSET for c in label):
                    print('Ignoring line {} from {}: Invalid label name ({})'.format(i, self.sym_file.name, label), file=sys.stderr)
                    continue

                try:
                    address = int(str_address, 16)
                    self.add_label(address, label)
                except ValueError:
                    print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.sym_file.name, str_address), file=sys.stderr)
            elif pre_comment:
                print('Ignoring line {} from {}: Invalid line'.format(i, self.sym_file.name), file=sys.stderr)

    def init_asm_functions(self):
        if not self.asm_functions_file:
            return

        address = None
        asm_function = None
        for i, line in enumerate(self.asm_functions_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                if address is None:
                    print('Ignoring line {} from {}: Attribute assignment without ASM_FUNCTION directive'.format(i, self.asm_functions_file.name), file=sys.stderr)
                    continue

                key, value = [s.strip() for s in pre_comment.split('=', maxsplit=1)]                
                if key == 'COMMENT':
                    self.asm_functions[address]['comment'] = value
                elif key == 'PARAMS':
                    self.asm_functions[address]['params'] = [s.strip() for s in value.split(',')]
                else:
                    print('Ignoring line {} from {}: Invalid attribute ({})'.format(i, self.asm_functions_file.name, key), file=sys.stderr)
            elif pre_comment:
                directive, arg = [s.strip() for s in pre_comment.split(maxsplit=1)]
                if directive == 'ASM_FUNCTION':
                    try:
                        address = int(arg, 16)
                        self.asm_functions[address] = {'comment': '', 'params': []}
                    except ValueError:
                        print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.asm_functions_file.name, arg), file=sys.stderr)
                else:
                    print('Ignoring line {} from {}: Unknown directive ({})'.format(i, self.asm_functions_file.name, directive), file=sys.stderr)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--symfile', help='the symbols definition file')
    parser.add_argument('-a', '--asmfuncfile', help='the asm functions definition file')
    parser.add_argument('romfile', help='the Earthbound or MOTHER 2 ROM file')
    parser.add_argument('outfile', help='the output file')
    args = parser.parse_args()

    rom_size = os.path.getsize(args.romfile)
    header_offset = rom_size % 0x010000
    if rom_size < 0x100000 or (header_offset != 0 and header_offset != 512):
        print("The file {} doesn't look like a valid Kirby Super Star ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    rom_file = open(args.romfile, 'rb')
    out_file = open(args.outfile, 'w')
    sym_file = open(args.symfile, 'r') if args.symfile else None
    asm_funcs_file = open(args.asmfuncfile, 'r') if args.asmfuncfile else None

    rom_file.seek(0x7FC0 + header_offset)
    rom_name = rom_file.read(21)

    if rom_name != b"KIRBY SUPER DELUXE".ljust(21):
        print("The file {} doesn't look like a valid Kirby Super Star ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    script_count = LAST_SCRIPT
    disassembler = Disassembler(rom_file, out_file, sym_file, asm_funcs_file, header_offset, script_count, indent=4)
    disassembler.disassemble_all()

    rom_file.close()
    out_file.close()
    if sym_file:
        sym_file.close()

    with open('symbols_out.txt', 'w') as f:
        for k, v in disassembler.symbols.items():
            f.write('{:06X}: {}\n'.format(k, v))

KSS/actionscript_dumper.py
CataLatas/hal-actionscript-extractor
#!/usr/bin/python3

# Some of the code you'll see here is bad, but it works.

import argparse
import os
import string
import sys

INES_SIZE = 16

LABEL_CHARSET = string.ascii_letters + string.digits + '_'  # Valid characters for labels

LAST_SCRIPT = 112  # Amount of scripts

# In the format (START, END)
SCRIPT_BLOCKS = (
    (0x14*0x2000 + 0x0000, 0x14*0x2000 + 0x00D5),
    (0x16*0x2000 + 0x0A3A, 0x16*0x2000 + 0x0A96),
    (0x16*0x2000 + 0x0B36, 0x16*0x2000 + 0x0B49),
    (0x16*0x2000 + 0x0BDC, 0x16*0x2000 + 0x0BF8),
    (0x16*0x2000 + 0x0C72, 0x16*0x2000 + 0x0C8D),
    (0x16*0x2000 + 0x0C9A, 0x16*0x2000 + 0x0CBD),
    (0x16*0x2000 + 0x0D15, 0x16*0x2000 + 0x0D1F),
    (0x16*0x2000 + 0x0D3C, 0x16*0x2000 + 0x0D63),
    (0x16*0x2000 + 0x0D97, 0x16*0x2000 + 0x0DEA),
    (0x18*0x2000 + 0x0311, 0x18*0x2000 + 0x040C),
    (0x3D*0x2000 + 0x0000, 0x3D*0x2000 + 0x02D5),
    (0x3D*0x2000 + 0x09EB, 0x3D*0x2000 + 0x0A4E),
    (0x3D*0x2000 + 0x09EB, 0x3D*0x2000 + 0x0A4E),
    (0x3D*0x2000 + 0x1F50, 0x3D*0x2000 + 0x1F57),
)

DATA_TYPE_SIZES = {
    'label_16': 2,  # 16-bit address (prefferably a label, when available)
    'label_24': 3,  # 24-bit address (prefferably a label, when available)
    'imm_8':    1,  # 8-bit immediate (hex representation)
    'imm_s8':   1,  # 8-bit signed immediate (decimal representation)
    'imm_u8':   1,  # 8-bit unsigned immediate (decimal representation)
    'imm_16':   2,  # 16-bit immediate (hex representation)
    'imm_s16':  2,  # 16-bit signed immediate (decimal representation)
    'imm_u16':  2,  # 16-bit unsigned immediate (decimal representation)
    'addr_16':  2,  # 16-bit address
    'addr_24':  3,  # 24-bit address
    'addr_32':  4,  # 32-bit address (high byte ignored)
    'obj_var':  1,  # Object variable (i.e. VAR0, VAR1, ..., VAR9)
    'reg':      0,  # The work register
    'nop':      0,  # Special, used only by opcode 0x0F (ONTICK NOP)
}

# (mnemonic, data types separated by spaces)
OPCODES = (
    ('END',         ''),                       # 00
    ('LOOP',        'imm_u8'),                 # 01
    ('ENDLOOP',     ''),                       # 02
    ('JML',         'label_24'),               # 03
    ('JSL',         'label_24'),               # 04
    ('RTL',         ''),                       # 05
    ('WAIT',        'imm_u8'),                 # 06
    ('TASK',        'label_16'),               # 07
    ('ONTICK',      'addr_24'),                # 08
    ('HALT',        ''),                       # 09
    ('JEQ',         'label_16'),               # 0A
    ('JNE',         'label_16'),               # 0B
    ('ENDTASK',     ''),                       # 0C
    ('MOV',         'obj_var imm_8'),          # 0D
    ('ONTICK',      'nop'),                    # 0E
    ('MULTIJMP',    'imm_u8'),                 # 0F
    ('MULTIJSR',    'imm_u8'),                 # 10
    ('MOV.b',       'addr_16 imm_8'),          # 11
    ('ENDLASTTASK', ''),                       # 12
    ('BINOP',       'obj_var imm_u8 imm_8'),   # 13
    ('BREAKEQ',     'label_16'),               # 14
    ('BREAKNE',     'label_16'),               # 15
    ('BINOP',       'addr_16 imm_u8 imm_8'),   # 16
    ('JMP',         'label_16'),               # 17
    ('JSR',         'label_16'),               # 18
    ('RTS',         ''),                       # 19
    ('SPRITEMAP',   'addr_24'),                # 1A
    ('MOV',         'reg imm_8'),              # 1B
    ('MOV',         'reg addr_16'),            # 1C
    ('MOV',         'obj_var reg'),            # 1D
    ('MOV',         'reg obj_var'),            # 1E
    ('WAIT',        'obj_var'),                # 1F
    ('ONDRAW',      'addr_16'),                # 20
    ('ONPOSITION',  'addr_16'),                # 21
    ('LOOP',        'reg'),                    # 22
    ('ONMOVE',      'addr_16'),                # 23
    ('SETPOSE',     'obj_var'),                # 24
    ('BINOP',       'reg imm_u8 imm_8'),       # 25
    ('ASMCALL.l',   'addr_24'),                # 26
    ('MOV.w',       'addr_16 imm_16'),         # 27
    ('SETBANK',     'imm_8'),                  # 28
    ('MULTICALL',   'imm_u8'),                 # 29
    ('SETXPOS',     'imm_16'),                 # 2A
    ('SETYPOS',     'imm_16'),                 # 2B
    ('ADDXPOS',     'imm_s16'),                # 2C
    ('ADDYPOS',     'imm_s16'),                # 2D
    ('ADDXVEL',     'imm_16'),                 # 2E
    ('ADDYVEL',     'imm_16'),                 # 2F
    ('UNK30',       'imm_8 imm_8'),            # 30
    ('UNK31',       'imm_8 imm_8'),            # 31
    ('UNK32',       'imm_8 imm_8'),            # 32
    ('UNK33',       'imm_8 imm_8'),            # 33
    ('UNK34',       'imm_8 imm_8'),            # 34
    ('UNK35',       'imm_8 imm_8'),            # 35
    ('UNK36',       'imm_8 imm_8'),            # 36
    ('UNK37',       'imm_8 imm_8'),            # 37
    ('ZEROVEL',     ''),                       # 38
    ('UNK39',       'imm_8'),                  # 39
    ('SETZPOS',     'imm_16'),                 # 3A
    ('ADDZPOS',     'imm_s16'),                # 3B
    ('SETZVEL',     'imm_16'),                 # 3C
    ('ADDZVEL',     'imm_16'),                 # 3D
    ('MULTIJSL',    'imm_u8'),                 # 3E
)

OPCODES_WAITED = (
    ('SETPOSE', 'imm_8'),   # 5x
    ('ADDPOSE', 'imm_s8'),  # 6x
    ('INCPOSE', ''),        # 7x
    ('DECPOSE', ''),        # 8x
    ('INC2POSE', ''),       # 9x
    ('DEC2POSE', ''),       # Ax
    ('SETXVEL', 'imm_16'),  # Bx
    ('SETYVEL', 'imm_16'),  # Cx
    ('ASMCALL', 'addr_16'), # Dx
)

BINOPS = ('AND', 'OR', 'ADD', 'XOR')

class Disassembler(object):
    def __init__(self, rom_file, out_file, sym_file, asm_functions_file, script_count, indent=4):
        self._indent = indent
        self.rom_file = rom_file
        self.out_file = out_file
        self.sym_file = sym_file
        self.asm_functions_file = asm_functions_file
        self.script_count = script_count
        self.indentation = self._indent
        self.pc = 0
        self.was_linebreak = False
        self.force_label = False
        self.symbols = dict()
        self.asm_functions = dict()
        self.init_asm_functions()
        self.init_symbols()

    @property
    def nes_pc(self):
        pc = self.pc - INES_SIZE
        return 0xA000 + (((pc >> 13) << 16) | (pc & 0x1FFF))

    def read_rom(self, amount, signed=False, byteorder='little'):
        b = self.rom_file.read(amount)
        return int.from_bytes(b, signed=signed, byteorder=byteorder)

    def datatype_to_str(self, data_type, bytes_):
        value = int.from_bytes(bytes_, byteorder='little')
        if data_type == 'label_16':
            return self.symbols.get((self.nes_pc & 0xFF0000) | value, '${:04X}'.format(value))
        elif data_type == 'label_24':
            return self.symbols.get(value, '${:06X}'.format(value))
        elif data_type == 'imm_8':
            return '#${:02X}'.format(value)
        elif data_type == 'imm_16':
            return '#${:04X}'.format(value)
        elif data_type in ('imm_s8', 'imm_s16'):
            value = int.from_bytes(bytes_, signed=True, byteorder='little')
            return '#{}'.format(value)
        elif data_type in ('imm_u8', 'imm_u16'):
            return '#{}'.format(value)
        elif data_type == 'addr_16':
            return '${:04X}'.format(value)
        elif data_type == 'addr_24':
            return '${:06X}'.format(value)
        elif data_type == 'addr_32':
            value = ((value & 0xFFFF0000) >> 16) | ((value & 0xFFFF) << 16)  # addr_32 has a weird-ass format
            return '${:06X}'.format(value)
        elif data_type == 'obj_var':
            if value > 10:
                print(f'WARNING: INVALID OBJ_VAR OPERAND {value} @{self.nes_pc:06X}')

            return 'VAR{}'.format(value)
        elif data_type in ('reg', 'nop'):
            return data_type.upper()
        else:
            raise TypeError('Unknown data type! ({})'.format(data_type))

    def try_add_label(self):
        label = self.symbols.get(self.nes_pc)

        if self.force_label:
            if not label:
                label = 'L_{:06X}'.format(self.nes_pc)
                self.symbols[self.nes_pc] = label

            label = '\n' + label

        if label:
            self.out_file.write(label + ':\n')

        return bool(label)

    def disassemble_all(self):
        l = list(self.symbols.keys())
        l.sort()

        self.bad_asmcall = set()
        for start in l:
        #for start, end in SCRIPT_BLOCKS:
            #end += INES_SIZE
            #size += end - start
            pcstart = (start >> 16) * 0x2000 + (start & 0x1FFF)
            self.rom_file.seek(pcstart + INES_SIZE)
            self.pc = self.rom_file.tell()

            while self.nes_pc in self.traversed:
                self.traversed.remove(self.nes_pc)

                was_label = self.try_add_label()
                self.was_linebreak = was_label or self.was_linebreak
                self.force_label = False

                opcode = self.read_rom(1)
                self.disasm_opcode(opcode)
                self.force_label = opcode in (0x00, 0x03, 0x05, 0x09, 0x0C, 0x17, 0x19)
                self.pc = self.rom_file.tell()
                if self.force_label:
                    break

        l = list(self.bad_asmcall)
        l.sort()
        for address in l:
            print(f"WARNING: I DON'T KNOW ANYTHING ABOUT ASMCALL ${address:06X}!")

    def disasm_opcode(self, op_byte):
        indentation = ' ' * self.indentation

        waited = None
        opcode = op_byte
        if opcode >= 0x50:
            waited = opcode & 0x0F
            opcode = ((opcode - 0x50) >> 4)  # Get the actual opcode from the "waited opcode" code
            valid = (opcode <= 0x0D-0x05)
        else:
            valid = (opcode < len(OPCODES))

        if not valid:
            to_write = (indentation + '.byte'.ljust(12) + '${:02X}'.format(op_byte)).ljust(40 + len(indentation))
            to_write += '; {:06X}/{:02X}\n'.format(self.nes_pc, op_byte)
            self.out_file.write(to_write)
            return

        op_addr = self.nes_pc

        if waited is None:
            mnemonic, types = OPCODES[opcode]
        else:
            mnemonic, types = OPCODES_WAITED[opcode]

        types = types.split()

        bytes_ = bytearray([op_byte])
        operands = []
        for data_type in types:
            b = self.rom_file.read(DATA_TYPE_SIZES[data_type])
            bytes_ += b
            operands.append(self.datatype_to_str(data_type, b))

        if waited is not None:
            operands.append(('' if len(operands) == 0 else ' ') + f'WAIT #{waited}')

        if not self.was_linebreak:
            mnemonic = '\n' + mnemonic  # Add leading newline if previous line didn't have a line break
            self.was_linebreak = True

        comment = ''
        extra = ''

        if mnemonic in ('ASMCALL', 'ASMCALL.l'):
            if mnemonic == 'ASMCALL.l':
                address = int.from_bytes(bytes_[1:], byteorder='little')
            elif self.pc in self.lower_asmcalls:
                address = self.lower_asmcalls[self.pc]
            else:
                address = int.from_bytes(bytes_[1:], byteorder='little')
                if address < 0xC000:
                    address = (self.nes_pc & 0xFF0000) | address

            asm_func = self.asm_functions.get(address, None)
            if asm_func:
                DIRECTIVES = ('.byte', '.word', '.long', '.dword')
                comment = ' // ' + asm_func['comment'] if asm_func['comment'] else ''
                
                # Fucking hell this is horrible
                for p in asm_func['params']:
                    if p == 'varargs':
                        count = self.read_rom(1)
                        
                        self.pc = self.rom_file.tell()
                        extra += (indentation + '.byte'.ljust(12) + str(count)).ljust(40 + len(indentation))
                        extra += '; {:06X}/{:02X}\n'.format(self.nes_pc, count)
                        
                        for i in range(count):
                            self.pc = self.rom_file.tell()
                            b = self.rom_file.read(2)
                            addr = int.from_bytes(b, byteorder='little')
                            extra += (indentation + '.word'.ljust(12) + str(addr)).ljust(40 + len(indentation))
                            extra += '; {:06X}/{}\n'.format(self.nes_pc, b.hex().upper())
                    else:
                        self.pc = self.rom_file.tell()
                        data_size = DATA_TYPE_SIZES[p]
                        b = self.rom_file.read(data_size)
                        arg = self.datatype_to_str(p, b)
                        arg = arg.replace('#', '')
                        extra += (indentation + DIRECTIVES[data_size - 1].ljust(12) + arg).ljust(40 + len(indentation))
                        extra += '; {:06X}/{}\n'.format(self.nes_pc, b.hex().upper())
                        
            elif address not in self.bad_asmcall:
                self.bad_asmcall.add(address)
        elif mnemonic == 'BINOP':
            # Now THIS is what I call hacky shit!
            if opcode == 0x25:
                op = bytes_[1]
            elif opcode == 0x13:
                op = bytes_[2]
            else:
                op = bytes_[3]

            if op in (0, 1, 2, 3):
                del operands[1]
                mnemonic = mnemonic.replace('BINOP', BINOPS[op])
                
                if op == 2:  # ADD, change operand type from HEXADECIMAL immediate to DECIMAL immediate
                    data_type = 'imm_s8'
                    size = DATA_TYPE_SIZES[data_type]
                    self.rom_file.seek(-size, os.SEEK_CUR)
                    operands[-1] = self.datatype_to_str(data_type, self.rom_file.read(size))
        elif mnemonic in ('MULTIJMP', 'MULTIJSR', 'MULTICALL', 'MULTIJSL'):
            count = bytes_[1]

            if mnemonic in ('MULTICALL', 'MULTIJSL'):
                data_type = 'label_24'
                directive = '.long'
                data_size = 3
            else:
                data_type = 'label_16'
                directive = '.word'
                data_size = 2

            for i in range(count):
                self.pc = self.rom_file.tell()
                b = self.rom_file.read(data_size)
                addr = self.datatype_to_str(data_type, b)
                extra += (indentation + directive.ljust(12) + addr).ljust(40 + len(indentation))
                extra += '; {:06X}/{}\n'.format(self.nes_pc, b.hex().upper())
        elif mnemonic in ('END', 'RTL', 'HALT', 'ENDTASK', 'RTS'):
            self.indentation = self._indent
            indentation = ' ' * self.indentation
            self.was_linebreak = True
        elif mnemonic == 'LOOP':
            self.indentation += self._indent
        elif mnemonic == 'ENDLOOP':
            self.indentation = max(self._indent, self.indentation - self._indent)
            indentation = ' ' * self.indentation
        elif mnemonic in ('JML', 'JMP'):
            self.was_linebreak = True
            self.indentation = max(self._indent, self.indentation - self._indent)

        to_write = (indentation + mnemonic.ljust(12) + ','.join(operands)).ljust(40 + len(indentation))
        to_write += '; {:06X}/{}{}\n'.format(op_addr, bytes_.hex().upper(), comment)
        self.out_file.write(to_write + extra)

    def init_symbols(self):
        self.traversed = set()
        self.traverse_queue = [] # wtf...

        if self.sym_file:
            self.parse_sym_file()

        self.rom_file.seek(0x18*0x2000 + 0x00A7 + INES_SIZE)
        lo = self.rom_file.read(self.script_count)
        hi = self.rom_file.read(self.script_count)
        bank = self.rom_file.read(self.script_count)

        entry_points = set()

        self.lower_asmcalls = {}
        # Add default symbols for each script entry point, if they've not been defined
        for i in range(self.script_count):
            address = (bank[i] << 16) | (hi[i] << 8) | lo[i]
            entry_points.add(address)
            print(f'Script{i:02X} @ ${address:06X}')
            self.add_label(address, f'Script{i:02X}')

        while self.traverse_queue:
            address = self.traverse_queue.pop()
            if address in entry_points:
                self.lower_prg = None

            self.traverse(address)

    def add_label(self, address, label, traverse=False):
        if (address >> 16) < 0x14 or (address >> 16) >= 0x3E or (address & 0xFFFF) < 0xA000:
            return

        if address not in self.symbols:
            self.symbols[address] = label
            self.traverse_queue.append(address)

    def traverse(self, start):
        old_pc = self.pc

        if (start & 0xFFFF) >= 0xC000:
            return

        address = (start >> 16) * 0x2000 + (start & 0x1FFF)
        self.rom_file.seek(address + INES_SIZE)
        self.pc = self.rom_file.tell()

        while self.nes_pc not in self.traversed:
            op_byte = self.read_rom(1)

            opcode = op_byte
            if opcode >= 0x50:
                waited = True
                opcode = ((opcode - 0x50) >> 4) # Get the actual opcode from the "waited opcode" code
                valid = (opcode <= 0x0D-0x05)
            else:
                waited = False
                valid = (opcode < len(OPCODES))

            if not valid:
                break

            self.traversed.add(self.nes_pc)

            if op_byte == 0x07:  # TASK
                addr = (self.nes_pc & 0xFF0000) | self.read_rom(2)
                self.add_label(addr, f'TASK_{addr:06X}')
            elif (op_byte & 0xF0 == 0xD0) or op_byte == 0x26:  # ASMCALL/ASMCALL.l
                if op_byte == 0x26:
                    addr = self.read_rom(3)
                else:
                    addr = self.read_rom(2)
                    if addr >= 0x8000 and addr < 0xA000:
                        bank = self.lower_prg
                        if self.lower_prg is None:
                            print(f'WTF???? LOWER PRG IS NONE AND WE ENCOUNTERED ASMCALL {addr:04X}')
                            bank = self.nes_pc >> 16

                        addr = (bank << 16) | addr
                        self.lower_asmcalls[self.pc] = addr
                    elif addr < 0xC000:
                        addr = (self.nes_pc & 0xFF0000) | addr

                asm_func = self.asm_functions.get(addr, None)
                if asm_func:
                    for p in asm_func['params']:
                        if p == 'varargs':
                            count = self.read_rom(1)
                            self.rom_file.read(count*2)  # Read and discard bytes
                        else:
                            data_size = DATA_TYPE_SIZES[p]
                            b = self.rom_file.read(data_size)  # Read and discard bytes
            elif op_byte in (0x0F, 0x10):  # MULTIJMP/MULTIJSR
                count = self.read_rom(1)
                for i in range(count):
                    addr = (self.nes_pc & 0xFF0000) | self.read_rom(2)
                    self.add_label(addr, f'L_{addr:06X}')
            elif op_byte in (0x29, 0x3E): # MULTICALL/MULTIJSL
                count = self.read_rom(1)
                for i in range(count):
                    addr = self.read_rom(3)
                    self.add_label(addr, f'L_{addr:06X}')
            elif op_byte == 0x28: # SETBANK
                self.lower_prg = self.read_rom(1)
            else:
                if not waited:
                    types = OPCODES[opcode][1].split()
                else:
                    types = OPCODES_WAITED[opcode][1].split()

                for data_type in types:
                    if data_type == 'label_16':
                        addr = (self.nes_pc & 0xFF0000) | self.read_rom(2)
                        self.add_label(addr, f'L_{addr:06X}')
                    elif data_type == 'label_24':
                        addr = self.read_rom(3)
                        self.add_label(addr, f'L_{addr:06X}')
                    else:
                        self.rom_file.read(DATA_TYPE_SIZES[data_type])  # Read and discard bytes

            self.pc = self.rom_file.tell()
            if op_byte in (0x00, 0x03, 0x05, 0x09, 0x0C, 0x17, 0x19):
                break


        self.rom_file.seek(old_pc)
        self.pc = self.rom_file.tell()

    def parse_sym_file(self):
        for i, line in enumerate(self.sym_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                label, str_address = [s.strip() for s in pre_comment.split('=', maxsplit=1)]

                if not all(c in LABEL_CHARSET for c in label):
                    print('Ignoring line {} from {}: Invalid label name ({})'.format(i, self.sym_file.name, label), file=sys.stderr)
                    continue

                try:
                    address = int(str_address, 16)
                    self.add_label(address, label)
                    #self.symbols[address] = label
                except ValueError:
                    print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.sym_file.name, str_address), file=sys.stderr)
            elif pre_comment:
                print('Ignoring line {} from {}: Invalid line'.format(i, self.sym_file.name), file=sys.stderr)

    def init_asm_functions(self):
        if not self.asm_functions_file:
            return

        address = None
        asm_function = None
        for i, line in enumerate(self.asm_functions_file):
            pre_comment = line.split(';', maxsplit=1)[0].strip()  # Get everything on the line before the comment
            if '=' in pre_comment:
                if address is None:
                    print('Ignoring line {} from {}: Attribute assignment without ASM_FUNCTION directive'.format(i, self.asm_functions_file.name), file=sys.stderr)
                    continue

                key, value = [s.strip() for s in pre_comment.split('=', maxsplit=1)]
                if key == 'COMMENT':
                    self.asm_functions[address]['comment'] = value
                elif key == 'PARAMS':
                    self.asm_functions[address]['params'] = [s.strip() for s in value.split(',')]
                else:
                    print('Ignoring line {} from {}: Invalid attribute ({})'.format(i, self.asm_functions_file.name, key), file=sys.stderr)
            elif pre_comment:
                directive, arg = [s.strip() for s in pre_comment.split(maxsplit=1)]
                if directive == 'ASM_FUNCTION':
                    try:
                        address = int(arg, 16)
                        self.asm_functions[address] = {'comment': '', 'params': []}
                    except ValueError:
                        print('Ignoring line {} from {}: Invalid address ({})'.format(i, self.asm_functions_file.name, arg), file=sys.stderr)
                else:
                    print('Ignoring line {} from {}: Unknown directive ({})'.format(i, self.asm_functions_file.name, directive), file=sys.stderr)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--symfile', help='the symbols definition file')
    parser.add_argument('-a', '--asmfuncfile', help='the asm functions definition file')
    parser.add_argument('romfile', help="the Kirby's Adventure ROM file")
    parser.add_argument('outfile', help='the output file')
    args = parser.parse_args()

    rom_size = os.path.getsize(args.romfile)
    if rom_size != 0xC0010:
        print("The file {} doesn't look like a valid Kirby's Adventure ROM!".format(args.romfile), file=sys.stderr)
        sys.exit(1)

    rom_file = open(args.romfile, 'rb')
    out_file = open(args.outfile, 'w')
    sym_file = open(args.symfile, 'r') if args.symfile else None
    asm_funcs_file = open(args.asmfuncfile, 'r') if args.asmfuncfile else None

    # TODO: Check if this is actually a Kirby's Adventure ROM (how??)

    disassembler = Disassembler(rom_file, out_file, sym_file, asm_funcs_file, LAST_SCRIPT, indent=4)
    disassembler.disassemble_all()

    rom_file.close()
    out_file.close()
    if sym_file:
        sym_file.close()

    with open('symbols_out.txt', 'w') as f:
        for k, v in disassembler.symbols.items():
            f.write('{:06X}: {}\n'.format(k, v))

KirbysAdventure/actionscript_dumper.py
CataLatas/hal-actionscript-extractor