Added methods to set content in both enc and dec form to content.py

src/libWiiPy/content.py

@@ -8,7 +8,7 @@ import sys
 import hashlib
 from typing import List
 from .types import ContentRecord
-from .crypto import decrypt_content
+from .crypto import decrypt_content, encrypt_content
 
 
 class ContentRegion:
@@ -25,11 +25,21 @@ class ContentRegion:
     def __init__(self, content_region, content_records: List[ContentRecord]):
         self.content_region = content_region
         self.content_records = content_records
-        self.content_region_size: int  # Size of the content region.
+        self.content_region_size: int = 0  # Size of the content region.
-        self.num_contents: int  # Number of contents in the content region.
+        self.num_contents: int = 0  # Number of contents in the content region.
         self.content_start_offsets: List[int] = [0]  # The start offsets of each content in the content region.
+        self.content_list: List[bytes] = []
+        # Call load() to set all of the attributes from the raw content region provided.
+        self.load()
 
-        with io.BytesIO(content_region) as content_region_data:
+    def load(self):
+        """Loads the raw content region and builds a list of all the contents.
+
+        Returns
+        -------
+        none
+        """
+        with io.BytesIO(self.content_region) as content_region_data:
             # Get the total size of the content region.
             self.content_region_size = sys.getsizeof(content_region_data)
             self.num_contents = len(self.content_records)
@@ -41,6 +51,48 @@ class ContentRegion:
                 if (content.content_size % 64) != 0:
                     start_offset += 64 - (content.content_size % 64)
                 self.content_start_offsets.append(start_offset)
+            # Build a list of all the encrypted content data.
+            for content in range(len(self.content_start_offsets)):
+                # Seek to the start of the content based on the list of offsets.
+                content_region_data.seek(self.content_start_offsets[content])
+                # Calculate the number of bytes we need to read by adding bytes up the nearest multiple of 16 if needed.
+                bytes_to_read = self.content_records[content].content_size
+                if (bytes_to_read % 16) != 0:
+                    bytes_to_read += 16 - (bytes_to_read % 16)
+                # Read the file based on the size of the content in the associated record, then append that data to
+                # the list of content.
+                content_enc = content_region_data.read(bytes_to_read)
+                self.content_list.append(content_enc)
+
+    def dump(self) -> bytes:
+        """Takes the list of contents and assembles them back into one content region. Returns this content region as a
+        bytes object and sets the raw content region variable to this result, then calls load() again to make sure the
+        content list matches the raw data.
+
+        Returns
+        -------
+        bytes
+            The full WAD file as bytes.
+        """
+        # Open the stream and begin writing data to it.
+        with io.BytesIO() as content_region_data:
+            for content in self.content_list:
+                # Calculate padding after this content before the next one.
+                padding_bytes = 0
+                if (len(content) % 64) != 0:
+                    padding_bytes = 64 - (len(content) % 64)
+                # Write content data, then the padding afterward if necessary.
+                content_region_data.write(content)
+                if padding_bytes > 0:
+                    content_region_data.write(b'\x00' * padding_bytes)
+            content_region_data.seek(0x0)
+            self.content_region = content_region_data.read()
+        # Clear existing lists.
+        self.content_start_offsets = [0]
+        self.content_list = []
+        # Reload object's attributes to ensure the raw data and object match.
+        self.load()
+        return self.content_region
 
     def get_enc_content_by_index(self, index: int) -> bytes:
         """Gets an individual content from the content region based on the provided index, in encrypted form.
@@ -55,16 +107,8 @@ class ContentRegion:
         bytes
             The encrypted content listed in the content record.
         """
-        with io.BytesIO(self.content_region) as content_region_data:
+        content_enc = self.content_list[index]
-            # Seek to the start of the requested content based on the list of offsets.
+        return content_enc
-            content_region_data.seek(self.content_start_offsets[index])
-            # Calculate the number of bytes we need to read by adding bytes up the nearest multiple of 16 if needed.
-            bytes_to_read = self.content_records[index].content_size
-            if (bytes_to_read % 16) != 0:
-                bytes_to_read += 16 - (bytes_to_read % 16)
-            # Read the file based on the size of the content in the associated record.
-            content_enc = content_region_data.read(bytes_to_read)
-            return content_enc
 
     def get_enc_content_by_cid(self, cid: int) -> bytes:
         """Gets an individual content from the content region based on the provided Content ID, in encrypted form.
@@ -100,11 +144,7 @@ class ContentRegion:
         List[bytes]
             A list containing all encrypted contents.
         """
-        enc_contents: List[bytes] = []
+        return self.content_list
-        # Iterate over every content and add it to a list, then return it.
-        for content in range(self.num_contents):
-            enc_contents.append(self.get_enc_content_by_index(content))
-        return enc_contents
 
     def get_content_by_index(self, index: int, title_key: bytes) -> bytes:
         """Gets an individual content from the content region based on the provided index, in decrypted form.
@@ -183,3 +223,71 @@ class ContentRegion:
         for content in range(self.num_contents):
             dec_contents.append(self.get_content_by_index(content, title_key))
         return dec_contents
+
+    def set_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int,
+                        content_hash: bytes) -> None:
+        """Sets the provided index to a new content with the provided Content ID. Hashes and size of the content are
+        set in the content record, with a new record being added if necessary.
+
+        Parameters
+        ----------
+        enc_content : bytes
+            The new encrypted content to set.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        index : int
+            The index to place the new content at.
+        content_type : int
+            The type of the new content.
+        content_size : int
+            The size of the new encrypted content when decrypted.
+        content_hash : bytes
+            The hash of the new encrypted content when decrypted.
+        """
+        # Save the number of contents currently in the content region and records.
+        num_contents = len(self.content_records)
+        # Check if a record already exists for this index. If it doesn't, create it.
+        if (index + 1) > num_contents:
+            # Ensure that you aren't attempting to create a gap before appending.
+            if (index + 1) > num_contents + 1:
+                raise ValueError("You are trying to set the content at position " + str(index) + ", but no content "
+                                 "exists at position " + str(index - 1) + "!")
+            self.content_records.append(ContentRecord(cid, index, content_type, content_size, content_hash))
+        # If it does, reassign the values in it.
+        else:
+            self.content_records[index].content_id = cid
+            self.content_records[index].content_type = content_type
+            self.content_records[index].content_size = content_size
+            self.content_records[index].content_hash = content_hash
+        # Check if a content already occupies the provided index. If it does, reassign it to the new content, if it
+        # doesn't, then append a new entry.
+        if (index + 1) > num_contents:
+            self.content_list.append(enc_content)
+        else:
+            self.content_list[index] = enc_content
+
+    def set_content(self, dec_content: bytes, cid: int, index: int, content_type: int, title_key: bytes) -> None:
+        """Sets the provided index to a new content with the provided Content ID. Hashes and size of the content are
+                set in the content record, with a new record being added if necessary.
+
+        Parameters
+        ----------
+        dec_content : bytes
+            The new decrypted content to set.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        index : int
+            The index to place the new content at.
+        content_type : int
+            The type of the new content.
+        title_key : bytes
+            The Title Key that matches the new decrypted content.
+        """
+        # Store the size of the new content.
+        dec_content_size = len(dec_content)
+        # Calculate the hash of the new content.
+        dec_content_hash = str.encode(hashlib.sha1(dec_content).hexdigest())
+        # Encrypt the content using the provided Title Key and index.
+        enc_content = encrypt_content(dec_content, title_key, index)
+        # Pass values to set_enc_content()
+        self.set_enc_content(enc_content, cid, index, content_type, dec_content_size, dec_content_hash)

src/libWiiPy/title.py

@@ -44,7 +44,18 @@ class Title:
             raise ValueError("The Title IDs of the TMD and Ticket in this WAD do not match. This WAD appears to be "
                              "invalid.")
 
-    def set_title_id(self, title_id: str):
+    def dump(self) -> bytes:
+        """Dumps all title components (TMD, ticket, and content) back into the WAD object, and then dumps the WAD back
+        into raw data and returns it.
+
+        Returns
+        -------
+        wad_data : bytes
+            The raw data of the WAD.
+        """
+        # Dump the TMD.
+
+    def set_title_id(self, title_id: str) -> None:
         """Sets the Title ID of the title in both the TMD and Ticket.
 
         Parameters
@@ -56,3 +67,50 @@ class Title:
             raise ValueError("Invalid Title ID! Title IDs must be 8 bytes long.")
         self.tmd.set_title_id(title_id)
         self.ticket.set_title_id(title_id)
+
+    def set_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int,
+                        content_hash: bytes) -> None:
+        """Sets the provided index to a new content with the provided Content ID. Hashes and size of the content are
+        set in the content record, with a new record being added if necessary. The TMD is also updated to match the new
+        records.
+
+        Parameters
+        ----------
+        enc_content : bytes
+            The new encrypted content to set.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        index : int
+            The index to place the new content at.
+        content_type : int
+            The type of the new content.
+        content_size : int
+            The size of the new encrypted content when decrypted.
+        content_hash : bytes
+            The hash of the new encrypted content when decrypted.
+        """
+        # Set the encrypted content.
+        self.content.set_enc_content(enc_content, cid, index, content_type, content_size, content_hash)
+        # Update the TMD to match.
+        self.tmd.content_records = self.content.content_records
+
+    def set_content(self, dec_content: bytes, cid: int, index: int, content_type: int) -> None:
+        """Sets the provided index to a new content with the provided Content ID. Hashes and size of the content are
+        set in the content record, with a new record being added if necessary. The Title Key is sourced from this
+        title's loaded ticket. The TMD is also updated to match the new records.
+
+        Parameters
+        ----------
+        dec_content : bytes
+            The new decrypted content to set.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        index : int
+            The index to place the new content at.
+        content_type : int
+            The type of the new content.
+        """
+        # Set the decrypted content.
+        self.content.set_content(dec_content, cid, index, content_type, self.ticket.get_title_key())
+        # Update the TMD to match.
+        self.tmd.content_records = self.content.content_records

src/libWiiPy/tmd.py

@@ -134,6 +134,7 @@ class TMD:
             tmd_data.seek(0x1E0)
             self.boot_index = int.from_bytes(tmd_data.read(2))
             # Get content records for the number of contents in num_contents.
+            self.content_records = []
             for content in range(0, self.num_contents):
                 tmd_data.seek(0x1E4 + (36 * content))
                 content_record_hdr = struct.unpack(">LHH4x4s20s", tmd_data.read(36))
@@ -214,6 +215,8 @@ class TMD:
             # Set the TMD attribute of the object to the new raw TMD.
             tmd_data.seek(0x0)
             self.tmd = tmd_data.read()
+        # Clear existing lists.
+        self.content_records = []
         # Reload object's attributes to ensure the raw data and object match.
         self.load()
         return self.tmd