From 5cff54592197d6696d69365cb3db42a5d6f1f52e Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Sat, 6 Jul 2024 20:39:52 +1000
Subject: [PATCH] Added --verbose and --wiiu to nus module --verbose prints
 NUSGet-style logging to the terminal during the download, and --wiiu enables
 the Wii U servers for faster downloads.

---
 modules/nus.py | 59 ++++++++++++++++++++++++++++++++++++++++++++++++--
 modules/wad.py |  5 +++++
 wiipy.py       |  6 +++++
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/modules/nus.py b/modules/nus.py
index 259e4bd..3acbba9 100644
--- a/modules/nus.py
+++ b/modules/nus.py
@@ -8,6 +8,15 @@ import libWiiPy
 def handle_nus(args):
     title_version = None
     file_path = None
+    tid = args.tid
+    if args.wiiu:
+        use_wiiu_servers = True
+    else:
+        use_wiiu_servers = False
+    if args.verbose:
+        verbose = True
+    else:
+        verbose = False
 
     # Check if --version was passed, because it'll be None if it wasn't.
     if args.version is not None:
@@ -26,8 +35,54 @@ def handle_nus(args):
         if file_path.suffix != ".wad":
             file_path = file_path.with_suffix(".wad")
 
-    # libWiiPy accepts a title version of "None" and will just use the latest available version if it gets it.
-    title = libWiiPy.title.download_title(args.tid, title_version)
+    # Download the title from the NUS. This is done "manually" (as opposed to using download_title()) so that we can
+    # provide verbose output if desired.
+    title = libWiiPy.title.Title()
+
+    # Announce the title being downloaded, and the version if applicable.
+    if verbose:
+        if title_version is not None:
+            print("Downloading title " + tid + " v" + str(title_version) + ", please wait...")
+        else:
+            print("Downloading title " + tid + " vLatest, please wait...")
+
+    # Download a specific TMD version if a version was specified, otherwise just download the latest TMD.
+    if verbose:
+        print(" - Downloading and parsing TMD...")
+    if title_version is not None:
+        title.load_tmd(libWiiPy.title.download_tmd(tid, title_version, wiiu_endpoint=use_wiiu_servers))
+    else:
+        title.load_tmd(libWiiPy.title.download_tmd(tid, wiiu_endpoint=use_wiiu_servers))
+
+    # Download and parse the Ticket.
+    if verbose:
+        print(" - Downloading and parsing Ticket...")
+    try:
+        title.load_ticket(libWiiPy.title.download_ticket(tid, wiiu_endpoint=use_wiiu_servers))
+    except ValueError:
+        # If libWiiPy returns an error, then no ticket is available, so we can't continue.
+        print("No Ticket is available for this title! Exiting...")
+        return
+
+    # Load the content records from the TMD, and begin iterating over the records.
+    title.load_content_records()
+    content_list = []
+    for content in range(len(title.tmd.content_records)):
+        if verbose:
+            print(" - Downloading content " + str(content + 1) + " of " +
+                  str(len(title.tmd.content_records)) + " (Content ID: " +
+                  str(title.tmd.content_records[content].content_id) + ", Size: " +
+                  str(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=use_wiiu_servers))
+        if verbose:
+            print("   - Done!")
+    title.content.content_list = content_list
+
+    # Get the WAD certificate chain.
+    if verbose:
+        print(" - Building certificate...")
+    title.wad.set_cert_data(libWiiPy.title.download_cert(wiiu_endpoint=use_wiiu_servers))
 
     # If we haven't gotten a name yet, make one from the TID and version.
     if file_path is None:
diff --git a/modules/wad.py b/modules/wad.py
index 4ed4bd2..f98d3b5 100644
--- a/modules/wad.py
+++ b/modules/wad.py
@@ -70,6 +70,11 @@ def handle_wad(args):
             # Method to ensure that the title's content records match between the TMD() and ContentRegion() objects.
             title.load_content_records()
 
+            # Nullify TMD/Ticket signatures here if the argument was passed.
+            if args.null_sigs:
+                title.tmd.signature = b'\x00' * 256
+                title.ticket.signature = b'\x00' * 256
+
             # Iterate over every file in the content_files list, and attempt to load it into the Title().
             for index in range(len(title.content.content_records)):
                 for content in range(len(content_files)):
diff --git a/wiipy.py b/wiipy.py
index 433ffa7..d223053 100644
--- a/wiipy.py
+++ b/wiipy.py
@@ -26,6 +26,8 @@ if __name__ == "__main__":
     wad_group.add_argument("-u", "--unpack", help="unpack a WAD file to a directory", action="store_true")
     wad_parser.add_argument("input", metavar="IN", type=str, help="input file")
     wad_parser.add_argument("output", metavar="OUT", type=str, help="output file")
+    wad_parser.add_argument("--null-sigs", help="nullify signatures in the TMD and Ticket (packing only)",
+                            action="store_true")
 
     # Argument parser for the NUS subcommand.
     nus_parser = subparsers.add_parser("nus", help="download a title from the NUS",
@@ -35,6 +37,10 @@ if __name__ == "__main__":
     nus_parser.add_argument("-v", "--version", metavar="VERSION", type=int,
                             help="version to download (optional)")
     nus_parser.add_argument("-o", "--output", metavar="OUT", type=str, help="output file (optional)")
+    nus_parser.add_argument("--verbose", help="output more information about the current download",
+                            action="store_true")
+    nus_parser.add_argument("-w", "--wiiu", help="use Wii U servers for faster downloads",
+                            action="store_true")
 
     # Argument parser for the U8 subcommand.
     u8_parser = subparsers.add_parser("u8", help="pack/unpack a U8 archive",