From f605b3a8a7c4e97da8586c5d15153bb3fcc02be5 Mon Sep 17 00:00:00 2001
From: Ethan Roseman <ethteck@gmail.com>
Date: Thu, 23 Apr 2020 14:53:08 -0400
Subject: [PATCH] fixbaserom cleanup

---
 fixbaserom.py | 40 ++++++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/fixbaserom.py b/fixbaserom.py
index 0d6777e87b..e00853e868 100644
--- a/fixbaserom.py
+++ b/fixbaserom.py
@@ -1,39 +1,48 @@
-import os.path
 from os import path
 import sys
 import struct
 import hashlib
 
 
+def get_str_hash(byte_array):
+    return str(hashlib.md5(byte_array).hexdigest())
+
+
+# If the baserom exists and is correct, we don't need to change anything
+if path.exists("baserom.z64"):
+    with open("baserom.z64", mode="rb") as file:
+        fileContent = bytearray(file.read())
+        if get_str_hash(fileContent) == "f0b7f35375f9cc8ca1b2d59d78e35405":
+            print("Found valid baserom - exiting early")
+            sys.exit(0)
+
 # Determine if we have a ROM file
 romFileName = ""
-if (path.exists("baserom.z64")):
-    romFileName = "baserom.z64"
-if (path.exists("baserom_original.z64")):
+if path.exists("baserom_original.z64"):
     romFileName = "baserom_original.z64"
-elif (path.exists("baserom_original.n64")):
+elif path.exists("baserom_original.n64"):
     romFileName = "baserom_original.n64"
 
 # Read in the original ROM
-if (romFileName != ""):
+if romFileName != "":
     print("File '" + romFileName + "' found.")
-    with open(romFileName, mode='rb') as file:
+    with open(romFileName, mode="rb") as file:
         fileContent = bytearray(file.read())
 
         # Check if ROM needs to be byte swapped
-        if (fileContent[0] == 0x40):
-             # Byte Swap ROM
+        if fileContent[0] == 0x40:
+            # Byte Swap ROM
             # TODO: This is pretty slow at the moment. Look into optimizing it later...
             print("ROM needs to be byte swapped...")
             i = 0
-            while (i < len(fileContent)):
+            while i < len(fileContent):
                 tmp = struct.unpack_from("BBBB", fileContent, i)
                 struct.pack_into("BBBB", fileContent, i + 0, tmp[3], tmp[2], tmp[1], tmp[0])
                 i += 4
 
                 perc = float(i) / float(len(fileContent))
 
-                if (i % (1024 * 1024 * 4) == 0):
+                if i % (1024 * 1024 * 4) == 0:
                     print(str(perc * 100) + "%")
         
             print("Byte swapping done.")
@@ -50,10 +59,10 @@ print("Patching header...")
 strippedContent[0x3E] = 0x50
 
 # Check to see if the ROM is a "vanilla" Debug ROM
-md5Hash = hashlib.md5(bytes(strippedContent)).hexdigest()
-
-if (str(md5Hash) != "f0b7f35375f9cc8ca1b2d59d78e35405"):
-    print("Error: Expected a hash of f0b7f35375f9cc8ca1b2d59d78e35405 but got " + str(md5Hash) + ". The baserom has probably been tampered, find a new one")
+str_hash = get_str_hash(strippedContent)
+if str_hash != "f0b7f35375f9cc8ca1b2d59d78e35405":
+    print("Error: Expected a hash of f0b7f35375f9cc8ca1b2d59d78e35405 but got " + str_hash + ". " +
+          "The baserom has probably been tampered, find a new one")
     sys.exit(1)
 
 # Write out our new ROM
@@ -62,4 +71,3 @@ with open("baserom.z64", mode="wb") as file:
     file.write(bytes(strippedContent))
 
 print("Done!")
-