Improved adding new content to title, fixed minor bugs

Dumping a title now properly updates the "number of contents" field in the TMD, so you're able to add more content than there was previously, and that new content will be added correctly.

src/libWiiPy/title/content.py

@@ -8,11 +8,19 @@ import io
 import hashlib
 from typing import List
 from dataclasses import dataclass as _dataclass
+from enum import IntEnum as _IntEnum
 from ..types import _ContentRecord
 from ..shared import _pad_bytes, _align_value
 from .crypto import decrypt_content, encrypt_content
 
 
+class ContentType(_IntEnum):
+    NORMAL = 1
+    HASH_TREE = 3
+    DLC = 16385
+    SHARED = 32769
+
+
 class ContentRegion:
     """
     A ContentRegion object to parse the continuous content region of a WAD. Allows for retrieving content from the
@@ -286,7 +294,7 @@ class ContentRegion:
         cid : int
             The Content ID to assign the new content in the content record.
         index : int
-            The index to place the new content at.
+            The index used when encrypting the new content.
         content_type : int
             The type of the new content.
         content_size : int
@@ -304,10 +312,11 @@ class ContentRegion:
         self.content_list.append(enc_content)
         self.content_records.append(_ContentRecord(cid, index, content_type, content_size, content_hash))
 
-    def add_content(self, dec_content: bytes, cid: int, index: int, content_type: int, title_key: bytes) -> None:
+    def add_content(self, dec_content: bytes, cid: int, content_type: int, title_key: bytes) -> None:
         """
-        Adds a new decrypted 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.
+        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
         provided Title Key before adding it to the ContentRegion.
@@ -318,13 +327,16 @@ class ContentRegion:
             The new decrypted content to add.
         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 other content in the ContentRegion.
         """
+        # Find the current highest content index and increment it for this content.
+        content_indices = []
+        for record in self.content_records:
+            content_indices.append(record.index)
+        index = max(content_indices) + 1
         content_size = len(dec_content)
         content_hash = str.encode(hashlib.sha1(dec_content).hexdigest())
         enc_content = encrypt_content(dec_content, title_key, index)

src/libWiiPy/title/title.py

@@ -76,7 +76,9 @@ class Title:
         if self.tmd.title_id == "0000000100000001":
             self.wad.wad_type = "ib"
         # Dump the TMD and set it in the WAD.
+        # This requires updating the content records and number of contents in the TMD first.
         self.tmd.content_records = self.content.content_records
+        self.tmd.num_contents = len(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())
@@ -247,7 +249,7 @@ class Title:
         """
         Sets the content at the provided content index to the provided new encrypted content. The provided hash and
         content size are set in the corresponding content record. A new Content ID or content type can also be
-        specified, but if it isn't than the current values are preserved.
+        specified, but if it isn't then the current values are preserved.
 
         This uses the content index, which 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
@@ -279,7 +281,7 @@ class Title:
         """
         Sets the content at the provided content index to the provided new decrypted content. The hash and content size
         of this content will be generated and then set in the corresponding content record. A new Content ID or content
-        type can also be specified, but if it isn't than the current values are preserved.
+        type can also be specified, but if it isn't then the current values are preserved.
 
         This also updates the content records in the TMD after the content is set.
 

src/libWiiPy/title/tmd.py

@@ -8,7 +8,7 @@ import binascii
 import hashlib
 import struct
 from typing import List
-from enum import IntEnum
+from enum import IntEnum as _IntEnum
 from ..types import _ContentRecord
 from ..shared import _bitmask
 from .util import title_ver_dec_to_standard, title_ver_standard_to_dec
@@ -390,7 +390,7 @@ class TMD:
             raise IndexError("Invalid content record! TMD lists '" + str(self.num_contents - 1) +
                              "' contents but index was '" + str(record) + "'!")
 
-    class AccessFlags(IntEnum):
+    class AccessFlags(_IntEnum):
         AHB = 0
         DVD_VIDEO = 1