Allow for custom endpoints for NUS downloads

NinjaCheetah · NinjaCheetah · commit 71c0726c4f78 · 2025-01-12T23:01:15.000-05:00
diff --git a/commands/title/info.py b/commands/title/info.py
@@ -43,6 +43,9 @@ def _print_tmd_info(tmd: libWiiPy.title.TMD, signing_cert=None):
     elif tmd.signature_issuer.find("CP00000007") != -1:
         print(f"  Certificate: CP00000007 (Development)")
         print(f"  Certificate Issuer: Root-CA00000002 (Development)")
+    elif tmd.signature_issuer.find("CP00000005") != -1:
+        print(f"  Certificate: CP00000005 (Development/Unknown)")
+        print(f"  Certificate Issuer: Root-CA00000002 (Development)")
     elif tmd.signature_issuer.find("CP10000000") != -1:
         print(f"  Certificate: CP10000000 (Arcade)")
         print(f"  Certificate Issuer: Root-CA10000000 (Arcade)")
@@ -128,6 +131,9 @@ def _print_ticket_info(ticket: libWiiPy.title.Ticket, signing_cert=None):
     elif ticket.signature_issuer.find("XS00000006") != -1:
         print(f"  Certificate: XS00000006 (Development)")
         print(f"  Certificate Issuer: Root-CA00000002 (Development)")
+    elif ticket.signature_issuer.find("XS00000004") != -1:
+        print(f"  Certificate: XS00000004 (Development/Unknown)")
+        print(f"  Certificate Issuer: Root-CA00000002 (Development)")
     else:
         print(f"  Certificate Info: {ticket.signature_issuer} (Unknown)")
     match ticket.common_key_index:
diff --git a/commands/title/nus.py b/commands/title/nus.py
@@ -8,16 +8,91 @@
 from modules.core import fatal_error
 
 
+def handle_nus_content(args):
+    tid = args.tid
+    cid = args.cid
+    version = args.version
+    if args.decrypt:
+        decrypt_content = True
+    else:
+        decrypt_content = False
+
+    # 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.
+    content_id = None
+    try:
+        content_id = int.from_bytes(binascii.unhexlify(cid))
+    except binascii.Error:
+        fatal_error("The provided Content ID is invalid! The Content ID must be in the format \"000000xx\"!")
+
+    # Use the supplied output path if one was specified, otherwise generate one using the Content ID.
+    if args.output is None:
+        content_file_name = f"{content_id:08X}".lower()
+        output_path = pathlib.Path(content_file_name)
+    else:
+        output_path = pathlib.Path(args.output)
+
+    # Ensure that a version was supplied before downloading, because we need the matching TMD for decryption to work.
+    if decrypt_content is True and version is None:
+        fatal_error("You must specify the version that the requested content belongs to for decryption!")
+
+    # Try to download the content, and catch the ValueError libWiiPy will throw if it can't be found.
+    print(f"Downloading content with Content ID {cid}...")
+    content_data = None
+    try:
+        content_data = libWiiPy.title.download_content(tid, content_id)
+    except ValueError:
+        fatal_error("The specified Title ID or Content ID could not be found!")
+
+    if decrypt_content is True:
+        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.
+        ticket = None
+        try:
+            ticket = libWiiPy.title.Ticket()
+            ticket.load(libWiiPy.title.download_ticket(tid, wiiu_endpoint=True))
+        except ValueError:
+            fatal_error("No Ticket is available! Content cannot be decrypted.")
+
+        content_hash = 'gggggggggggggggggggggggggggggggggggggggg'
+        content_size = 0
+        content_index = 0
+        for record in tmd.content_records:
+            if record.content_id == content_id:
+                content_hash = record.content_hash.decode()
+                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':
+            fatal_error("Content was not found in the TMD for the specified version! Content cannot be decrypted.")
+
+        # 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:
+            fatal_error("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))
+        output_path.write_bytes(content_dec)
+    else:
+        output_path.write_bytes(content_data)
+
+    print(f"Downloaded content with Content ID \"{cid}\"!")
+
+
 def handle_nus_title(args):
     title_version = None
     wad_file = None
     output_dir = None
     can_decrypt = False
     tid = args.tid
-    if args.wii:
-        wiiu_nus_enabled = False
-    else:
-        wiiu_nus_enabled = True
+    wiiu_nus_enabled = False if args.wii else True
+    endpoint_override = args.endpoint if args.endpoint else None
 
     # Check if --version was passed, because it'll be None if it wasn't.
     if args.version is not None:
@@ -53,9 +128,11 @@ def handle_nus_title(args):
     print(" - Downloading and parsing TMD...")
     # Download a specific TMD version if a version was specified, otherwise just download the latest TMD.
     if title_version is not None:
-        title.load_tmd(libWiiPy.title.download_tmd(tid, title_version, wiiu_endpoint=wiiu_nus_enabled))
+        title.load_tmd(libWiiPy.title.download_tmd(tid, title_version, wiiu_endpoint=wiiu_nus_enabled,
+                                                   endpoint_override=endpoint_override))
     else:
-        title.load_tmd(libWiiPy.title.download_tmd(tid, wiiu_endpoint=wiiu_nus_enabled))
+        title.load_tmd(libWiiPy.title.download_tmd(tid, wiiu_endpoint=wiiu_nus_enabled,
+                                                   endpoint_override=endpoint_override))
         title_version = title.tmd.title_version
     # Write out the TMD to a file.
     if output_dir is not None:
@@ -64,7 +141,8 @@ def handle_nus_title(args):
     # Download the ticket, if we can.
     print(" - Downloading and parsing Ticket...")
     try:
-        title.load_ticket(libWiiPy.title.download_ticket(tid, wiiu_endpoint=wiiu_nus_enabled))
+        title.load_ticket(libWiiPy.title.download_ticket(tid, wiiu_endpoint=wiiu_nus_enabled,
+                                                         endpoint_override=endpoint_override))
         can_decrypt = True
         if output_dir is not None:
             output_dir.joinpath("tik").write_bytes(title.ticket.dump())
@@ -87,7 +165,8 @@ def handle_nus_title(args):
               f"(Content ID: {title.tmd.content_records[content].content_id}, "
               f"Size: {title.tmd.content_records[content].content_size} bytes)...")
         content_list.append(libWiiPy.title.download_content(tid, title.tmd.content_records[content].content_id,
-                                                            wiiu_endpoint=wiiu_nus_enabled))
+                                                            wiiu_endpoint=wiiu_nus_enabled,
+                                                            endpoint_override=endpoint_override))
         print("   - Done!")
         # If we're supposed to be outputting to a folder, then write these files out.
         if output_dir is not None:
@@ -110,7 +189,8 @@ def handle_nus_title(args):
     if wad_file is not None:
         # Get the WAD certificate chain.
         print(" - Building certificate...")
-        title.load_cert_chain(libWiiPy.title.download_cert_chain(wiiu_endpoint=wiiu_nus_enabled))
+        title.load_cert_chain(libWiiPy.title.download_cert_chain(wiiu_endpoint=wiiu_nus_enabled,
+                                                                 endpoint_override=endpoint_override))
         # Ensure that the path ends in .wad, and add that if it doesn't.
         print("Packing WAD...")
         if wad_file.suffix != ".wad":
@@ -121,83 +201,6 @@ def handle_nus_title(args):
     print(f"Downloaded title with Title ID \"{args.tid}\"!")
 
 
-def handle_nus_content(args):
-    tid = args.tid
-    cid = args.cid
-    version = args.version
-    if args.decrypt:
-        decrypt_content = True
-    else:
-        decrypt_content = False
-
-    # 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.
-    content_id = None
-    try:
-        content_id = int.from_bytes(binascii.unhexlify(cid))
-    except binascii.Error:
-        fatal_error("The provided Content ID is invalid! The Content ID must be in the format \"000000xx\"!")
-
-    # Use the supplied output path if one was specified, otherwise generate one using the Content ID.
-    if args.output is None:
-        content_file_name = f"{content_id:08X}".lower()
-        output_path = pathlib.Path(content_file_name)
-    else:
-        output_path = pathlib.Path(args.output)
-
-    # Ensure that a version was supplied before downloading, because we need the matching TMD for decryption to work.
-    if decrypt_content is True and version is None:
-        fatal_error("You must specify the version that the requested content belongs to for decryption!")
-
-    # Try to download the content, and catch the ValueError libWiiPy will throw if it can't be found.
-    print(f"Downloading content with Content ID {cid}...")
-    content_data = None
-    try:
-        content_data = libWiiPy.title.download_content(tid, content_id)
-    except ValueError:
-        fatal_error("The specified Title ID or Content ID could not be found!")
-
-    if decrypt_content is True:
-        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.
-        ticket = None
-        try:
-            ticket = libWiiPy.title.Ticket()
-            ticket.load(libWiiPy.title.download_ticket(tid, wiiu_endpoint=True))
-        except ValueError:
-            fatal_error("No Ticket is available! Content cannot be decrypted.")
-
-        content_hash = 'gggggggggggggggggggggggggggggggggggggggg'
-        content_size = 0
-        content_index = 0
-        for record in tmd.content_records:
-            if record.content_id == content_id:
-                content_hash = record.content_hash.decode()
-                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':
-            fatal_error("Content was not found in the TMD for the specified version! Content cannot be decrypted.")
-
-        # 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:
-            fatal_error("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))
-        output_path.write_bytes(content_dec)
-    else:
-        output_path.write_bytes(content_data)
-
-    print(f"Downloaded content with Content ID \"{cid}\"!")
-
-
 def handle_nus_tmd(args):
     tid = args.tid
 
diff --git a/commands/title/wad.py b/commands/title/wad.py
@@ -81,7 +81,6 @@ def handle_wad_convert(args):
     else:
         fatal_error("No valid encryption target was specified!")
 
-    output_path = pathlib.Path(args.output)
     if args.output is None:
         match target:
             case "development":
@@ -92,6 +91,8 @@ def handle_wad_convert(args):
                 output_path = pathlib.Path(input_path.stem + "_vWii" + input_path.suffix)
             case _:
                 fatal_error("No valid encryption target was specified!")
+    else:
+        output_path = pathlib.Path(args.output)
 
     if not input_path.exists():
         fatal_error(f"The specified WAD file \"{input_path}\" does not exist!")
diff --git a/wiipy.py b/wiipy.py
@@ -186,8 +186,10 @@
                                      help="download the title to a folder")
     nus_title_out_group.add_argument("-w", "--wad", metavar="WAD", type=str,
                                      help="pack a wad with the provided name")
-    nus_title_parser.add_argument("--wii", help="use original Wii NUS instead of the Wii U servers",
+    nus_title_parser.add_argument("--wii", help="use the original Wii NUS endpoint instead of the Wii U endpoint",
                                   action="store_true")
+    nus_title_parser.add_argument("-e", "--endpoint", metavar="ENDPOINT", type=str,
+                                  help="use the specified NUS endpoint instead of the official one")
     # Content NUS subcommand.
     nus_content_parser = nus_subparsers.add_parser("content", help="download a specific content from the NUS",
                                                    description="download a specific content from the NUS")
