Removed fakesigning code, as it is now part of libWiiPy instead

New command uses its own subcommands. Currently only offers "title", which supports downloading a WAD using --wad or downloading to a folder using --output. More subcommands will be added.
Verbose output is also the default now and --verbose has been removed.

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,25 @@
 # "nus.py" from WiiPy by NinjaCheetah
 # https://github.com/NinjaCheetah/WiiPy
 
+import os
+import hashlib
+import pathlib
+import binascii
 import libWiiPy
 
 
-def handle_nus(args):
+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
+
+    # 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,12 +27,198 @@ def handle_nus(args):
             print("Enter a valid integer for the Title Version.")
             return
 
-    title = libWiiPy.title.download_title(args.tid, title_version)
+    # If --wad was passed, check to make sure the path is okay.
+    if args.wad is not None:
+        wad_file = pathlib.Path(args.wad)
+        if wad_file.suffix != ".wad":
+            wad_file = wad_file.with_suffix(".wad")
 
-    file_name = args.tid + "-v" + str(title.tmd.title_version) + ".wad"
+    # If --output was passed, make sure the directory either doesn't exist or is empty.
+    if args.output is not None:
+        output_dir = pathlib.Path(args.output)
+        if output_dir.exists():
+            if output_dir.is_dir() and next(os.scandir(output_dir), None):
+                raise ValueError("Output folder is not empty!")
+            elif output_dir.is_file():
+                raise ValueError("A file already exists with the provided directory name!")
+        else:
+            os.mkdir(output_dir)
 
-    wad_file = open(file_name, "wb")
-    wad_file.write(title.dump_wad())
-    wad_file.close()
+    # Download the title from the NUS. This is done "manually" (as opposed to using download_title()) so that we can
+    # provide verbose output.
+    title = libWiiPy.title.Title()
+
+    # Announce the title being downloaded, and the version if applicable.
+    if title_version is not None:
+        print("Downloading title " + tid + " v" + str(title_version) + ", please wait...")
+    else:
+        print("Downloading title " + tid + " vLatest, please wait...")
+    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))
+    else:
+        title.load_tmd(libWiiPy.title.download_tmd(tid, wiiu_endpoint=wiiu_nus_enabled))
+        title_version = title.tmd.title_version
+    # Write out the TMD to a file.
+    if output_dir is not None:
+        tmd_out = open(output_dir.joinpath("tmd." + str(title_version)), "wb")
+        tmd_out.write(title.tmd.dump())
+        tmd_out.close()
+
+    # 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))
+        can_decrypt = True
+        if output_dir is not None:
+            ticket_out = open(output_dir.joinpath("tik"), "wb")
+            ticket_out.write(title.ticket.dump())
+            ticket_out.close()
+    except ValueError:
+        # If libWiiPy returns an error, then no ticket is available. Log this, and disable options requiring a
+        # ticket so that they aren't attempted later.
+        print("  - No Ticket is available!")
+        if wad_file is not None and output_dir is None:
+            print("--wad was passed, but this title cannot be packed into a WAD!")
+            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)):
+        # Generate the content file name by converting the Content ID to hex and then removing the 0x.
+        content_file_name = hex(title.tmd.content_records[content].content_id)[2:]
+        while len(content_file_name) < 8:
+            content_file_name = "0" + content_file_name
+        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=wiiu_nus_enabled))
+        print("   - Done!")
+        # If we're supposed to be outputting to a folder, then write these files out.
+        if output_dir is not None:
+            enc_content_out = open(output_dir.joinpath(content_file_name), "wb")
+            enc_content_out.write(content_list[content])
+            enc_content_out.close()
+    title.content.content_list = content_list
+
+    # Try to decrypt the contents for this title if a ticket was available.
+    if can_decrypt is True and output_dir is not None:
+        for content in range(len(title.tmd.content_records)):
+            print(" - Decrypting content " + str(content + 1) + " of " + str(len(title.tmd.content_records)) +
+                  " (Content ID: " + str(title.tmd.content_records[content].content_id) + ")...")
+            dec_content = title.get_content_by_index(content)
+            content_file_name = hex(title.tmd.content_records[content].content_id)[2:]
+            while len(content_file_name) < 8:
+                content_file_name = "0" + content_file_name
+            content_file_name = content_file_name + ".app"
+            dec_content_out = open(output_dir.joinpath(content_file_name), "wb")
+            dec_content_out.write(dec_content)
+            dec_content_out.close()
+    else:
+        print("Title has no Ticket, so content will not be decrypted!")
+
+    # If --wad was passed, pack a WAD and output that.
+    if wad_file is not None:
+        # Get the WAD certificate chain.
+        print(" - Building certificate...")
+        title.wad.set_cert_data(libWiiPy.title.download_cert(wiiu_endpoint=wiiu_nus_enabled))
+        # Ensure that the path ends in .wad, and add that if it doesn't.
+        print("Packing WAD...")
+        if wad_file.suffix != ".wad":
+            wad_file = wad_file.with_suffix(".wad")
+        # Have libWiiPy dump the WAD, and write that data out.
+        file = open(wad_file, "wb")
+        file.write(title.dump_wad())
+        file.close()
 
     print("Downloaded title with Title ID \"" + args.tid + "\"!")
+
+
+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
+
+    # 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
+
+    # 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)
+
+    # 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:
+        # 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, wiiu_endpoint=True))
+        except ValueError:
+            print("No Ticket is available! Content cannot be decrypted!")
+            return
+
+        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':
+            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(output_path, "wb")
+        file.write(content_dec)
+        file.close()
+    else:
+        file = open(output_path, "wb")
+        file.write(content_data)
+        file.close()
+
+    print("Downloaded content with Content ID \"" + cid + "\"!")

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():
-            print("Error: Specified output directory already exists!")
-            return
-
+        # Output path is deliberately not checked in any way because libWiiPy already has those checks, and it's easier
+        # and cleaner to only have one component doing all the checks.
         libWiiPy.archive.extract_u8(u8_data, str(output_path))
 
         print("U8 archive unpacked!")

modules/wad.py

@@ -1,7 +1,9 @@
 # "wad.py" from WiiPy by NinjaCheetah
 # https://github.com/NinjaCheetah/WiiPy
 
+import os
 import pathlib
+import hashlib
 import binascii
 import libWiiPy
 
@@ -10,48 +12,76 @@ 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]
-        if not tmd_file.exists():
-            raise FileNotFoundError("Cannot find a TMD! Exiting...")
+        # Get a list of all files ending in .tmd, and then make sure that that list has *only* 1 entry. More than 1
+        # means we can't pack a WAD because we couldn't really tell which TMD is intended for this WAD.
+        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]
-        if not ticket_file.exists():
-            raise FileNotFoundError("Cannot find a Ticket! Exiting...")
+        # Repeat the same process as above for all .tik files.
+        ticket_list = list(input_path.glob('*.tik'))
+        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]
-        if not cert_file.exists():
-            raise FileNotFoundError("Cannot find a cert! Exiting...")
+        # And one more time for all .cert files.
+        cert_list = list(input_path.glob('*.cert'))
+        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()
+            # Fakesign the TMD and Ticket using the trucha bug, if enabled. This is built-in in libWiiPy v0.4.1+.
+            if args.fakesign:
+                title.fakesign()
 
-            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)):
-                    dec_content = open(content_list[content], "rb").read()
+                for content in range(len(content_files)):
+                    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():
-            output_path.mkdir()
+        # Check if the output path already exists, and if it does, ensure that it is both a directory and empty.
+        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 directory 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.2.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,42 @@ 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("--fakesign", help="fakesign the TMD and Ticket (trucha bug)",
+                            action="store_true")
 
-    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)")
+    # Argument parser for the NUS subcommand.
+    nus_parser = subparsers.add_parser("nus", help="download data from the NUS",
+                                       description="download from the NUS")
+    nus_subparsers = nus_parser.add_subparsers(dest="subcommand", required=True)
+    # Title NUS subcommand.
+    nus_title_parser = nus_subparsers.add_parser("title", help="download a title from the NUS",
+                                                 description="download a title from the NUS")
+    nus_title_parser.set_defaults(func=handle_nus_title)
+    nus_title_parser.add_argument("tid", metavar="TID", type=str, help="Title ID to download")
+    nus_title_parser.add_argument("-v", "--version", metavar="VERSION", type=int,
+                                  help="version to download (optional)")
+    nus_title_out_group_label = nus_title_parser.add_argument_group(title="output types (required)")
+    nus_title_out_group = nus_title_out_group_label.add_mutually_exclusive_group(required=True)
+    nus_title_out_group.add_argument("-o", "--output", metavar="OUT", type=str,
+                                     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",
+                                  action="store_true")
+    # 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")
+    nus_content_parser.set_defaults(func=handle_nus_content)
+    nus_content_parser.add_argument("tid", metavar="TID", type=str, help="Title ID the content belongs to")
+    nus_content_parser.add_argument("cid", metavar="CID", type=str,
+                                    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.
     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 +71,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 +85,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)