Add IOSPatcher to apply patches to IOS WADs loaded into a Title()

Also fixes a MAJOR bug with WAD packing where changes to the content records would be dropped when dumping a WAD (!!)

src/libWiiPy/title/__init__.py

@@ -3,6 +3,7 @@
 
 from .content import *
 from .crypto import *
+from .iospatcher import *
 from .nus import *
 from .ticket import *
 from .title import *

src/libWiiPy/title/iospatcher.py

@@ -0,0 +1,168 @@
+# "title/iospatcher.py" from libWiiPy by NinjaCheetah & Contributors
+# https://github.com/NinjaCheetah/libWiiPy
+#
+# Module for applying patches to IOS WADs via a Title().
+
+import io
+from .title import Title
+
+
+class IOSPatcher:
+    """
+    An IOSPatcher object that allows for applying patches to IOS WADs loaded into Title objects.
+
+    Attributes
+    ----------
+    """
+    def __init__(self):
+        self.title: Title = Title()
+        self.es_module_index: int = -1
+
+    def load(self, title: Title) -> None:
+        """
+        Loads a Title object containing an IOS WAD and locates the content containing the ES module that needs to be
+        patched.
+
+        Parameters
+        ----------
+        title : Title
+            A Title object containing the IOS to be patched.
+        """
+        # Check to ensure that this Title contains IOS. IOS always has a TID high of 00000001, and any TID low after
+        # 00000002.
+        tid = title.tmd.title_id
+        if tid[:8] != "00000001" or tid[8:] == "00000001" or tid[8:] == "00000002":
+            raise ValueError("This Title does not contain an IOS! Cannot load Title for patching.")
+
+        # Now that we know this is IOS, we need to go ahead and check all of its contents until we find the one that
+        # contains the ES module, since that's what we're patching.
+        es_content_index = -1
+        for content in range(len(title.content.content_records)):
+            target_content = title.get_content_by_index(title.content.content_records[content].index)
+            es_offset = target_content.find(b'\x45\x53\x3A')  # This is looking for "ES:"
+            if es_offset != -1:
+                es_content_index = title.content.content_records[content].index
+                break
+
+        # If we get here with no content index, then ES wasn't found. That probably means that this isn't IOS.
+        if es_content_index == -1:
+            raise Exception("ES module could not be found! Please ensure that this is an intact copy of an IOS.")
+
+        self.title = title
+        self.es_module_index = es_content_index
+
+    def dump(self) -> Title:
+        """
+        Returns the patched Title object.
+
+        Returns
+        -------
+        Title
+            The patched Title object.
+        """
+        return self.title
+
+    def patch_all(self) -> None:
+        """
+        Applies all patches to patch in fakesigning, ES_Identify access, /dev/flash access, and the version patch.
+        """
+        self.patch_fakesigning()
+        self.patch_es_identify()
+        self.patch_nand_access()
+        self.patch_version_patch()
+
+    def patch_fakesigning(self) -> None:
+        """
+        Patches the trucha/fakesigning bug back into the IOS' ES module to allow it to accept fakesigned TMDs and
+        Tickets.
+        """
+        if self.es_module_index == -1:
+            raise Exception("No valid IOS is loaded! Patching cannot continue.")
+
+        target_content = self.title.get_content_by_index(self.es_module_index)
+
+        patch_count = 0
+        patch_sequences = [b'\x20\x07\x23\xa2', b'\x20\x07\x4b\x0b']
+        for sequence in patch_sequences:
+            start_offset = target_content.find(sequence)
+            if start_offset != -1:
+                with io.BytesIO(target_content) as content_data:
+                    content_data.seek(start_offset + 1)
+                    content_data.write(b'\x00')
+                    content_data.seek(0)
+                    target_content = content_data.read()
+                    patch_count += 1
+
+        # If neither structure was found, then no patches could be applied, so return an error.
+        if patch_count == 0:
+            raise Exception("No patches could be applied because the required data could not be found!")
+
+        self.title.set_content(target_content, self.es_module_index)
+
+    def patch_es_identify(self) -> None:
+        """
+        Patches the ability to call ES_Identify back into the IOS' ES module to allow for changing the permissions of a
+        title.
+        """
+        if self.es_module_index == -1:
+            raise Exception("No valid IOS is loaded! Patching cannot continue.")
+
+        target_content = self.title.get_content_by_index(self.es_module_index)
+
+        patch_sequence = b'\x28\x03\xd1\x23'
+        start_offset = target_content.find(patch_sequence)
+        if start_offset != -1:
+            with io.BytesIO(target_content) as content_data:
+                content_data.seek(start_offset + 2)
+                content_data.write(b'\x00\x00')
+                content_data.seek(0)
+                target_content = content_data.read()
+        else:
+            raise Exception("No patches could be applied because the required data could not be found!")
+
+        self.title.set_content(target_content, self.es_module_index)
+
+    def patch_nand_access(self) -> None:
+        """
+        Patches the ability to directly access /dev/flash back into the IOS' ES module to allow for raw access to the
+        Wii's filesystem.
+        """
+        if self.es_module_index == -1:
+            raise Exception("No valid IOS is loaded! Patching cannot continue.")
+
+        target_content = self.title.get_content_by_index(self.es_module_index)
+
+        patch_sequence = b'\x42\x8b\xd0\x01\x25\x66'
+        start_offset = target_content.find(patch_sequence)
+        if start_offset != -1:
+            with io.BytesIO(target_content) as content_data:
+                content_data.seek(start_offset + 2)
+                content_data.write(b'\xe0')
+                content_data.seek(0)
+                target_content = content_data.read()
+        else:
+            raise Exception("No patches could be applied because the required data could not be found!")
+
+        self.title.set_content(target_content, self.es_module_index)
+
+    def patch_version_patch(self) -> None:
+        """
+        Patches the ability to idk man do something I guess back into IOS' ES module. (Awaiting a real explanation)
+        """
+        if self.es_module_index == -1:
+            raise Exception("No valid IOS is loaded! Patching cannot continue.")
+
+        target_content = self.title.get_content_by_index(self.es_module_index)
+
+        patch_sequence = b'\xd2\x01\x4e\x56'
+        start_offset = target_content.find(patch_sequence)
+        if start_offset != -1:
+            with io.BytesIO(target_content) as content_data:
+                content_data.seek(start_offset)
+                content_data.write(b'\xe0')
+                content_data.seek(0)
+                target_content = content_data.read()
+        else:
+            raise Exception("No patches could be applied because the required data could not be found!")
+
+        self.title.set_content(target_content, self.es_module_index)

src/libWiiPy/title/title.py

@@ -74,6 +74,7 @@ class Title:
         if self.tmd.title_id == "0000000100000001":
             self.wad.wad_type = "ib"
         # Dump the TMD and set it in the WAD.
+        self.tmd.content_records = self.content.content_records
         self.wad.set_tmd_data(self.tmd.dump())
         # Dump the Ticket and set it in the WAD.
         self.wad.set_ticket_data(self.ticket.dump())