Added methods to get a content index from a CID and add content from a Title()

src/libWiiPy/title/content.py

@@ -158,16 +158,11 @@ class ContentRegion:
         bytes
             The encrypted content listed in the content record.
         """
-        # Get a list of the current Content IDs, so we can make sure the target one exists.
-        content_ids = []
-        for record in self.content_records:
-            content_ids.append(record.content_id)
-        if cid not in content_ids:
-            raise ValueError("You are trying to get a content with Content ID " + str(cid) + ", but no content with "
-                             "that ID exists!")
-        # Get the content index associated with the CID we now know exists.
-        target_index = content_ids.index(cid)
-        content_index = self.content_records[target_index].index
+        try:
+            content_index = self.get_index_from_cid(cid)
+        except ValueError:
+            raise ValueError(f"You are trying to get a content with Content ID {cid}, "
+                             f"but no content with that ID exists!")
         content_enc = self.get_enc_content_by_index(content_index)
         return content_enc
 
@@ -246,16 +241,11 @@ class ContentRegion:
         bytes
             The decrypted content listed in the content record.
         """
-        # Get a list of the current Content IDs, so we can make sure the target one exists.
-        content_ids = []
-        for record in self.content_records:
-            content_ids.append(record.content_id)
-        if cid not in content_ids:
-            raise ValueError("You are trying to get a content with Content ID " + str(cid) + ", but no content with "
-                             "that ID exists!")
-        # Get the content index associated with the CID we now know exists.
-        target_index = content_ids.index(cid)
-        content_index = self.content_records[target_index].index
+        try:
+            content_index = self.get_index_from_cid(cid)
+        except ValueError:
+            raise ValueError(f"You are trying to get a content with Content ID {cid}, "
+                             f"but no content with that ID exists!")
         content_dec = self.get_content_by_index(content_index, title_key, skip_hash)
         return content_dec
 
@@ -281,6 +271,32 @@ class ContentRegion:
             dec_contents.append(self.get_content_by_index(content, title_key, skip_hash))
         return dec_contents
 
+    def get_index_from_cid(self, cid: int) -> int:
+        """
+        Gets the content index of a content by its Content ID. The returned index is the value tied to each content and
+        used as the IV for encryption, rather than the literal index in the array of content, because sometimes the
+        contents end up out of order in a WAD while still retaining the original indices.
+
+        Parameters
+        ----------
+        cid : int
+            The Content ID to get the index of.
+
+        Returns
+        -------
+        int
+            The content index.
+        """
+        # Get a list of the current Content IDs, so we can make sure the target one exists.
+        content_ids = []
+        for record in self.content_records:
+            content_ids.append(record.content_id)
+        if cid not in content_ids:
+            raise ValueError("The specified Content ID does not exist!")
+        literal_index = content_ids.index(cid)
+        target_index = self.content_records[literal_index].index
+        return target_index
+
     def add_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int,
                         content_hash: bytes) -> None:
         """
@@ -534,16 +550,11 @@ class ContentRegion:
         cid : int
             The Content ID of the content you want to remove.
         """
-        # Get a list of the current Content IDs, so we can make sure the target one exists.
-        content_ids = []
-        for record in self.content_records:
-            content_ids.append(record.content_id)
-        if cid not in content_ids:
-            raise ValueError("You are trying to remove content with Content ID " + str(cid) + ", but no content with "
-                             "that ID exists!")
-        # Remove the content index associated with the CID we now know exists.
-        target_index = content_ids.index(cid)
-        content_index = self.content_records[target_index].index
+        try:
+            content_index = self.get_index_from_cid(cid)
+        except ValueError:
+            raise ValueError(f"You are trying to remove content with Content ID {cid}, "
+                             f"but no content with that ID exists!")
         self.remove_content_by_index(content_index)
 
 

src/libWiiPy/title/title.py

@@ -244,6 +244,55 @@ class Title:
         blocks = math.ceil(title_size_bytes / 131072)
         return blocks
 
+    def add_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int,
+                        content_hash: bytes) -> None:
+        """
+        Adds a new encrypted content to the ContentRegion, and adds the provided Content ID, index, content type,
+        content size, and content hash to a new record in the ContentRecord list.
+
+        Parameters
+        ----------
+        enc_content : bytes
+            The new encrypted content to add.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        index : int
+            The index used when encrypting the new content.
+        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.
+        """
+        # Add the encrypted content.
+        self.content.add_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 add_content(self, dec_content: bytes, cid: int, content_type: int) -> None:
+        """
+        Adds a new decrypted content to the end of the ContentRegion, and adds the provided Content ID, content type,
+        content size, and content hash to a new record in the ContentRecord list. The index will be automatically
+        assigned by incrementing the current highest index in the records.
+
+        This first gets the content hash and size from the provided data, and then encrypts the content with the
+        Title Key before adding it to the ContentRegion.
+
+        Parameters
+        ----------
+        dec_content : bytes
+            The new decrypted content to add.
+        cid : int
+            The Content ID to assign the new content in the content record.
+        content_type : int
+            The type of the new content.
+        """
+        # Add the decrypted content.
+        self.content.add_content(dec_content, cid, content_type, self.ticket.get_title_key())
+        # Update the TMD to match.
+        self.tmd.content_records = self.content.content_records
+
     def set_enc_content(self, enc_content: bytes, index: int, content_size: int, content_hash: bytes, cid: int = None,
                         content_type: int = None) -> None:
         """