Python script to trim and untrim Nintendo 3DS game cart backups

3dstrim.py

#!/usr/bin/env python

# See https://www.3dbrew.org/wiki/NCSD for details on the file format

import io
import os
import logging


MEDIA_UNIT = 0x200
PAD_BYTE = b"\xFF"
BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE
_PAD_BUFFER = PAD_BYTE * BUFFER_SIZE


__log__ = logging.getLogger(__name__)


def is_pow_2(num):
    return (num & (num - 1)) == 0


def read_u32(fp):
    return int.from_bytes(fp.read(4), byteorder="little", signed=False)


def get_trimmed_size(fp):
    fp.seek(0x120)  # partition table start
    partition_end = 0
    for x in range(8):  # max of 8 paritions
        offset = read_u32(fp)
        size = read_u32(fp)
        partition_end = max(partition_end, offset + size)

    return partition_end * MEDIA_UNIT


def check_remaining(fp, end_offset):
    fp.seek(end_offset)
    while True:
        data = fp.read(BUFFER_SIZE)
        if not data:
            break
        ld = len(data)
        if data != (_PAD_BUFFER[:ld] if ld < BUFFER_SIZE else _PAD_BUFFER):
            return False
    return True


def humanize_bytes(num):
    num = abs(num)
    for unit in ("B", "KB", "MB", "GB"):
        if num < 1024:
            break
        num /= 1024.0
    return f"{num:.1f} {unit}"


def trim(path, *, trim=True, check=True):

    with open(path, "r+b") as fp:

        cur_size = fp.seek(0, io.SEEK_END)

        fp.seek(0x100)
        if fp.read(4) != b"NCSD":
            __log__.error("'%s': not a valid backup (missing NCSD magic)", path)
            return

        fp.seek(0x104)
        original_size = read_u32(fp) * MEDIA_UNIT
        if not is_pow_2(original_size):
            __log__.error(
                "'%s': reported original size is not a power of 2 - corrupt file?", path
            )
            return

        end_offset = get_trimmed_size(fp)

        if end_offset > cur_size:
            __log__.error(
                "'%s': smaller than the calculated trimmed size - corrupt?",
                path,
            )
            return

        if trim and end_offset == cur_size:
            __log__.info(
                "'%s': already been trimmed (current size: %s, untrimmed: %s)",
                path,
                humanize_bytes(cur_size),
                humanize_bytes(original_size),
            )
            return

        if trim:
            if check and not check_remaining(fp, end_offset):
                __log__.error("'%s': data to be trimmed isn't just 0x%s bytes", path, PAD_BYTE.hex().upper())
                return

            fp.truncate(end_offset)
            __log__.info(
                "'%s': trimmed from %s to %s (saved %s)",
                path,
                humanize_bytes(cur_size),
                humanize_bytes(end_offset),
                humanize_bytes(cur_size - end_offset),
            )
        else:
            required = original_size - cur_size
            if required < 0:
                __log__.warning("'%s': not padding file - already too big", path)
                return
            elif required == 0:
                __log__.info(
                    "'%s': file is already the correct original size (%s)",
                    path,
                    humanize_bytes(cur_size),
                )
                return

            fp.seek(cur_size)
            while required >= BUFFER_SIZE:
                fp.write(_PAD_BUFFER)
                required -= BUFFER_SIZE
            if required:
                fp.write(_PAD_BUFFER[:required])

            __log__.info(
                "'%s': padded out to %s from %s",
                path,
                humanize_bytes(original_size),
                humanize_bytes(cur_size),
            )


def main():
    import argparse

    logging.basicConfig(
        level=logging.INFO, style="{", format="{asctime} {levelname:>7}: {message}"
    )

    parser = argparse.ArgumentParser(description="[un]trim 3DS cart backups")
    parser.add_argument(
        "files", metavar="file", nargs="+", help="The cart backups to process"
    )
    parser.add_argument(
        "--no-check",
        dest="check",
        action="store_false",
        help="Verify that the data to trim is all null before trimming (slower, safer)",
    )
    parser.add_argument(
        "--untrim",
        action="store_true",
        help="The default is to trim files, use this to untrim instead",
    )

    args = parser.parse_args()
    for path in args.files:
        trim(path, trim=not args.untrim, check=args.check)


if __name__ == "__main__":
    main()