From 5c33550a0c649eb507c9ba50e5d066a701ed709b Mon Sep 17 00:00:00 2001
From: Connor Olding <cloningdonor@gmail.com>
Date: Sun, 1 Mar 2015 19:52:12 -0800
Subject: [PATCH] knock off TODOs

---
 n64.py     | 17 ++++++++++++++++-
 z64dump.py | 28 +++++++++++++++++++++-------
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/n64.py b/n64.py
index 244b399..dab3d59 100644
--- a/n64.py
+++ b/n64.py
@@ -1,5 +1,9 @@
 # Based on uCON64's N64 checksum algorithm by Andreas Sterbenz
 
+from zlib import crc32
+
+MAX32 = 0xFFFFFFFF
+
 crc_seeds = {
     6101: 0xF8CA4DDC,
     6102: 0xF8CA4DDC,
@@ -8,7 +12,13 @@ crc_seeds = {
     6106: 0x1FEA617A,
 }
 
-MAX32 = 0xFFFFFFFF
+bootcode_crcs = {
+    0x6170A4A1: 6101,
+    0x90BB6CB5: 6102,
+    0x0B050EE0: 6103,
+    0x98BC2C86: 6105,
+    0xACC8580A: 6106,
+}
 
 def ROL(i, b):
     return ((i << b) | (i >> (32 - b))) & MAX32
@@ -67,3 +77,8 @@ def crc(f, bootcode=6105):
         crc1 = t6 ^ t4 ^ t3
         crc2 = t5 ^ t2 ^ t1
     return crc1 & MAX32, crc2 & MAX32
+
+def bootcode_version(f):
+    f.seek(0x40)
+    return bootcode_crcs[crc32(f.read(0x1000 - 0x40)) & MAX32]
+
diff --git a/z64dump.py b/z64dump.py
index 8c6d7a0..31b9cb9 100755
--- a/z64dump.py
+++ b/z64dump.py
@@ -3,7 +3,8 @@
 
 import os, os.path
 import sys
-import struct
+import io
+import struct, array
 import hashlib
 
 import n64
@@ -105,15 +106,28 @@ def z_dump(f):
     while z_dump_file(f, '{:05} '.format(i)):
         i += 1
 
+def swap_order(f, size='H'):
+    f.seek(0)
+    a = array.array(size, f.read())
+    a.byteswap()
+    f.seek(0)
+    f.write(a.tobytes())
+
 def dump_rom(fn):
     with open(fn, 'rb') as f:
         data = f.read()
 
-        if data[:4] != b'\x80\x37\x12\x40':
-            # TODO: check if it's a .n64 (2 byte swap) and convert
+    with io.BytesIO(data) as f:
+        start = f.read(4)
+        if start == b'\x37\x80\x40\x12':
+            swap_order(f)
+        elif start != b'\x80\x37\x12\x40':
             lament('not a .z64:', fn)
             return
 
+        f.seek(0)
+        data = f.read()
+
         outdir = hashlib.sha1(data).hexdigest()
         del data
 
@@ -123,8 +137,7 @@ def dump_rom(fn):
 
 def z_read_file(path, fn=None):
     if fn == None:
-        # TODO: infer from path
-        return False
+        fn = os.path.basename(path)
     if len(fn) < 37:
         return False
 
@@ -196,8 +209,9 @@ def create_rom(d):
         assert(f.tell() <= (pe or ve))
 
         # fix makerom (n64 header)
-        # TODO: don't assume bootcode is 6105
-        crc1, crc2 = n64.crc(f)
+        bootcode = n64.bootcode_version(f)
+        lament('bootcode:', bootcode)
+        crc1, crc2 = n64.crc(f, bootcode)
         lament('crcs: {:08X} {:08X}'.format(crc1, crc2))
         f.seek(0x10)
         f.write(W4(crc1))