Python 3 version of Greg Kennedy's iNES Header Fixer

ines-header-fixer.py

#!/usr/local/bin/python
"""
Converts NES ROMs to iNES format, applying correct iNES header.
Usage: ines-fix <infile>
Supported infile formats are .nes, .pas (headerless .nes)

ROM data is recognized by CRC32 using BootGod's master XML database
  so make sure you have a local copy

This script is a python3-ified version of Greg Kennedy's original:
  https://greg-kennedy.com/wordpress/2012/05/30/ines-header-fixer/
"""

import sys
import struct
from binascii import crc32
from xml.etree import ElementTree

##### String holding location of cart db
# cart_xml = "NesCarts (2011-09-10).xml"
# uncomment next line to use Nestopia's DB instead
cart_xml = "NstDatabase.xml"

# Other required vars
i_fmt = 'p'
blob = None
found = 0
oldHeader = b""


# Parse command-line

if (len(sys.argv) != 2):
    print("Usage: " + sys.argv[0] + " <infile>")
    sys.exit(0)

# Open rom database
#print "Attempting to open cart db " + cart_xml
tree = ElementTree.parse(cart_xml)
#print "DB opened!"

# Attempt to open supplied rom file

try:
    with open(sys.argv[1], "rb") as f:
        tag = f.read(4)
        f.seek(0)
        if (tag == bytes("NES\x1A", encoding='ascii')):
            i_fmt = 'i'
            oldHeader = f.read(16)

#        print "Detected " + i_fmt + " format for input file"


        blob = f.read()

except IOError as xxx_todo_changeme1:
    (errno, strerror) = xxx_todo_changeme1.args
    print("Error opening " + sys.argv[1] + ": ")
    print("I/O error({0}): {1}".format(errno, strerror))
    sys.exit(2)

if (len(blob) > 0):
    format_crc32 = format(crc32(blob) & 0xFFFFFFFF, '08X')
    print("CRC check: " + format_crc32)

    # go look up crc32 in db
    game_list = tree.findall("game")
    for game in game_list:
        cart_list = game.findall("cartridge")
        for cart in cart_list:
            if (cart.attrib.get('crc') == format_crc32):
                found=1
                break
        if (found):
            break


    if (found == 0):
        print(sys.argv[1])
        print("*** CART NOT FOUND IN DB")
        print("----------------------------------------------------")
        sys.exit(2)

if game.attrib.get("name"):
    print("Found CRC match: " + game.attrib.get("name").encode('ascii', 'ignore'))
#ElementTree.tostring(game)

# retrieve data from game
board = cart.find("board")
mapper = int(board.attrib.get("mapper"))

prg_size = 0
prg_list = board.findall("prg")
for prg in prg_list:
    prg_size = prg_size + int(prg.attrib.get("size") [:-1])

chr_size = 0
chr_list = board.findall("chr")
for chr in chr_list:
    chr_size = chr_size + int(chr.attrib.get("size") [:-1])


battery = 0
wram_list = board.findall("wram")
for wram in wram_list:
    if (wram.attrib.get("battery") is not None):
        battery = int(wram.attrib.get("battery"))

chip_list = board.findall("chip")
for chip in chip_list:
    if (chip.attrib.get("battery") is not None):
        battery = int(chip.attrib.get("battery"))


mirror_4 = 0
mirror_v = 0
pad = board.find("pad")
if (format_crc32 == "CD50A092" or \
    format_crc32 == "EC968C51" or \
    format_crc32 == "404B2E8B"):
    mirror_4 = 1
elif (pad is not None):
# the "h" pad means "v" mirror
    mirror_v = int(pad.attrib.get("h"))


mapper_lo = int(mapper & 0x0F)
mapper_hi = int(mapper & 0xF0)

newHeader = bytes("NES\x1A", encoding='ascii') + \
              struct.pack("BBBB", \
                int( prg_size / 16 ), \
                int( chr_size / 8 ), \
                (mapper_lo << 4) + (mirror_4 << 3) + (battery << 1) + (mirror_v), \
                (mapper_hi) ) + (bytes( '\0', encoding='ascii') * 8 )

if (newHeader != oldHeader):
    print("*** HEADER UPDATED ***\noldHeader: " + oldHeader.hex().upper())
    print("newHeader: " + newHeader.hex().upper())

    # write new file
    try:
        with open(sys.argv[1], "wb") as f:
            f.write( newHeader )
            f.write( blob )


    except IOError as xxx_todo_changeme:
        (errno, strerror) = xxx_todo_changeme.args
        print("Error opening " + sys.argv[1])
        print("I/O error({0}): {1}".format(errno, strerror))
        sys.exit(2)

    print("All done.  Wrote new file " + sys.argv[1])
else:
    print("*** Header unchanged: not writing replacement file.")
print("----------------------------------------------------")