Change --wiiu to --wii and made Wii U servers the default

--verbose prints NUSGet-style logging to the terminal during the download, and --wiiu enables the Wii U servers for faster downloads.

modules/ash.py

@@ -9,10 +9,14 @@ def handle_ash(args):
     input_path = pathlib.Path(args.input)
     output_path = pathlib.Path(args.output)
 
+    # Code for if --compress was passed.
+    # ASH compression has not been implemented in libWiiPy yet, but it'll be filled in here when it has.
     if args.compress:
         print("Compression is not implemented yet.")
 
+    # Code for if --decompress was passed.
     elif args.decompress:
+        # These default to 9 and 11, respectively, so we can always read them.
         sym_tree_bits = args.sym_bits
         dist_tree_bits = args.dist_bits
 
@@ -23,6 +27,7 @@ def handle_ash(args):
         ash_data = ash_file.read()
         ash_file.close()
 
+        # Decompress ASH file using the provided symbol/distance tree widths.
         ash_decompressed = libWiiPy.archive.decompress_ash(ash_data, sym_tree_bits=sym_tree_bits,
                                                            dist_tree_bits=dist_tree_bits)
 

modules/nus.py

@@ -1,11 +1,24 @@
 # "nus.py" from WiiPy by NinjaCheetah
 # https://github.com/NinjaCheetah/WiiPy
 
+import pathlib
 import libWiiPy
 
 
 def handle_nus(args):
     title_version = None
+    file_path = None
+    tid = args.tid
+    if args.wii:
+        use_wiiu_servers = False
+    else:
+        use_wiiu_servers = True
+    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:
         try:
             title_version = int(args.version)
@@ -13,11 +26,69 @@ def handle_nus(args):
             print("Enter a valid integer for the Title Version.")
             return
 
-    title = libWiiPy.title.download_title(args.tid, title_version)
+    # If --output was passed, then save the file to the specified path (as long as it's valid).
+    if args.output is not None:
+        file_path = pathlib.Path(args.output)
+        if not file_path.parent.exists() or not file_path.parent.is_dir():
+            print("The specified output path does not exist!")
+            return
+        if file_path.suffix != ".wad":
+            file_path = file_path.with_suffix(".wad")
 
-    file_name = args.tid + "-v" + str(title.tmd.title_version) + ".wad"
+    # 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()
 
-    wad_file = open(file_name, "wb")
+    # 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:
+        file_path = pathlib.Path(args.tid + "-v" + str(title.tmd.title_version) + ".wad")
+
+    wad_file = open(file_path, "wb")
     wad_file.write(title.dump_wad())
     wad_file.close()
 

modules/u8.py

@@ -9,6 +9,7 @@ def handle_u8(args):
     input_path = pathlib.Path(args.input)
     output_path = pathlib.Path(args.output)
 
+    # Code for if the --pack argument was passed.
     if args.pack:
         try:
             u8_data = libWiiPy.archive.pack_u8(input_path)
@@ -22,16 +23,15 @@ def handle_u8(args):
 
         print("U8 archive packed!")
 
+    # Code for if the --unpack argument was passed.
     elif args.unpack:
         if not input_path.exists():
             raise FileNotFoundError(args.input)
 
         u8_data = open(input_path, "rb").read()
 
-        if output_path.exists():
+        # Output path is deliberately not checked in any way because libWiiPy already has those checks, and it's easier
-            print("Error: Specified output directory already exists!")
+        # and cleaner to only have one component doing all the checks.
-            return
-
         libWiiPy.archive.extract_u8(u8_data, str(output_path))
 
         print("U8 archive unpacked!")

modules/wad.py

@@ -1,6 +1,7 @@
 # "wad.py" from WiiPy by NinjaCheetah
 # https://github.com/NinjaCheetah/WiiPy
 
+import os
 import pathlib
 import binascii
 import libWiiPy
@@ -10,48 +11,77 @@ def handle_wad(args):
     input_path = pathlib.Path(args.input)
     output_path = pathlib.Path(args.output)
 
+    # Code for if the --pack argument was passed.
     if args.pack:
+        # Make sure input path both exists and is a directory. Separate checks because this provides more relevant
+        # errors than just a NotADirectoryError if the actual issue is that there's nothing at all.
         if not input_path.exists():
             raise FileNotFoundError(input_path)
         if not input_path.is_dir():
             raise NotADirectoryError(input_path)
 
-        tmd_file = list(input_path.glob("*.tmd"))[0]
+        # Get a list of all files ending in .tmd, and then make sure that that list has *only* 1 entry. More than 1
-        if not tmd_file.exists():
+        # means we can't pack a WAD because we couldn't really tell which TMD is intended for this WAD.
-            raise FileNotFoundError("Cannot find a TMD! Exiting...")
+        tmd_list = list(input_path.glob('*.tmd'))
+        if len(tmd_list) > 1:
+            raise FileExistsError("More than one TMD file was found! Only one TMD can be packed into a WAD.")
+        elif len(tmd_list) == 0:
+            raise FileNotFoundError("No TMD file found! Cannot pack WAD.")
+        else:
+            tmd_file = tmd_list[0]
 
-        ticket_file = list(input_path.glob("*.tik"))[0]
+        # Repeat the same process as above for all .tik files.
-        if not ticket_file.exists():
+        ticket_list = list(input_path.glob('*.tik'))
-            raise FileNotFoundError("Cannot find a Ticket! Exiting...")
+        if len(ticket_list) > 1:
+            raise FileExistsError("More than one Ticket file was found! Only one Ticket can be packed into a WAD.")
+        elif len(ticket_list) == 0:
+            raise FileNotFoundError("No Ticket file found! Cannot pack WAD.")
+        else:
+            ticket_file = ticket_list[0]
 
-        cert_file = list(input_path.glob("*.cert"))[0]
+        # And one more time for all .cert files.
-        if not cert_file.exists():
+        cert_list = list(input_path.glob('*.cert'))
-            raise FileNotFoundError("Cannot find a cert! Exiting...")
+        if len(cert_list) > 1:
+            raise FileExistsError("More than one certificate file was found! Only one certificate can be packed into a "
+                                  "WAD.")
+        elif len(cert_list) == 0:
+            raise FileNotFoundError("No certificate file found! Cannot pack WAD.")
+        else:
+            cert_file = cert_list[0]
 
+        # Make sure that there's at least one content to pack.
         content_files = list(input_path.glob("*.app"))
         if not content_files:
-            raise FileNotFoundError("Cannot find any contents! Exiting...")
+            raise FileNotFoundError("No contents found! Cannot pack WAD.")
 
+        # Open the output file, and load all the component files that we've now verified we have into a libWiiPy Title()
+        # object.
         with open(output_path, "wb") as output_path:
             title = libWiiPy.title.Title()
 
             title.load_tmd(open(tmd_file, "rb").read())
             title.load_ticket(open(ticket_file, "rb").read())
             title.wad.set_cert_data(open(cert_file, "rb").read())
+            # Footers are not super common and are not required, so we don't care about one existing until we get to
+            # the step where we'd pack it.
             footer_file = list(input_path.glob("*.footer"))[0]
             if footer_file.exists():
                 title.wad.set_meta_data(open(footer_file, "rb").read())
+            # Method to ensure that the title's content records match between the TMD() and ContentRegion() objects.
             title.load_content_records()
 
-            title_key = title.ticket.get_title_key()
+            # 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
 
-            content_list = list(input_path.glob("*.app"))
+            # 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_list)):
+                for content in range(len(content_files)):
-                    dec_content = open(content_list[content], "rb").read()
+                    dec_content = open(content_files[content], "rb").read()
                     try:
                         # Attempt to load the content into the correct index.
-                        title.content.load_content(dec_content, index, title_key)
+                        title.load_content(dec_content, index)
                         break
                     except ValueError:
                         # Wasn't the right content, so try again.
@@ -61,12 +91,20 @@ def handle_wad(args):
 
         print("WAD file packed!")
 
+    # Code for if the --unpack argument was passed.
     elif args.unpack:
         if not input_path.exists():
             raise FileNotFoundError(input_path)
-        if not output_path.is_dir():
+        # Check if the output path already exists, and if it does, ensure that it is both a directory and empty.
-            output_path.mkdir()
+        if output_path.exists():
+            if output_path.is_dir() and next(os.scandir(output_path), None):
+                raise ValueError("Output folder is not empty!")
+            elif output_path.is_file():
+                raise ValueError("A file already exists with the provided name!")
+        else:
+            os.mkdir(output_path)
 
+        # Step through each component of a WAD and dump it to a file.
         with open(args.input, "rb") as wad_file:
             title = libWiiPy.title.Title()
             title.load_wad(wad_file.read())

wiipy.py

@@ -10,12 +10,14 @@ from modules.u8 import *
 from modules.ash import *
 
 if __name__ == "__main__":
+    # Main argument parser.
     parser = argparse.ArgumentParser(
         description="WiiPy is a simple command line tool to manage file formats used by the Wii.")
     parser.add_argument("--version", action="version",
-                        version=f"WiiPy v1.0.0, based on libWiiPy v{version('libWiiPy')} (from branch \'main\')")
+                        version=f"WiiPy v1.1.0, based on libWiiPy v{version('libWiiPy')} (from branch \'main\')")
     subparsers = parser.add_subparsers(dest="subcommand", required=True)
 
+    # Argument parser for the WAD subcommand.
     wad_parser = subparsers.add_parser("wad", help="pack/unpack a WAD file",
                                        description="pack/unpack a WAD file")
     wad_parser.set_defaults(func=handle_wad)
@@ -24,14 +26,23 @@ 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",
                                        description="download a title from the NUS")
     nus_parser.set_defaults(func=handle_nus)
     nus_parser.add_argument("tid", metavar="TID", type=str, help="Title ID to download")
     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", "--wii", help="use original Wii NUS instead of the Wii U servers",
+                            action="store_true")
 
+    # Argument parser for the U8 subcommand.
     u8_parser = subparsers.add_parser("u8", help="pack/unpack a U8 archive",
                                       description="pack/unpack a U8 archive")
     u8_parser.set_defaults(func=handle_u8)
@@ -41,6 +52,7 @@ if __name__ == "__main__":
     u8_parser.add_argument("input", metavar="IN", type=str, help="input file")
     u8_parser.add_argument("output", metavar="OUT", type=str, help="output file")
 
+    # Argument parser for the ASH subcommand.
     ash_parser = subparsers.add_parser("ash", help="compress/decompress an ASH file",
                                        description="compress/decompress an ASH file")
     ash_parser.set_defaults(func=handle_ash)
@@ -54,6 +66,6 @@ if __name__ == "__main__":
     ash_parser.add_argument("--dist-bits", metavar="DIST_BITS", type=int,
                             help="number of bits in each distance tree leaf (default: 11)", default=11)
 
+    # Parse all the args, and call the appropriate function with all of those args if a valid subcommand was passed.
     args = parser.parse_args()
-
     args.func(args)