inject.py

# Copyright 2010 Dolphin Emulator Project
# Licensed under GPLv2+

# Copyright 2007,2008  Segher Boessenkool  <segher@kernel.crashing.org>
# Licensed under the terms of the GNU GPL, version 2
# http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt

# Copyright 2021 Zephiles, Seeky

from hashlib import sha1
from sys import argv

from Crypto.Cipher import AES

from binarybuffer import BinaryBuffer
import commondefs
import ec

# Based off of SaveFileHack.py by Zephiles
def generateRelLoader(VersionString):
    # Set version-specific values
    if (VersionString == "eu0") or (VersionString == "eu1"):
        InitAsmFunctionPointer = 0x80526294
        TextBufferPointerOffset = 0x1B0
        TextBufferPointer = 0x805256FC
        BinVersion = "EU"
        ItemId = 0x6E59
    elif VersionString == "jp0":
        InitAsmFunctionPointer = 0x804B8594
        TextBufferPointerOffset = 0x17C
        TextBufferPointer = 0x804B79C8
        BinVersion = "JP_0"
        ItemId = 0x6D0A
    elif VersionString == "jp1":
        InitAsmFunctionPointer = 0x804B9B94
        TextBufferPointerOffset = 0x174
        TextBufferPointer = 0x804B8FC0
        BinVersion = "JP_1"
        ItemId = 0x6D24
    elif VersionString == "kr0":
        InitAsmFunctionPointer = 0x8055DBF4
        TextBufferPointerOffset = 0x170
        TextBufferPointer = 0x8055D01C
        BinVersion = "KR"
        ItemId = 0x70D9
    elif VersionString == "us0":
        InitAsmFunctionPointer = 0x804E3294
        TextBufferPointerOffset = 0x1A4
        TextBufferPointer = 0x804E26F0
        BinVersion = "US_0"
        ItemId = 0x6D08
    elif VersionString == "us1":
        InitAsmFunctionPointer = 0x804E4B14
        TextBufferPointerOffset = 0x1AC
        TextBufferPointer = 0x804E3F78
        BinVersion = "US_1"
        ItemId = 0x6D26
    elif VersionString == "us2":
        InitAsmFunctionPointer = 0x804E4C94
        TextBufferPointerOffset = 0x174
        TextBufferPointer = 0x804E40C0
        BinVersion = "US_2"
        ItemId = 0x6D24

    # Create a cleared file
    ret = BinaryBuffer(bytearray(0x25b8))

    # Set the default efb width and height, as they can apparently effect being able to open the item menu
    ret.writeHalfword(0x20, 608)
    ret.writeHalfword(0x22, 480)

    # Write the new file name
    ret.writeStr(0x28, "REL Loader")

    # Write the map name
    ret.writeStr(0x4c, "dos_01")
    
    # Set Mario's level to 1, to prevent leveling up immediately
    ret.writeWord(0x1b14, 1)

    # Set the flip timer to 10 to prevent counting up immediately
    ret.writeWord(0x1b24, 10)

    # Write the item id
    ret.writeHalfword(0x1b70, ItemId)

    # Write the pointer to text buffer
    ret.writeWord(TextBufferPointerOffset, TextBufferPointer)

    # Write the text buffer
    ret.writeat(TextBufferPointerOffset + 4, [0x33] * 0x94)

    # Write the pointer to the init asm function
    ret.writeWord(TextBufferPointerOffset + 0x98, InitAsmFunctionPointer)

    # Write the init asm function
    # The init function is the same for all versions except for Korean
    if VersionString == "kr0":
        InitAsmFuncBinName = "Init_KR"
    else:
        InitAsmFuncBinName = "Init_Main"

    InitAsmFuncOffset = 0xD4C
    with open("bin/" + InitAsmFuncBinName + ".bin", "rb") as g:
        ret.writeat(InitAsmFuncOffset, g.read())

    # Write the main asm function
    with open("bin/Main_" + BinVersion + ".bin", "rb") as g:
        ret.writeat(InitAsmFuncOffset + 0x24, g.read())

    # Get the sum of the bytes for the data field
    DataFieldSum = sum(ret.data[:0x25b0]) + (0xff * 4)

    # Set the checksum of the bytes for the data field
    ret.writeWord(0x25b0, DataFieldSum)

    # Set the inverted checksum of the bytes for the data field
    ret.writeWord(0x25b4, ~DataFieldSum & 0xffffffff)

    return ret.data

def makeBlankEccCert(issuer, name, publicKey, keyId):
    cert = BinaryBuffer(bytearray(0x180))
    cert.writeWord(commondefs.CERT_SIG_TYPE_OFFSET, commondefs.CERT_SIG_TYPE_ECC) # sig type ECC
    cert.writeStr(commondefs.CERT_ISSUER_OFFSET, issuer)
    cert.writeWord(commondefs.CERT_KEY_TYPE_OFFSET, commondefs.CERT_KEY_TYPE_ECC) # key type ECC
    cert.writeStr(commondefs.CERT_NAME_OFFSET, name)
    cert.writeWord(commondefs.CERT_KEY_ID_OFFSET, keyId) # key id
    cert.writeat(commondefs.CERT_PUBLIC_KEY_OFFSET, publicKey) # public key
    return cert

class SaveFile(BinaryBuffer):
    def __init__(self, name, data, perm, attr):
        self.name = name
        self.data = bytearray(data)
        self.perm = perm
        self.attr = attr

def unpackDataBin(dataBinPath):
    #########
    # Setup #
    #########

    dataBin = open(dataBinPath, 'rb')
    files = []

    ###############
    # Main Header #
    ###############

    header = dataBin.read(commondefs.HEADER_SIZE)
    
    #############
    # Bk Header #
    #############

    _bk = dataBin.read(commondefs.BK_SIZE)
    bk = BinaryBuffer(_bk)

    fileCount = bk.readWord(commondefs.BK_FILE_COUNT_OFFSET)
    _titleId = bk.readat(commondefs.BK_TITLE_ID_OFFSET, 8)
    titleId = ""
    for c in _titleId:
        titleId += f"{c:02x}"

    #########
    # Files #
    #########

    for i in range(0, fileCount):
        curHeader = BinaryBuffer(dataBin.read(commondefs.FILE_HEADER_SIZE))

        # directories currently aren't supported, since SPM doesn't use them
        assert curHeader.readByte(commondefs.FILE_TYPE_OFFSET) == commondefs.FILE_TYPE_FILE, "Directories are not supported"

        curSize = curHeader.readWord(commondefs.FILE_SIZE_OFFSET)
        trueSize = (curSize + 63) & ~63

        curName = curHeader.readStr(commondefs.FILE_NAME_OFFSET)
        curIV = curHeader.readat(commondefs.FILE_IV_OFFSET, commondefs.IV_SIZE)        
        
        cipher = AES.new(commondefs.SD_KEY, AES.MODE_CBC, iv=curIV)
        curData = cipher.decrypt(dataBin.read(trueSize))

        files.append(SaveFile(
            curName,
            curData[:curSize],
            curHeader.readByte(commondefs.FILE_PERM_OFFSET),
            curHeader.readByte(commondefs.FILE_ATTR_OFFSET)
        ))

    dataBin.close()

    return header, titleId, files

def repackDataBin(outPath, header, titleId, files):
    #########
    # Setup #
    #########

    outBin = open(outPath, 'wb')
    filesHashWork = sha1()

    ###############
    # Main Header #
    ###############

    outBin.write(header)

    #################
    # Prepare Files #
    #################

    fileSize = 0
    for file in files:
        fileSize += ((len(file.data) + 63) & ~63) + commondefs.FILE_HEADER_SIZE

    #############
    # Bk Header #
    #############

    # Unknown fields and mac address left as 0
    bk = BinaryBuffer(bytearray(commondefs.BK_SIZE))
    bk.writeWord(commondefs.BK_SIZE_OFFSET, commondefs.BK_TRUE_SIZE)
    bk.writeWord(commondefs.BK_MAGIC_OFFSET, commondefs.BK_MAGIC)
    bk.writeWord(commondefs.BK_NG_ID_OFFSET, commondefs.DEFAULT_DEVICE_ID)
    bk.writeWord(commondefs.BK_FILE_COUNT_OFFSET, len(files))
    bk.writeWord(commondefs.BK_FILE_SIZE_OFFSET, fileSize)
    bk.writeWord(commondefs.BK_TOTAL_SIZE_OFFSET, commondefs.BK_SIZE + fileSize + commondefs.TOTAL_CERTS_SIZE)
    for i in range(0, len(titleId), 2): # convert from hex string to bytes
        h = titleId[i] + titleId[i+1]
        bk.writeByte(commondefs.BK_TITLE_ID_OFFSET+(i//2), int(h, 16))
    outBin.write(bk.data)
    filesHashWork.update(bk.data)

    #############
    # Add Files #
    #############

    for i, file in enumerate(files):
        curHeader = BinaryBuffer(bytearray(commondefs.FILE_HEADER_SIZE))
        curHeader.writeWord(commondefs.FILE_MAGIC_OFFSET, commondefs.FILE_MAGIC)
        curHeader.writeWord(commondefs.FILE_SIZE_OFFSET, len(file.data))
        curHeader.writeByte(commondefs.FILE_PERM_OFFSET, files[i].perm)
        curHeader.writeByte(commondefs.FILE_ATTR_OFFSET, files[i].attr)
        # directories currently aren't supported, since SPM doesn't use them
        curHeader.writeByte(commondefs.FILE_TYPE_OFFSET, commondefs.FILE_TYPE_FILE)
        curHeader.writeStr(commondefs.FILE_NAME_OFFSET, files[i].name)
        curHeader.writeat(commondefs.FILE_IV_OFFSET, commondefs.SD_IV)

        outBin.write(curHeader.data)
        filesHashWork.update(curHeader.data)

        trueSize = (len(file.data) + 63) & ~63
        while len(file.data) < trueSize:
            file.data.append(0)

        cipher = AES.new(commondefs.SD_KEY, AES.MODE_CBC, iv=commondefs.SD_IV)
        enc = cipher.encrypt(file.data)
        outBin.write(enc)
        filesHashWork.update(enc)

    ###########
    # Signing #
    ###########

    apPriv = bytearray(30)
    apPriv[0x1d] = 1
    
    # # apPub = ec.privToPub(apPriv)
    # apPub = bytearray([ # generating the public key can take a while, so it's pre-calculated
    #     0x00, 0xfa, 0xc9, 0xdf, 0xcb, 0xac, 0x83, 0x13,
    #     0xbb, 0x21, 0x39, 0xf1, 0xbb, 0x75, 0x5f, 0xef,
    #     0x65, 0xbc, 0x39, 0x1f, 0x8b, 0x36, 0xf8, 0xf8,
    #     0xeb, 0x73, 0x71, 0xfd, 0x55, 0x8b, 0x01, 0x00,
    #     0x6a, 0x08, 0xa4, 0x19, 0x03, 0x35, 0x06, 0x78,
    #     0xe5, 0x85, 0x28, 0xbe, 0xbf, 0x8a, 0x0b, 0xef,
    #     0xf8, 0x67, 0xa7, 0xca, 0x36, 0x71, 0x6f, 0x7e,
    #     0x01, 0xf8, 0x10, 0x52
    # ])
    # apCert = makeBlankEccCert("Root-CA00000001-MS00000002-NG0403ac68", "AP0000000100000002", apPub, 0)
    # apHashWork = sha1()
    # apHashWork.update(apCert.data[0x80:])
    # apHash = apHashWork.digest()
    # apCert.writeat(commondefs.CERT_SIG_OFFSET, ec.sign(commondefs.DEFAULT_PRIVATE_KEY, apHash))

    apCert = BinaryBuffer([ # That code can take a while, so the cert is pre-calculated
        0x00, 0x01, 0x00, 0x02, 0x00, 0x31, 0x6d, 0xe7, 0xe7, 0xca, 0x8d, 0xdf, 0xfd, 0xf1, 0x3f, 0xe0,
        0x35, 0x54, 0xf3, 0x9e, 0x35, 0x92, 0x5d, 0x3e, 0xb3, 0x40, 0x85, 0x06, 0x11, 0x13, 0xdd, 0x37,
        0x11, 0xe5, 0x00, 0xf9, 0x60, 0x22, 0xe0, 0xde, 0x92, 0x09, 0x92, 0xf4, 0xa3, 0x9b, 0xe3, 0xc5,
        0x6f, 0xf7, 0x9c, 0xdd, 0x0b, 0x1a, 0x3e, 0x3b, 0x2a, 0x8f, 0x4f, 0xea, 0xc7, 0xe4, 0x0d, 0xa7,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x52, 0x6f, 0x6f, 0x74, 0x2d, 0x43, 0x41, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x31, 0x2d,
        0x4d, 0x53, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x32, 0x2d, 0x4e, 0x47, 0x30, 0x34, 0x30,
        0x33, 0x61, 0x63, 0x36, 0x38, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x02, 0x41, 0x50, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x30, 0x31, 0x30, 0x30,
        0x30, 0x30, 0x30, 0x30, 0x30, 0x32, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xfa, 0xc9, 0xdf, 0xcb, 0xac, 0x83, 0x13,
        0xbb, 0x21, 0x39, 0xf1, 0xbb, 0x75, 0x5f, 0xef, 0x65, 0xbc, 0x39, 0x1f, 0x8b, 0x36, 0xf8, 0xf8,
        0xeb, 0x73, 0x71, 0xfd, 0x55, 0x8b, 0x01, 0x00, 0x6a, 0x08, 0xa4, 0x19, 0x03, 0x35, 0x06, 0x78,
        0xe5, 0x85, 0x28, 0xbe, 0xbf, 0x8a, 0x0b, 0xef, 0xf8, 0x67, 0xa7, 0xca, 0x36, 0x71, 0x6f, 0x7e,
        0x01, 0xf8, 0x10, 0x52, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
    ])

    filesHash = filesHashWork.digest()
    sigHashWork = sha1()
    sigHashWork.update(filesHash)
    sigHash = sigHashWork.digest()
    apSig = ec.sign(apPriv, sigHash)

    deviceCert = makeBlankEccCert("Root-CA00000001-MS00000002", "NG0403ac68", commondefs.DEFAULT_PUBLIC_KEY, commondefs.DEFAULT_KEY_ID)
    deviceCert.writeat(commondefs.CERT_SIG_OFFSET, commondefs.DEFAULT_SIGNATURE)

    outBin.write(apSig)
    outBin.write(int.to_bytes(commondefs.SIGNATURE_END_MAGIC, 4, 'big'))
    outBin.write(deviceCert.data)  
    outBin.write(apCert.data)

    outBin.close()

def getNameEncoding(version):
    if version[:2] in ["eu", "us", "kr"]:
        return 'ansi'
    else: # jp
        return 'shift-jis'

if __name__ == "__main__":
    ##########
    # Unpack #
    ##########

    # Get input data.bin path
    if len(argv) > 1:
        inBinName = argv[1]
    else:
        inBinName = input("Enter path to input data.bin: ")

    # Unpack input data.bin
    print("Unpacking data.bin...")
    header, titleId, files = unpackDataBin(inBinName)

    #################
    # Inject loader #
    #################

    # Get game version
    if len(argv) > 2:
        gameVer = argv[2].lower()
    else:
        gameVer = ""
    while not gameVer in ["eu0", "eu1", "jp0", "jp1", "us0", "us1", "us2", "kr0"]:
        gameVer = input("Enter game version (eu0/eu1/jp0/jp1/us0/us1/us2/kr0): ").lower()

    # Scan for wiimario save files
    encoding = getNameEncoding(gameVer)
    fileNames = [None] * 4
    filePositions = [None] * 4
    for i, file in enumerate(files):
        if not file.name.startswith("wiimario"):
            continue
        idx = int(file.name[len("wiimario"):])
        if not 0 <= idx <= 3:
            continue
        name = file.readStr(0x8 + 0x20, encoding)
        try:
            name = file.readStr(0x8 + 0x20, encoding)
            if len(name) == 0:
                name = "<Empty>"
        except:
            name = "<Unable to read name>"
        fileNames[idx] = name
        filePositions[idx] = i
    
    
    # Get save file id to replace
    if len(argv) > 3:
        try:
            saveNumber = int(argv[3])
        except:
            saveNumber = 0
    else:
        saveNumber = 0
    
    if not 1 <= saveNumber <= 4:
        print("Saves found: ")
        for i, info in enumerate(fileNames):
            print(f"    {i+1}: {info}")
        saveNumber = 0
        while not 1 <= saveNumber <= 4:
            inp = input("Enter file id to replace with rel loader (1-4): ")
            try:
                saveNumber = int(inp)
            except:
                pass
    
    # Inject rel loader
    print("Injecting loader...")
    files[filePositions[saveNumber - 1]].data = generateRelLoader(gameVer)

    ##############
    # Inject rel #
    ##############

    # Remove existing rel if present
    for i, file in enumerate(files):
        if file.name == "pcrel.bin":
            del files[i]
            break

    # Get input rel path
    if len(argv) > 4:
        inRelName = argv[4]
    else:
        inRelName = input("Enter path to input rel: ")

    print("Injecting rel...")
    with open(inRelName, 'rb') as inRel:
        relData = inRel.read()
        files.append(SaveFile(
            "pcrel.bin", int.to_bytes(len(relData), 4, 'big') + bytes(0x1c) + relData, 0x3f, 0
        ))

    ##########
    # Repack #
    ##########

    # Get output data.bin path
    if len(argv) > 5:
        outBinName = argv[5]
    else:
        outBinName = input("Enter path to output data.bin: ")
    
    # Pack output data.bin
    print("Repacking data.bin... (this can take a while)")
    repackDataBin(outBinName, header, titleId, files)

    print("Done!")