From 183498025a4b4e92800a946f9bd9614ab0f2539f Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Thu, 11 Jul 2024 21:11:58 +1000
Subject: [PATCH] Add checks to nus content subcommand, and improve error
 handling

---
 modules/nus.py | 52 ++++++++++++++++++++++++++++++++++++++++----------
 wiipy.py       |  2 ++
 2 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/modules/nus.py b/modules/nus.py
index 0fd645d..f0ee122 100644
--- a/modules/nus.py
+++ b/modules/nus.py
@@ -142,28 +142,52 @@ def handle_nus_content(args):
     tid = args.tid
     cid = args.cid
     version = args.version
+    out = args.output
     if args.decrypt:
         decrypt_content = True
     else:
         decrypt_content = False
 
-    content_id = int.from_bytes(binascii.unhexlify(cid))
+    # Only accepting the 000000xx format because it's the one that would be most commonly known, rather than using the
+    # actual integer that the hex Content ID translates to.
+    try:
+        content_id = int.from_bytes(binascii.unhexlify(cid))
+    except binascii.Error:
+        print("Invalid Content ID! Content ID must be in format \"000000xx\"!")
+        return
 
-    content_file_name = hex(content_id)[2:]
-    while len(content_file_name) < 8:
-        content_file_name = "0" + content_file_name
+    # Use the supplied output path if one was specified, otherwise generate one using the Content ID.
+    if out is None:
+        content_file_name = hex(content_id)[2:]
+        while len(content_file_name) < 8:
+            content_file_name = "0" + content_file_name
+        output_path = pathlib.Path(content_file_name)
+    else:
+        output_path = pathlib.Path(out)
 
-    content_data = libWiiPy.title.download_content(tid, content_id)
+    # Try to download the content, and catch the ValueError libWiiPy will throw if it can't be found.
+    print("Downloading content with Content ID " + cid + "...")
+    try:
+        content_data = libWiiPy.title.download_content(tid, content_id)
+    except ValueError:
+        print("The Title ID or Content ID you specified could not be found!")
+        return
 
     if decrypt_content is True:
-        content_file_name = content_file_name + ".app"
+        # Ensure that a version was supplied, because we need the matching TMD for decryption to work.
+        if version is None:
+            print("You must specify the version that the requested content belongs to for decryption!")
+            return
+
+        output_path = output_path.with_suffix(".app")
         tmd = libWiiPy.title.TMD()
         tmd.load(libWiiPy.title.download_tmd(tid, version))
+        # Try to get a Ticket for the title, if a common one is available.
         try:
             ticket = libWiiPy.title.Ticket()
-            ticket.load(libWiiPy.title.download_ticket(tid))
+            ticket.load(libWiiPy.title.download_ticket(tid, wiiu_endpoint=True))
         except ValueError:
-            print("  - No Ticket is available!")
+            print("No Ticket is available! Content cannot be decrypted!")
             return
 
         content_hash = 'gggggggggggggggggggggggggggggggggggggggg'
@@ -175,17 +199,25 @@ def handle_nus_content(args):
                 content_size = record.content_size
                 content_index = record.index
 
+        # If the default hash never changed, then a content record matching the downloaded content couldn't be found,
+        # which most likely means that the wrong version was specified.
+        if content_hash == 'gggggggggggggggggggggggggggggggggggggggg':
+            print("Content was not found in the TMD from the specified version! Content cannot be decrypted!")
+            return
+
+        # Manually decrypt the content and verify its hash, which is what libWiiPy's get_content() methods do. We just
+        # can't really use that here because that require setting up a lot more of the title than is necessary.
         content_dec = libWiiPy.title.decrypt_content(content_data, ticket.get_title_key(), content_index, content_size)
         content_dec_hash = hashlib.sha1(content_dec).hexdigest()
         if content_hash != content_dec_hash:
             raise ValueError("The decrypted content provided does not match the record at the provided index. \n"
                              "Expected hash is: {}\n".format(content_hash) +
                              "Actual hash is: {}".format(content_dec_hash))
-        file = open(content_file_name, "wb")
+        file = open(output_path, "wb")
         file.write(content_dec)
         file.close()
     else:
-        file = open(content_file_name, "wb")
+        file = open(output_path, "wb")
         file.write(content_data)
         file.close()
 
diff --git a/wiipy.py b/wiipy.py
index 7b00388..0daa531 100644
--- a/wiipy.py
+++ b/wiipy.py
@@ -57,6 +57,8 @@ if __name__ == "__main__":
                                     help="Content ID to download (in \"000000xx\" format)")
     nus_content_parser.add_argument("-v", "--version", metavar="VERSION", type=int,
                                     help="version this content belongs to (required for decryption)")
+    nus_content_parser.add_argument("-o", "--output", metavar="OUT", type=str,
+                                    help="path to download the content to (optional)")
     nus_content_parser.add_argument("-d", "--decrypt", action="store_true", help="decrypt this content")
 
     # Argument parser for the U8 subcommand.