ntrbootbanner.py - Extracts/Injects a banner from/to an ntrboot backup

ntrbootbanner.py

#!/usr/bin/env python3

"""
This is free and unencumbered software released into the public domain.
Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.
In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.
For more information, please refer to <http://unlicense.org/>
"""

from argparse import ArgumentParser, FileType
from io import SEEK_CUR, SEEK_END
from struct import unpack


BANNER_SIZES = {
	0x0001: 0x840,
	0x0002: 0x940,
	0x0003: 0xA40,
	0x0103: 0x23C0
}


def crc16(data):
	crc = 0xFFFF
	for byte in bytearray(data):
		crc ^= byte
		for __ in range(8):
			crc = (crc >> 1) ^ (0xA001 if (crc & 0x0001) else 0)

	return crc


def checkBackup(backup, force=False):
	# NDS file starts at 0x2000, banner offset is 0x68 within that
	print("Info: Searching for banner...")

	banners = []

	backup.seek(0, SEEK_END)
	fileSize = backup.tell()
	backup.seek(0)

	print("Info: ^c to stop search early")

	try:
		while backup.tell() < fileSize:
			offset = backup.tell()

			if checkBanner(backup, True, force):
				print("\rInfo: Banner found at 0x%X" % offset)
				banners.append(offset)

			backup.seek(0x200, SEEK_CUR)

			print("\r0x%07X / 0x%07X" % (offset, fileSize), end="")
	except KeyboardInterrupt:
		pass

	print("\r", end="")

	if len(banners) > 0:
		return banners
	else:
		print("Error: No valid banner found in ntrboot backup file")
		backup.seek(0)
		return None


def checkBanner(banner, fromBackup=False, force=False):
	start = banner.tell()
	data = banner.read(0x23C0)
	if len(data) < 0x840:
		return False

	ver, = unpack("<H", data[:2])
	size = banner.tell()
	banner.seek(start)

	if not fromBackup:  # For checking the banner in the backup
		if ver not in BANNER_SIZES or size != BANNER_SIZES[ver]:
			print("Error: Incorrect banner size")
			return False

	if crc16(data[0x20:0x840]) != unpack("<H", data[2:4])[0]:
		if not fromBackup:
			print("Error: Incorrect banner version 1 CRC")
		return False

	if ver & 2 and crc16(data[0x20:0x940]) != unpack("<H", data[4:6])[0]:
		print("Warn: Incorrect banner version 2 CRC")
		if not force:
			return False

	if (ver & 3) == 3 and crc16(data[0x20:0xA40]) != unpack("<H", data[6:8])[0]:
		print("Warn: Incorrect banner version 3 CRC")
		if not force:
			return False

	if ver & 0x0100 and crc16(data[0x1240:0x23C0]) != unpack("<H", data[8:10])[0]:
		print("Warn: Incorrect banner DSi icon CRC")
		if not force:
			return False

	return True


def selectBanner(banners, backup):
	titles = []

	for banner in banners:
		backup.seek(banner + 0x340)
		titles.append(backup.read(0x100).decode("utf-16le"))

	selection = None
	while not selection:
		for i in range(len(banners)):
			print("%d: %s (0x%X)" % (i + 1, titles[i][:titles[i].find("\n")], banners[i]))

		try:
			selection = int(input("> "))
		except ValueError:
			pass

		if not selection or selection < 1 or selection > len(banners):
			print("Please number between 1 and %d." % len(banners))
			selection = None

	return banners[selection - 1]


def extractBanner(backup, banner):
	# Backup should be seeked to the banner
	offset = backup.tell()

	# Read banner version, to know the size of it
	bannerVer, = unpack("<H", backup.read(2))
	backup.seek(offset)

	# Extract banner
	banner.write(backup.read(BANNER_SIZES[bannerVer]))
	print("Info: Banner extracted successfully!")

	return True


def injectBanner(backup, banner):
	# Backup should be seeked to the banner
	offset = backup.tell()

	# Read existing banner version
	oldBannerVer, = unpack("<H", backup.read(2))
	backup.seek(offset)

	# Check that new banner size matches
	newBannerVer, = unpack("<H", banner.read(2))
	banner.seek(0)
	if newBannerVer != oldBannerVer:
		print("Error: New banner version (0x%04X) does not match old one (0x%04X)" % (newBannerVer, oldBannerVer))
		return False

	# Inject banner
	backup.write(banner.read())
	print("Info: Banner injected successfully!")

	return True


def main():
	parser = ArgumentParser(description="Extracts/Injects a banner from/to an ntrboot backup")
	parser.add_argument("-x", "--extract", metavar="banner.bin", type=FileType("wb"), help="banner file to extract")
	parser.add_argument("-i", "--inject", metavar="banner.bin", type=FileType("rb"), help="banner file to inject")
	parser.add_argument("-f", "--force", action="store_true", help="allow ignoring some checksums")
	parser.add_argument("backup", metavar="backup.bin", type=FileType("rb+"), help="flashcard backup to extract/inject from/to")

	args = parser.parse_args()

	banners = checkBackup(args.backup)
	if not banners:
		print("Error: %s is not a valid ntrboot backup" % args.backup.name)
		exit(1)

	args.backup.seek(selectBanner(banners, args.backup))

	if args.extract and not args.inject:
		extractBanner(args.backup, args.extract)
	elif args.inject and not args.extract:
		if not checkBanner(args.inject, force=args.force):
			print("Error: %s is not a valid banner file" % args.inject.name)
			exit(2)

		injectBanner(args.backup, args.inject)
	else:
		print("Error: You must do *one* of -x or -i")


if __name__ == "__main__":
	main()