organizer2.py

#!/usr/bin/env python

VERBOSE = False

import os
import subprocess
import shutil
import errno
import hashlib

# CNMT parser shamelessly stolen from https://github.com/Rikikooo/CDNSP/blob/master/lib/CNMT.py
from struct import pack as pk, unpack as upk
def read_at(fp, off, len):
    fp.seek(off)
    return fp.read(len)

def read_u8(fp, off):
    return upk('<B', read_at(fp, off, 1))[0]

def read_u16(fp, off):
    return upk('<H', read_at(fp, off, 2))[0]

def read_u32(fp, off):
    return upk('<I', read_at(fp, off, 4))[0]

def read_u48(f, off):
    s = upk('<HI', read_at(f, off, 6))
    return s[1] << 16 | s[0]

def read_u64(fp, off):
    return upk('<Q', read_at(fp, off, 8))[0]

class CNMT:
    title_types = {
        0x1: 'SystemProgram',
        0x2: 'SystemData',
        0x3: 'SystemUpdate',
        0x4: 'BootImagePackage',
        0x5: 'BootImagePackageSafe',
        0x80:'Application',
        0x81:'Patch',
        0x82:'AddOnContent',
        0x83:'Delta'
    }
    nca_types  = {
        0:'Meta',
        1:'Program',
        2:'Data',
        3:'Control',
        4:'HtmlDocument',
        5:'LegalInformation',
        6:'DeltaFragment'
    }
    def __init__(self, fp):
        self.f = fp
        self._parse()

    def _parse(self):
        self.title_type = self.title_types[read_u8(self.f, 0xC)]
        self.tid        = read_u64(self.f, 0x0)
        self.ver        = read_u32(self.f, 0x8)
        self.sysver     = read_u64(self.f, 0x28)
        self.dlsysver   = read_u64(self.f, 0x18)

        self.data = {}
        if self.title_type == 'SystemUpdate':
            self.titles_nb = 0
            entries_nb = read_u16(self.f, 0x12)
            for n in range(entries_nb):
                self.titles_nb += 1
                offset = 0x20 + 0x10 * n
                tid        = read_u64(self.f, offset)
                ver        = read_u32(self.f, offset + 0x8)
                title_type = self.title_types[read_u8(self.f, offset + 0xC)]
                self.data[tid] = {
                    'Version': ver,
                    'Type':    title_type
                }
        else:
            self.files_nb   = 0
            self.title_size = 0
            for nca_type in list(self.nca_types.values()):
                self.data[nca_type] = {}
            table_offset = read_u16(self.f,0xE)
            entries_nb = read_u16(self.f, 0x10)
            for n in range(entries_nb):
                offset = 0x20 + table_offset + 0x38 * n
                hash = read_at(self.f, offset, 0x20)
                ncaid = int.from_bytes(read_at(self.f, offset + 0x20, 0x10), byteorder='big')
                size = read_u48(self.f, offset + 0x30)
                nca_type = self.nca_types[read_u16(self.f, offset+0x36)]
                self.data[nca_type][ncaid] = {
                    'Size': size,
                    'Hash': hash
                }
                self.files_nb   += 1
                self.title_size += size

def mkdirp(path):
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise


def symlink_force(target, link_name):
    try:
        os.symlink(target, link_name)
    except OSError as e:
        if e.errno == errno.EEXIST:
            os.remove(link_name)
            os.symlink(target, link_name)
        else:
            raise e

versions = {
    450:       "1.0.0",
    65796:     "2.0.0",
    131162:    "2.1.0",
    196628:    "2.2.0",
    262164:    "2.3.0",
    201327002: "3.0.0",
    201392178: "3.0.1",
    201457684: "3.0.2",
    268435656: "4.0.0",
    268501002: "4.0.1",
    269484082: "4.1.0",
    335544750: "5.0.0",
    335609886: "5.0.1",
    335675432: "5.0.2",
    336592976: "5.1.0",
    402653494: "6.0.0-4",
    402653514: "6.0.0",
    402653544: "6.0.0",
    402718730: "6.0.1",
    403701850: "6.1.0",
    404750376: "6.2.0",
    469762248: "7.0.0",
    469827614: "7.0.1",
    536871502: "8.0.0"
}

def get_ncaid(filename):
    BLOCKSIZE = 65536
    hasher = hashlib.sha256()
    with open(filename, 'rb') as afile:
        buf = afile.read(BLOCKSIZE)
        while len(buf) > 0:
            hasher.update(buf)
            buf = afile.read(BLOCKSIZE)
    return hasher.hexdigest()[:32]

def print_verbose(*args, **kwargs):
    if VERBOSE:
        print(*args, **kwargs)

for version in os.listdir("."):
    # Ignore files (only treat directories)
    if os.path.isfile(version):
        continue

    if version == "updaters":
        continue

    # Rename CNMTs to make them easier to find. Also, if nca is a folder,
    # get its content. Also, get the hash, and give it the proper ncaid name
    print(f"===== Handling {version} =====")

    HACTOOL_PROGRAM = "hactool"
    if os.path.isfile(version + "/dev"):
        HACTOOL_PROGRAM += " --dev"

    print(f"# Normalizing the by-hash folder")
    for nca in os.listdir(version + "/by-hash"):
        ncaFull = version + "/by-hash/" + nca
        # Fix "folder-as-file" files when dumped from Switch NAND
        if os.path.isdir(ncaFull):
            print_verbose(f"{ncaFull}/00 -> {ncaFull}")
            os.rename(ncaFull, ncaFull + "_folder")
            os.rename(ncaFull + "_folder/00", ncaFull)
            os.rmdir(ncaFull + "_folder")
        # Ensure the NCAID is correct (It's wrong when dumped from the
        # Placeholder folder on a Switch NAND
        ncaid = get_ncaid(ncaFull)
        newName = version + "/by-hash/" + ncaid + "." + ".".join(os.path.basename(ncaFull).split(".")[1:])
        print_verbose(f"{ncaFull} -> {newName}")
        os.rename(ncaFull, newName)
        ncaFull = newName

        # Ensure meta files have .cnmt.nca extension
        process = subprocess.Popen(["sh", "-c", HACTOOL_PROGRAM + " '" + ncaFull + "' | grep 'Content Type:' | awk '{print $3}'"], stdout=subprocess.PIPE)
        contentType = process.communicate()[0].split(b"\n")[0].decode('utf-8')
        if contentType == "Meta" and not nca.endswith(".cnmt.nca"):
            print_verbose(ncaFull + " -> " + ".".join(ncaFull.split(".")[:-1]) + ".cnmt.nca")
            shutil.move(ncaFull, ".".join(ncaFull.split(".")[:-1]) + ".cnmt.nca")

    print("# Setup by-title and by-name")
    for nca in os.listdir(version + "/by-hash"):
        ncaFull = version + "/by-hash/" + nca
        process = subprocess.Popen(["sh", "-c", HACTOOL_PROGRAM + " '" + ncaFull + "' | grep 'Title ID:' | awk '{print $3}'"], stdout=subprocess.PIPE)
        titleId = process.communicate()[0].split(b"\n")[0].decode('utf-8')
        process = subprocess.Popen(["sh", "-c", HACTOOL_PROGRAM + " '" + ncaFull + "' | grep 'Content Type:' | awk '{print $3}'"], stdout=subprocess.PIPE)
        contentType = process.communicate()[0].split(b"\n")[0].decode('utf-8')

        mkdirp(version + "/by-title/" + titleId)

        print_verbose(version + "/by-title/" + titleId + "/" + contentType + ".nca -> " + "../../by-hash/" + nca)
        symlink_force("../../by-hash/" + nca, version + "/by-title/" + titleId + "/" + contentType + ".nca")

        process = subprocess.Popen(["sh", "-c", HACTOOL_PROGRAM + " '" + ncaFull + "' | grep 'Title Name:' | awk '{print $3}'"], stdout=subprocess.PIPE)
        titleName = process.communicate()[0].split(b"\n")[0].decode('utf-8')
        if titleName != "":
            mkdirp(version + "/by-name")
            print_verbose(version + "/by-name/" + titleName + " -> " + "../by-title/" + titleId)
            symlink_force("../by-title/" + titleId, version + "/by-name/" + titleName)

    print("# Extracting all ExeFS and RomFS")
    mkdirp(version + "/by-name")
    for titlename in os.listdir(version + "/by-name"):
        ncaParent = version + "/by-name/" + titlename
        ncaFull = ncaParent + "/Program.nca"
        print_verbose(f"Extracting {ncaFull}")
        process = subprocess.Popen(["sh", "-c", f"{HACTOOL_PROGRAM} '{ncaFull}' -tnca --exefsdir '{ncaParent}/exefs' --romfsdir '{ncaParent}/romfs'"], stdout=subprocess.DEVNULL)

    print("# Extracting firmware title")
    ncaParent = version + "/by-title/0100000000000819"
    ncaFull = ncaParent + "/Data.nca"
    process = subprocess.Popen(["sh", "-c", f"{HACTOOL_PROGRAM} '{ncaFull}' -tnca --romfsdir '{ncaParent}/romfs'"], stdout=subprocess.DEVNULL)
    process.wait()
    subprocess.Popen(["sh", "-c", f"{HACTOOL_PROGRAM} '{ncaParent}/romfs/nx/package1' -tpk11 --package1dir '{ncaParent}/package1'"], stdout=subprocess.DEVNULL)
    subprocess.Popen(["sh", "-c", f"{HACTOOL_PROGRAM} '{ncaParent}/romfs/nx/package2' -tpk21 --package2dir '{ncaParent}/package2' --ini1dir '{ncaParent}/ini1'"], stdout=subprocess.DEVNULL)

    print("# Verifying the dump is complete")
    process = subprocess.run(["sh", "-c", f"{HACTOOL_PROGRAM} '{version}/by-title/0100000000000816/Meta.nca' -tnca --section0dir '{version}/by-title/0100000000000816/section0'"], stdout=subprocess.DEVNULL)
    cnmt = CNMT(open(version + "/by-title/0100000000000816/section0/SystemUpdate_0100000000000816.cnmt", "rb"))
    # Verify the update
    fails = []
    for key, value in cnmt.data.items():
        hexkey = '0' + hex(key)[2:]
        ncaParent = f"{version}/by-title/{hexkey}"
        ncaFull = f"{ncaParent}/Meta.nca"
        print_verbose(f"Extracting {ncaFull}")
        process = subprocess.run(["sh", "-c", f"{HACTOOL_PROGRAM} '{ncaFull}' -tnca --section0dir '{ncaParent}/meta'"], stdout=subprocess.DEVNULL)
        try:
            programcnmt = CNMT(open(f"{ncaParent}/meta/{value['Type']}_{hexkey}.cnmt", "rb"))
            if programcnmt.ver != value["Version"]:
                fails.append(f"Title {hexkey} has an unmatching Version")

            # Ensure each data type matches the NCA ID
            for ty, v in programcnmt.data.items():
                for ncaid, v in v.items():
                    ncaid = "{0:0{1}x}".format(ncaid, 32)
                    if not os.path.exists(f"{ncaParent}/{ty}.nca") and ty == "Data":
                        ty = "Unknown"
                    realnca = os.path.realpath(f"{ncaParent}/{ty}.nca")
                    if not realnca.endswith(ncaid + ".nca"):
                        fails.append(f"Title {hexkey} has wrong NCAID (should be {ncaid}.nca, is {realnca})")
        except FileNotFoundError:
            fails.append(f"Title {hexkey} is missing. It should be at version {value['Version']}")

    # TODO: Get "real" SystemVersion
    if cnmt.ver in versions:
        updatename = "update " + versions[cnmt.ver]
    else:
        updatename = "unknown update " + str(cnmt.ver)

    if len(fails) != 0:
        print(f"Incomplete {updatename}:")
        for fail in fails:
            print(fail)
    else:
        print(f"Successfully organized data for {updatename}")