From df7c8361fe19ea4a9013fff77c540df6af901fc5 Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Mon, 24 Jun 2024 22:50:03 -0400 Subject: [PATCH] Improved wad packing code and added lots of needed comments --- modules/ash.py | 5 +++++ modules/nus.py | 3 +++ modules/u8.py | 4 ++++ modules/wad.py | 58 ++++++++++++++++++++++++++++++++++++-------------- wiipy.py | 7 +++++- 5 files changed, 60 insertions(+), 17 deletions(-) diff --git a/modules/ash.py b/modules/ash.py index d79fc7e..ea57f3d 100644 --- a/modules/ash.py +++ b/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) diff --git a/modules/nus.py b/modules/nus.py index 39180c4..ee6dc31 100644 --- a/modules/nus.py +++ b/modules/nus.py @@ -6,6 +6,8 @@ import libWiiPy def handle_nus(args): title_version = None + + # 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,6 +15,7 @@ def handle_nus(args): print("Enter a valid integer for the Title Version.") return + # 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) file_name = args.tid + "-v" + str(title.tmd.title_version) + ".wad" diff --git a/modules/u8.py b/modules/u8.py index 1a24cc1..68f6267 100644 --- a/modules/u8.py +++ b/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,12 +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() + # Ensure the output directory doesn't already exist, because libWiiPy wants to create a new one to ensure that + # the contents of the U8 archive are extracted correctly. if output_path.exists(): print("Error: Specified output directory already exists!") return diff --git a/modules/wad.py b/modules/wad.py index 95c8f47..018efc1 100644 --- a/modules/wad.py +++ b/modules/wad.py @@ -10,48 +10,72 @@ 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() - - 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 +85,14 @@ 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() + # 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()) diff --git a/wiipy.py b/wiipy.py index 156ddd4..2d4b375 100644 --- a/wiipy.py +++ b/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\')") 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) @@ -25,6 +27,7 @@ if __name__ == "__main__": wad_parser.add_argument("input", metavar="IN", type=str, help="input file") wad_parser.add_argument("output", metavar="OUT", type=str, help="output file") + # 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) @@ -32,6 +35,7 @@ if __name__ == "__main__": nus_parser.add_argument("-v", "--version", metavar="VERSION", type=int, help="version to download (optional)") + # 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 +45,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 +59,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)