From 74192f8febb15c4a96f4282877ecd36f6fefd76b Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Wed, 6 Mar 2024 22:48:51 -0500
Subject: [PATCH] Experimental content re-encrypting code

---
 src/libWiiPy/content.py |  2 +-
 src/libWiiPy/crypto.py  | 47 ++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/src/libWiiPy/content.py b/src/libWiiPy/content.py
index ebd55dc..f769e2b 100644
--- a/src/libWiiPy/content.py
+++ b/src/libWiiPy/content.py
@@ -91,7 +91,7 @@ class ContentRegion:
         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! The incorrect Title Key may"
+            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()))
diff --git a/src/libWiiPy/crypto.py b/src/libWiiPy/crypto.py
index a112e2f..9ea1256 100644
--- a/src/libWiiPy/crypto.py
+++ b/src/libWiiPy/crypto.py
@@ -39,7 +39,8 @@ def decrypt_title_key(title_key_enc, common_key_index, title_id) -> bytes:
 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.
+    This requires the index of the content to decrypt as it is used as the IV, as well as the content length to adjust
+    padding as necessary.
 
     Parameters
     ----------
@@ -72,3 +73,47 @@ def decrypt_content(content_enc, title_key, content_index, content_length) -> by
     while len(content_dec) > content_length:
         content_dec = content_dec[:-1]
     return content_dec
+
+
+def encrypt_content(content_dec, title_key, content_index) -> bytes:
+    """Gets the encrypted version of the decrypted content.
+
+    This requires the index of the content to encrypt as it is used as the IV, as well as the content length to adjust
+    padding as necessary.
+
+    Parameters
+    ----------
+    content_dec : bytes
+        The decrypted 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.
+
+    Returns
+    -------
+    bytes
+        The decrypted content.
+    """
+    # Generate the IV from the Content Index of the content to be decrypted.
+    content_index_bin = struct.pack('>H', content_index)
+    while len(content_index_bin) < 16:
+        content_index_bin += b'\x00'
+    # Align content to 64 bytes to ensure that all the data is being encrypted, and so it works with AES encryption.
+    bytes_added = None
+    if (len(content_dec) % 64) != 0:
+        bytes_added = len(b'\x00' * (64 - (len(content_dec) % 64)))
+        print(bytes_added)
+        content_dec = content_dec + (b'\x00' * (64 - (len(content_dec) % 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)
+    # Encrypt the content using the AES object.
+    content_enc = aes.encrypt(content_dec)
+    # Remove any bytes added.
+    if bytes_added:
+        while bytes_added:
+            content_enc = content_enc[:-1]
+            bytes_added -= 1
+            print("removing " + str(bytes_added))
+    print(str(len(content_enc)))
+    return content_enc