PoC Mario Kart Texture Replacement

#!/usr/bin/env python

# order is:
#  offset to 0x50 resource virtual addresses
#  offset to resource entry tables
#  offset to uncompressed root files
#  offset to compressed root files
#  pairs for start of first file
#  pairs for start of second file (and end of first)
#  pairs for start of third file (and end of second)
#  pairs for start of fourth file (and end of third)
#  pairs for root table
#  pairs for resource table
#  pairs for whatever the float table is
#  pairs for bank 0A
#  pairs for bank 0B (TKMK)
#  pairs for ???
#  pairs for ???
#  pairs for end of ??? / start of first offsets in res. tbl. entries
#  pairs for start of res. tbl. entry 4's first offset
#  pairs for end of res. tbl. entry 4's first offset
#  pairs for sound.ctl
#  pairs for sound.tbl
#  pairs of hardcoded resource virtual addresses
NKTE = (0xE60D0,
     0x2AE80,
     0xDE7F4,
     0xDEAC0,
    ((0x1CF0, 0x1CF8),),
    ((0x1C78, 0x1C80),
     (0x1CF4, 0x1CFC),),
    ((0x1C7C, 0x1C84),
     (0x1F10, 0x1F20),),
    ((0x1F14, 0x1F18),
     (0x1F30, 0x1F3C),),
    ((0x1F34, 0x1F38),
     (0x5B2C, 0x5B28),
     (0x27DB8, 0x27DBC),
     (0x27EC4, 0x27EC8),
     (0x27FDC, 0x27FE0),
     (0x28100, 0x28104),
     (0x282D8, 0x282DC),
     (0x283CC, 0x283D0),
     (0x284CC, 0x284D0),
     (0x285D8, 0x285DC),
     (0x286C0, 0x286D8),
     (0x28768, 0x28780),
     (0x28814, 0x28820),
     (0x288AC, 0x288B8),),
    ((0x6F9D4, 0x6F9E0),
     (0x75190, 0x7519C),
     (0x1119A8, 0x1119B4),
     (0x111A80, 0x111A8C),
     (0x111B58, 0x111B64),
     (0x111C34, 0x111C40),),
    ((0x1ED0, 0x1EDC),),
    ((0x99C30, 0x99C3C),
     (0x9A7C4, 0x9A7CC),
     (0x9AB78, 0x9AB80),),
    ((0x99CB8, 0x99CC4),),
    ((0x123B20, 0x123B2C),
     (0x125080, 0x12508C),),
    ((0x9280C, 0x92818),
     (0x123B24, 0x123B28),
     (0x125084, 0x125088),),
    ((0x92810, 0x92814),),
    ((0x1250A4, 0x1250B4),),
    ((0x1250A8, 0x1250AC),),
    ((0xBCFE0, 0xBCFEC),),
    ((0xBD0D4, 0xBD0D8),),
    ((0xBCF64, 0xBCF68),),
    ((0xBD150, 0xBD158),),
    ((0x107970, 0x107974),
     (0x107984, 0x107988),
     (0x107780, 0x107784),
     (0x107798, 0x1077A0),
     (0x1077B0, 0x1077B4),
     (0x1077C4, 0x1077C8),
     (0x1077D8, 0x1077DC),
     (0x1077EC, 0x1077F0),
     (0x107800, 0x107804),
     (0x107814, 0x107818),
     (0x107828, 0x10782C),
     (0x107840, 0x107848),
     (0x107858, 0x10785C),
     (0x10786C, 0x107870),
     (0x107880, 0x107884),
     (0x107894, 0x107898),
     (0x1078A8, 0x1078AC),
     (0x1078BC, 0x1078C0),
     (0x3D10, 0x3D14),
     (0x1079C0, 0x1079C4),
     (0x107AB4, 0x107AB8),
     (0x107B00, 0x107B04),
     (0x3DCC, 0x3DD0),
     (0x107C1C, 0x107C20),
     (0x3DE0, 0x3DE4),
     (0x107C30, 0x107C34),
     (0x107BEC, 0x107BF0),
     (0x107C00, 0x107C04),
     (0x107B14, 0x107B18),
     (0x107A98, 0x107A9C),
     (0x107AD0, 0x107AD4),
     (0x107AE4, 0x107AE8),
     (0x3DF4, 0x3DF8),
     (0x107C44, 0x107C48),
     (0x3E08, 0x3E0C),
     (0x107C58, 0x107C5C),
     (0x3E1C, 0x3E20),
     (0x107C6C, 0x107C70),
     (0x3E30, 0x3E34),
     (0x107C80, 0x107C84),
     (0x3E44, 0x3E48),
     (0x107C94, 0x107C98),
     (0x3E58, 0x3E5C),
     (0x107CA8, 0x107CAC),
     (0x3E6C, 0x3E70),
     (0x107CBC, 0x107CC0),
     (0x3E80, 0x3E84),
     (0x107CD0, 0x107CD4),
     (0x3E94, 0x3E98),
     (0x107CE4, 0x107CE8),
     (0x3EA8, 0x3EAC),
     (0x107CF8, 0x107CFC),
     (0x107D14, 0x107D18),
     (0x107D28, 0x107D2C),
     (0x107D3C, 0x107D40),
     (0x107D50, 0x107D54),
     (0x107D64, 0x107D68),
     (0x1078D0, 0x1078D4),
     (0x1078E4, 0x1078E8),
     (0x1078F8, 0x1078FC),
     (0x10790C, 0x107910),
     (0x107920, 0x107924),
     (0x107934, 0x107938),
     (0x107948, 0x10794C),
     (0x10795C, 0x107960),
     (0x1079D4, 0x1079D8),
     (0x107B28, 0x107B2C),
     (0x1079EC, 0x1079F4),
     (0x107B40, 0x107B48),
     (0x107A04, 0x107A08),
     (0x107B58, 0x107B5C),
     (0x107A18, 0x107A1C),
     (0x107B6C, 0x107B70),
     (0x107A2C, 0x107A30),
     (0x107B80, 0x107B84),
     (0x107A40, 0x107A44),
     (0x107B94, 0x107B98),
     (0x107A54, 0x107A58),
     (0x107BA8, 0x107BAC),
     (0x107A68, 0x107A6C),
     (0x107BBC, 0x107BC0),
     (0x107A7C, 0x107A80),
     (0x107BD0, 0x107BD4),
     (0x107D80, 0x107D84),
     (0x107D94, 0x107D98),
     (0x107DA8, 0x107DAC),
     (0x5DEB0, 0x5DEB8),
     (0x5DECC, 0x5DED4),
     (0x5DEE8, 0x5DEF0),
     (0x5DFA0, 0x5DFA4),
     (0x5DF3C, 0x5DF44),
     (0x5DF20, 0x5DF28),
     (0x5E010, 0x5E018),
     (0x5DF04, 0x5DF0C),
     (0x5DF58, 0x5DF60),
     (0x5DF74, 0x5DF7C),
     (0x5DE98, 0x5DE9C),
     (0x5DFF4, 0x5DFFC),
     (0x5E0D4, 0x5E0DC),
     (0x5E0F0, 0x5E0F8),
     (0x5E10C, 0x5E114),
     (0x5E128, 0x5E130),
     (0x5E144, 0x5E14C),
     (0x5E160, 0x5E168),
     (0x5E17C, 0x5E184),
     (0x5E198, 0x5E1A0),
     (0x5E1B4, 0x5E1BC),
     (0x5E1D0, 0x5E1D8),
     (0x5E1EC, 0x5E1F4),
     (0x5E208, 0x5E210),
     (0x5E224, 0x5E22C),
     (0x5E240, 0x5E248),
     (0x5E25C, 0x5E264),
     (0x5E09C, 0x5E0A4),
     (0x5E0B8, 0x5E0C0),
     (0x5E064, 0x5E06C),
     (0x5E080, 0x5E088),
     (0x5E02C, 0x5E034),
     (0x5E048, 0x5E050),
     (0x5DFBC, 0x5DFC4),
     (0x5DFD8, 0x5DFE0),
     (0x5E278, 0x5E280),
     (0x5E294, 0x5E29C),
     (0x79EFC, 0x79F08),
     (0x7A254, 0x7A260),
     (0x7AC78, 0x7AC84),
     (0x7AE40, 0x7AE4C),
     (0x7B00C, 0x7B018),
     (0x7A5C4, 0x7A5D0),
     (0x6FDE4, 0x6FDE8),
     (0x7DA1C, 0x7DA30),
     (0x6FE48, 0x6FE4C),
     (0x70070, 0x70074),
     (0x70178, 0x7017C),
     (0x6FFE8, 0x6FFEC),
     (0x6FEFC, 0x6FF00),
     (0x6FF60, 0x6FF64),
     (0x6FCE4, 0x6FCE8),
     (0x70114, 0x70120),
     (0x6FA4C, 0x6FA58),)
    )

# Format: filename, old entry, size replaced.
edits = (("-hi1.NKTJ.mio", (0xF04B420, 0x4A1, 0x1000, 0), 0x4A4),
    ("Yos-.NKTJ.mio", (0xF04AF40, 0x4DF, 0x1000, 0), 0x4E0),
    ("-endo.NKTJ.mio", (0xF02C3A0, 0x2F7, 0x1000, 0), 0x2F8),
    ("Ninte-.NKTJ.mio", (0xF02C0F8, 0x2A6, 0x1000, 0), 0x2A8),
    ("-oro.NKTJ.mio", (0xF02A584, 0x2B1, 0x1000, 0), 0x2B4),
    ("Mario-.NKTJ.mio", (0xF02A2B0, 0x2D2, 0x1000, 0), 0x2D4),
    ("-igip.NKTJ.mio", (0xF02A000, 0x2AF, 0x1000, 0), 0x2B0),
    ("Lu-.NKTJ.mio", (0xF029D78, 0x287, 0x1000, 0), 0x288),
    ("-a_Air.NKTJ.mio", (0xF022D40, 0x304, 0x1000, 0), 0x304),
    ("Koop-.NKTJ.mio", (0xF0229E0, 0x360, 0x1000, 0), 0x360),
    ("64.NKTJ.mio", (0xF000B18, 0x567, 0x1000, 0), 0x568),
     )

def main():
    """This is buggered to just do the file insertion and fix the offsets."""
    import N64
    from struct import Struct
    resentry = Struct(">12L")
    helpo = Struct(">4L")
    statictbl = Struct(">20L")

    with open(r"Mario Kart 64 (U) [!].z64", 'rb') as f:
        cart = N64.Cart(f.read())
    # Get a list of the original locations.
    i = cart.ID
    name = "{}{}{}{}".format(i.frmt, i.name, i.region, '' if not i.version else str(i.version))
    ##tbl = MK64.get(name, None)
    tbl = NKTE
    if not tbl:
        raise NotImplementedError("%s is not supported!" % name)
    # Use any entry in tbl[9] to pull offset to resource table in order to calculate correction.
    for i in edits:
        with open(i[0], 'rb') as f:
            data = f.read()
        orgv = i[1][0]
        orgp = (orgv & 0xFFFFFF) + cart.getASMvalue(tbl[9][0][0], tbl[9][0][1])
        sz = i[1][1]
        sz += 3
        sz &= ~3
        l = len(data)
        padding = (4 - (l&3)) & 3
        correct = l + padding - sz
        new = (orgv, l, int.from_bytes(data[4:8], 'big') if data[0:4] == b'MIO0' else l+padding, 0)
        # walk through and update resource entries.
        # Use tbl[1] to pull the resource entries from first file.
        p = tbl[1] + cart.getASMvalue(tbl[4][0][0], tbl[4][0][1])
        while True:
            v = list(resentry.unpack_from(cart.rom, p))
            if not v[0]:
                break
            for k in range(v[4], v[5], 16):
                e = helpo.unpack_from(cart.rom, k)
                if not e[0]:
                    break
                if e == i[1]:
                    helpo.pack_into(cart.rom, k, *new)
                elif e[0] > orgv:
                    helpo.pack_into(cart.rom, k, e[0]+correct, e[1], e[2], e[3])
            # After fiddling with tables, update the entry and write it back.
            for k in range(6):
                if v[k] > orgp:
                    v[k] += correct
            resentry.pack_into(cart.rom, p, *v)
            p += 48
        # tbls 9, 10, 11, 12, 13, 14, 15 get pushed around
        for k in range(4, len(tbl)-1):
            for j in tbl[k]:
                v = cart.getASMvalue(j[0], j[1])
                if v <= orgp:
                    break
                v += correct
                cart.setASMvalue(v, j[0], j[1])
        for j in tbl[-1]:
            v = cart.getASMvalue(j[0], j[1])
            if v > orgv:
                v += correct
                cart.setASMvalue(v, j[0], j[1])
        v = list(statictbl.unpack_from(cart.rom, tbl[0]))
        for j in range(len(v)):
            if v[j] > orgv:
                v[j] += correct
        statictbl.pack_into(cart.rom, tbl[0], *v)
        # Now replace the file.
        cart.rom[orgp:orgp+i[2]] = data + bytes(padding)
    cart.calccrc(fix=True)
    with open("hacked.n64", 'wb') as f:
        f.write(cart.rom)

if __name__ == '__main__':
    main()