From 98666285db169444b805197d28011b7a0b8e07b2 Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Fri, 3 May 2024 22:32:51 -0400
Subject: [PATCH] Improved comments, moved TID to IV conversion into a function
 in shared.py

---
 src/libWiiPy/crypto.py | 29 +++--------------------------
 src/libWiiPy/shared.py | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/src/libWiiPy/crypto.py b/src/libWiiPy/crypto.py
index d3194ec..066cb7b 100644
--- a/src/libWiiPy/crypto.py
+++ b/src/libWiiPy/crypto.py
@@ -1,9 +1,8 @@
 # "crypto.py" from libWiiPy by NinjaCheetah & Contributors
 # https://github.com/NinjaCheetah/libWiiPy
-import binascii
-
 import struct
 from .commonkeys import get_common_key
+from .shared import convert_tid_to_iv
 
 from Crypto.Cipher import AES
 
@@ -31,18 +30,7 @@ def decrypt_title_key(title_key_enc: bytes, common_key_index: int, title_id: byt
     # Load the correct common key for the title.
     common_key = get_common_key(common_key_index)
     # Convert the IV into the correct format based on the type provided.
-    title_key_iv = b''
-    if type(title_id) is bytes:
-        if len(title_id) == 16:
-            title_key_iv = binascii.unhexlify(title_id)
-        elif len(title_id) == 8:
-            pass
-        else:
-            raise ValueError("Title ID is not valid!")
-    elif type(title_id) is str:
-        title_key_iv = binascii.unhexlify(title_id)
-    else:
-        raise TypeError("Title ID is not valid! It must be either type str or bytes.")
+    title_key_iv = convert_tid_to_iv(title_id)
     # The IV will always be in the same format by this point, so add the last 8 bytes.
     title_key_iv = title_key_iv + (b'\x00' * 8)
     # Create a new AES object with the values provided.
@@ -75,18 +63,7 @@ def encrypt_title_key(title_key_dec: bytes, common_key_index: int, title_id: byt
     # Load the correct common key for the title.
     common_key = get_common_key(common_key_index)
     # Convert the IV into the correct format based on the type provided.
-    title_key_iv = b''
-    if type(title_id) is bytes:
-        if len(title_id) == 16:
-            title_key_iv = binascii.unhexlify(title_id)
-        elif len(title_id) == 8:
-            pass
-        else:
-            raise ValueError("Title ID is not valid!")
-    elif type(title_id) is str:
-        title_key_iv = binascii.unhexlify(title_id)
-    else:
-        raise TypeError("Title ID is not valid! It must be either type str or bytes.")
+    title_key_iv = convert_tid_to_iv(title_id)
     # The IV will always be in the same format by this point, so add the last 8 bytes.
     title_key_iv = title_key_iv + (b'\x00' * 8)
     # Create a new AES object with the values provided.
diff --git a/src/libWiiPy/shared.py b/src/libWiiPy/shared.py
index dd00602..d54af86 100644
--- a/src/libWiiPy/shared.py
+++ b/src/libWiiPy/shared.py
@@ -4,6 +4,8 @@
 # This file defines general functions that may be useful in other modules of libWiiPy. Putting them here cuts down on
 # clutter in other files.
 
+import binascii
+
 def align_value(value, alignment=64) -> int:
     """
     Aligns the provided value to the set alignment (defaults to 64).
@@ -45,3 +47,24 @@ def pad_bytes_stream(data, alignment=64) -> bytes:
     while (data.getbuffer().nbytes % alignment) != 0:
         data.write(b'\x00')
     return data
+
+
+def convert_tid_to_iv(title_id: str) -> bytes:
+    title_key_iv = b''
+    if type(title_id) is bytes:
+        # This catches the format b'0000000100000002'
+        if len(title_id) == 16:
+            title_key_iv = binascii.unhexlify(title_id)
+        # This catches the format b'\x00\x00\x00\x01\x00\x00\x00\x02'
+        elif len(title_id) == 8:
+            pass
+        # If it isn't one of those lengths, it cannot possibly be valid, so reject it.
+        else:
+            raise ValueError("Title ID is not valid!")
+    # Allow for a string like "0000000100000002"
+    elif type(title_id) is str:
+        title_key_iv = binascii.unhexlify(title_id)
+    # If the Title ID isn't bytes or a string, it isn't valid and is rejected.
+    else:
+        raise TypeError("Title ID type is not valid! It must be either type str or bytes.")
+    return title_key_iv