Mario Golf: Super Rush NRO extractor

mgsr_nro_extractor.py

# a massive kludge by Ninji
# extracts the encrypted NRO from Mario Golf: Super Rush
#
# requires Python 3 and the lz4 and unicorn modules

import struct
import lz4.block
from unicorn import *
from unicorn.arm64_const import *

emu = Uc(UC_ARCH_ARM64, UC_MODE_ARM)
def map_section(hdr, file_size, is_compressed):
    file_offset, mem_offset, decomp_size = hdr
    print(f'mapping offset={mem_offset:x} size={decomp_size:x}')
    emu.mem_map(mem_offset, (decomp_size + 0xFFF) & ~0xFFF)

    nso.seek(file_offset)
    blob = nso.read(file_size)
    if is_compressed:
        blob = lz4.block.decompress(blob, uncompressed_size=decomp_size)
    emu.mem_write(mem_offset, blob)
    return blob

alloc_addr = 0x10000000
alloc_history = []
def alloc(size):
    global alloc_addr
    rounded_size = (size + 0xFFF) & ~0xFFF
    addr = alloc_addr
    alloc_addr += rounded_size
    alloc_history.append((addr, size))
    emu.mem_map(addr, rounded_size)
    return addr

# read NSO0
nso = open('main', 'rb')
header = struct.unpack('<4sIII', nso.read(0x10))
text_hdr = struct.unpack('<III', nso.read(0xC))
module_name_offset = struct.unpack('<I', nso.read(4))[0]
rodata_hdr = struct.unpack('<III', nso.read(0xC))
module_name_size = struct.unpack('<I', nso.read(4))[0]
data_hdr = struct.unpack('<III', nso.read(0xC))
bss_size = struct.unpack('<I', nso.read(4))[0]
module_id = nso.read(0x20)
text_file_size = struct.unpack('<I', nso.read(4))[0]
ro_file_size = struct.unpack('<I', nso.read(4))[0]
data_file_size = struct.unpack('<I', nso.read(4))[0]
_ = nso.read(0x1C)
api_info_hdr = struct.unpack('<II', nso.read(8))
dynstr_hdr = struct.unpack('<II', nso.read(8))
dynsym_hdr = struct.unpack('<II', nso.read(8))

text = map_section(text_hdr, text_file_size, (header[3] & 1) != 0)
rodata = map_section(rodata_hdr, ro_file_size, (header[3] & 2) != 0)
data = map_section(data_hdr, data_file_size, (header[3] & 4) != 0)

# parse the dynsym
print(f'dynsym: {dynsym_hdr[0]:x},{dynsym_hdr[1]:x}')
print(f'dynstr: {dynstr_hdr[0]:x},{dynstr_hdr[1]:x}')
dynsym = rodata[dynsym_hdr[0]:dynsym_hdr[0]+dynsym_hdr[1]]
dynstr = rodata[dynstr_hdr[0]:dynstr_hdr[0]+dynstr_hdr[1]]

# parse the MOD0
dynamic_offset = struct.unpack_from('<I', text, 0xC)[0] + 8
print(f'dynamic: {dynamic_offset:x}')

# parse the dynamic section
rela = 0
rela_size = 0
rela_plt = 0
rela_plt_size = 0
while True:
    tag, value = struct.unpack('<QQ', emu.mem_read(dynamic_offset, 16))
    print(f'tag {tag} = {value:x}')
    if tag == 0:
        break
    elif tag == 7:
        rela = value
    elif tag == 8:
        rela_size = value // 0x18
    elif tag == 23:
        rela_plt = value
    elif tag == 2:
        rela_plt_size = value // 0x18
    dynamic_offset += 16

symbols = []
symbols_by_name = {}
symbols_by_addr = {}
hle_hooks = {}

for i in range(len(dynsym) // 0x18):
    st_name, st_info, st_other, st_shndx, st_value, st_size = struct.unpack_from('<IBBHQQ', dynsym, i * 0x18)
    st_name = dynstr[st_name:dynstr.index(0, st_name)].decode('ascii')
    symbols.append(st_name)
    if st_shndx == 2:
        symbols_by_name[st_name] = st_value
        symbols_by_addr[st_value] = st_name
    else:
        pass #print(f'{i}: {st_shndx:04x} {st_value:08x} {st_size:08x} {st_name}')

plt_surrogates = alloc(0x10000)

# handle rela
for i in range(rela_size):
    offset, info, addend = struct.unpack_from('<QQQ', emu.mem_read(rela + i * 0x18, 0x18))
    r_type = info & 0xffffffff
    r_sym = info >> 32
    if info == 0x403:
        emu.mem_write(offset, struct.pack('<Q', addend))

# handle rela plt
for i in range(rela_plt_size):
    addr, _, sym_index = struct.unpack_from('<QII', emu.mem_read(rela_plt + i * 0x18, 0x10))
    #print(f'{addr:x} -> {symbols[sym_index]}')
    surrogate = plt_surrogates + sym_index * 4
    emu.mem_write(addr, struct.pack('<Q', surrogate))
    symbols_by_name[symbols[sym_index]] = surrogate
    symbols_by_addr[surrogate] = symbols[sym_index]

# execute some code
def hook_block(uc, address, size, user_data):
    hook = hle_hooks.get(address)
    if hook != None:
        hook()
        emu.reg_write(UC_ARM64_REG_PC, emu.reg_read(UC_ARM64_REG_LR))
        return
    if address >= plt_surrogates and address < (plt_surrogates + 0x10000):
        raise ValueError(f'unhandled import {symbols_by_addr[address]}')

    try:
        name = symbols_by_addr[address]
        print(f'{name}@{address:x} size={size:x}')
    except KeyError:
        pass #print(f'block@{address:x} size={size:x}')

def hook_aligned_alloc():
    size = emu.reg_read(UC_ARM64_REG_X1)
    addr = alloc(size)
    print(f'alloc size:{size:x} to {addr:x}')
    emu.reg_write(UC_ARM64_REG_X0, addr)
hle_hooks[symbols_by_name['aligned_alloc']] = hook_aligned_alloc

def hook_null():
    pass

class Thread:
    def __init__(self, pc, stack, arg):
        old_ctx = emu.context_save()
        emu.reg_write(UC_ARM64_REG_PC, pc)
        emu.reg_write(UC_ARM64_REG_SP, stack)
        emu.reg_write(UC_ARM64_REG_X0, arg)
        self.ctx = emu.context_save()
        emu.context_restore(old_ctx)

    def exec(self):
        emu.context_restore(self.ctx)
        pc = emu.reg_read(UC_ARM64_REG_PC)
        print(f'  started at pc={pc:x}')
        emu.emu_start(pc, symbols_by_name['_ZN2nn2os12AwaitBarrierEPNS0_11BarrierTypeE'], 0, 0)
        emu.reg_write(UC_ARM64_REG_PC, emu.reg_read(UC_ARM64_REG_LR))
        pc = emu.reg_read(UC_ARM64_REG_PC)
        print(f'  ended at pc={pc:x}')
        emu.context_update(self.ctx)


threads = []
barriers = {}
def hook_InitBarrier():
    barrier_addr = emu.reg_read(UC_ARM64_REG_X0)
    barrier_count = emu.reg_read(UC_ARM64_REG_W1)
    barriers[barrier_addr] = barrier_count

def hook_CreateThread():
    func_addr = emu.reg_read(UC_ARM64_REG_X1)
    arg = emu.reg_read(UC_ARM64_REG_X2)
    threads.append(Thread(func_addr, alloc(0x4000) + 0x4000, arg))

hle_hooks[symbols_by_name['_ZN2nn2os13GetSystemTickEv']] = hook_null
hle_hooks[symbols_by_name['_ZN2nn2os17InitializeBarrierEPNS0_11BarrierTypeEi']] = hook_InitBarrier
hle_hooks[symbols_by_name['_ZN2nn2os12CreateThreadEPNS0_10ThreadTypeEPFvPvES3_S3_mii']] = hook_CreateThread
hle_hooks[symbols_by_name['_ZN2nn2os11StartThreadEPNS0_10ThreadTypeE']] = hook_null

stack_addr = alloc(2 * 1024 * 1024)
emu.reg_write(UC_ARM64_REG_SP, stack_addr + 1024 * 1024)

emu.hook_add(UC_HOOK_BLOCK, hook_block)
emu.emu_start(symbols_by_name['nnMain'], symbols_by_name['_ZN2nn2fs12SetAllocatorEPFPvmEPFvS1_mE'], 0, 0)

for i in range(len(barriers)):
    for j, thread in enumerate(threads):
        print(f'{i} - exec thread {j}:')
        thread.exec()

# find the largest alloc
alloc_history.sort(key=lambda a: a[1])
addr, size = alloc_history[-1]
with open('dumped.nro', 'wb') as f:
    f.write(emu.mem_read(addr, size))