From 7ca46372b0c1c7caccfe1273b59f3ab9c87608f1 Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Mon, 4 Mar 2024 20:48:38 -0500
Subject: [PATCH] Full content extraction is working!

---
 src/libWiiPy/content.py | 24 ++++++++++++++----------
 src/libWiiPy/crypto.py  | 29 ++++++++++++++---------------
 src/libWiiPy/ticket.py  |  2 +-
 src/libWiiPy/types.py   | 10 +++++-----
 4 files changed, 34 insertions(+), 31 deletions(-)

diff --git a/src/libWiiPy/content.py b/src/libWiiPy/content.py
index 5eebe2b..96dbce2 100644
--- a/src/libWiiPy/content.py
+++ b/src/libWiiPy/content.py
@@ -34,8 +34,12 @@ class ContentRegion:
             self.content_region_size = sys.getsizeof(content_region_data)
             self.num_contents = len(self.content_records)
             # Calculate the offsets of each content in the content region.
+            # Content is aligned to 16 bytes, however a new content won't start until the next multiple of 64 bytes.
+            # Because of this, we need to add bytes to the next 64 byte offset if the previous content wasn't that long.
             for content in self.content_records[:-1]:
-                start_offset = int(64 * round(content.content_size / 64)) + self.content_start_offsets[-1]
+                start_offset = content.content_size + self.content_start_offsets[-1]
+                if (content.content_size % 64) != 0:
+                    start_offset += 64 - (content.content_size % 64)
                 self.content_start_offsets.append(start_offset)
 
     def get_enc_content(self, index: int) -> bytes:
@@ -54,8 +58,10 @@ class ContentRegion:
         with io.BytesIO(self.content_region) as content_region_data:
             # Seek to the start of the requested content based on the list of offsets.
             content_region_data.seek(self.content_start_offsets[index])
-            # Calculate the number of bytes we need to read by rounding the size to the nearest 16 bytes.
-            bytes_to_read = int(64 * round(self.content_records[index].content_size / 64))
+            # 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
@@ -77,18 +83,16 @@ class ContentRegion:
         """
         # Load the encrypted content at the specified index and then decrypt it with the Title Key.
         content_enc = self.get_enc_content(index)
-        content_dec = decrypt_content(content_enc, title_key, self.content_records[index].index, self.content_records[index].content_size)
+        content_dec = decrypt_content(content_enc, title_key, self.content_records[index].index,
+                                      self.content_records[index].content_size)
         # Hash the decrypted content and ensure that the hash matches the one in its Content Record.
         # If it does not, then something has gone wrong in the decryption, and an error will be thrown.
         content_dec_hash = hashlib.sha1(content_dec)
         content_record_hash = str(self.content_records[index].content_hash.decode())
+        # Compare the hash and throw a ValueError if the hash doesn't match.
         if content_dec_hash.hexdigest() != content_record_hash:
-            #raise ValueError("Content hash did not match the expected hash in its record! This most likely means that "
-                             #"the incorrect Title Key was used for this content.\n"
-                             #"Expected hash is: {}\n".format(content_record_hash) +
-                             #"Actual hash is: {}".format(content_dec_hash.hexdigest()))
-            print("Content hash did not match the expected hash in its record! This most likely means that "
-                             "the incorrect Title Key was used for this content.\n"
+            raise ValueError("Content hash did not match the expected hash in its record! The incorrect Title Key may"
+                             "have been used!.\n"
                              "Expected hash is: {}\n".format(content_record_hash) +
                              "Actual hash is: {}".format(content_dec_hash.hexdigest()))
         return content_dec
diff --git a/src/libWiiPy/crypto.py b/src/libWiiPy/crypto.py
index 4e8abef..2a6d538 100644
--- a/src/libWiiPy/crypto.py
+++ b/src/libWiiPy/crypto.py
@@ -1,7 +1,5 @@
 # "crypto.py" from libWiiPy by NinjaCheetah & Contributors
 # https://github.com/NinjaCheetah/libWiiPy
-#
-# See https://wiibrew.org/wiki/Ticket for details about the TMD format
 
 import struct
 from .commonkeys import get_common_key
@@ -9,7 +7,7 @@ from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
 
 
-def decrypt_title_key(title_key_enc, common_key_index, title_id):
+def decrypt_title_key(title_key_enc, common_key_index, title_id) -> bytes:
     """Gets the decrypted version of the encrypted Title Key provided.
 
     Requires the index of the common key to use, and the Title ID of the title that the Title Key is for.
@@ -21,7 +19,7 @@ def decrypt_title_key(title_key_enc, common_key_index, title_id):
     common_key_index : int
         The index of the common key to be returned.
     title_id : bytes
-        The title ID of the tite that the key is for.
+        The title ID of the title that the key is for.
 
     Returns
     -------
@@ -39,7 +37,7 @@ def decrypt_title_key(title_key_enc, common_key_index, title_id):
     return title_key
 
 
-def decrypt_content(content_enc, title_key, content_index, content_length):
+def decrypt_content(content_enc, title_key, content_index, content_length) -> bytes:
     """Gets the decrypted version of the encrypted content.
 
     Requires the index of the common key to use, and the Title ID of the title that the Title Key is for.
@@ -50,6 +48,10 @@ def decrypt_content(content_enc, title_key, content_index, content_length):
         The encrypted content.
     title_key : bytes
         The Title Key for the title the content is from.
+    content_index : int
+        The index in the TMD's content record of the content being decrypted.
+    content_length : int
+        The length in the TMD's content record of the content being decrypted.
 
     Returns
     -------
@@ -60,18 +62,15 @@ def decrypt_content(content_enc, title_key, content_index, content_length):
     content_index_bin = struct.pack('>H', content_index)
     while len(content_index_bin) < 16:
         content_index_bin += b'\x00'
-    # In CBC mode, content must be padded out to a 16-byte boundary, so do that here, and then remove bytes added after.
-    padded = False
-    if (len(content_enc) % 128) != 0:
-        print("needs padding to 16 bytes")
-        content_enc = pad(content_enc, 128, "pkcs7")
-        padded = True
+    # Align content to 64 bytes to ensure that all the data is being decrypted, and so it works with AES encryption.
+    if (len(content_enc) % 64) != 0:
+        print("needs padding to 64 bytes")
+        content_enc = content_enc + (b'\x00' * (64 - (len(content_enc) % 64)))
     # Create a new AES object with the values provided, with the content's unique ID as the IV.
     aes = AES.new(title_key, AES.MODE_CBC, content_index_bin)
     # Decrypt the content using the AES object.
     content_dec = aes.decrypt(content_enc)
-    # Remove padding bytes, if any were added.
-    if padded:
-        while len(content_dec) > content_length:
-            content_dec = content_dec[:-1]
+    # Trim additional bytes that may have been added so the content is the correct size.
+    while len(content_dec) > content_length:
+        content_dec = content_dec[:-1]
     return content_dec
diff --git a/src/libWiiPy/ticket.py b/src/libWiiPy/ticket.py
index 85a99de..ee9f11b 100644
--- a/src/libWiiPy/ticket.py
+++ b/src/libWiiPy/ticket.py
@@ -1,7 +1,7 @@
 # "ticket.py" from libWiiPy by NinjaCheetah & Contributors
 # https://github.com/NinjaCheetah/libWiiPy
 #
-# See https://wiibrew.org/wiki/Ticket for details about the TMD format
+# See https://wiibrew.org/wiki/Ticket for details about the ticket format
 
 import io
 from .crypto import decrypt_title_key
diff --git a/src/libWiiPy/types.py b/src/libWiiPy/types.py
index 85dd87b..da78e14 100644
--- a/src/libWiiPy/types.py
+++ b/src/libWiiPy/types.py
@@ -22,8 +22,8 @@ class ContentRecord:
     content_hash
         The SHA-1 hash of the decrypted content.
     """
-    content_id: int  # Unique ID for the current content
-    index: int  # Index in the list of contents
-    content_type: int  # Type of content, possible values of: 0x0001: Normal, 0x4001: DLC, 0x8001: Shared
-    content_size: int  # Size of the current content
-    content_hash: bytes  # SHA1 hash of the current content
+    content_id: int  # The unique ID of the content.
+    index: int  # The index of this content in the content record.
+    content_type: int  # Type of content, possible values of: 0x0001: Normal, 0x4001: DLC, 0x8001: Shared.
+    content_size: int  # Size of the content when decrypted.
+    content_hash: bytes  # SHA-1 hash of the content when decrypted.