From 5e1bf6ed4eb2b539bf757547693c8e9d069cfe30 Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Mon, 24 Jun 2024 00:41:07 -0400 Subject: [PATCH] Rewrote entire source to be based around argparse This does change a large amount of the syntax for using the CLI, but I think that's for the better. This new system also allows for having help pages for each sub command, making the tool a lot easier to use. It also allows for having more arguments available for each subcommand, which is especially necessary for the ASH module. --- modules/ash.py | 36 +++++++----- modules/nus.py | 12 ++-- modules/u8.py | 44 ++++++++------ modules/wad.py | 154 +++++++++++++++++++++++-------------------------- wiipy.py | 98 +++++++++++++++---------------- 5 files changed, 175 insertions(+), 169 deletions(-) diff --git a/modules/ash.py b/modules/ash.py index 85a055e..d79fc7e 100644 --- a/modules/ash.py +++ b/modules/ash.py @@ -1,23 +1,33 @@ # "ash.py" from WiiPy by NinjaCheetah # https://github.com/NinjaCheetah/WiiPy -import os +import pathlib import libWiiPy -def decompress_ash(in_file: str, out_file: str = None): - if not os.path.isfile(in_file): - raise FileNotFoundError(in_file) +def handle_ash(args): + input_path = pathlib.Path(args.input) + output_path = pathlib.Path(args.output) - ash_file = open(in_file, "rb") - ash_data = ash_file.read() - ash_file.close() + if args.compress: + print("Compression is not implemented yet.") - ash_decompressed = libWiiPy.archive.decompress_ash(ash_data) + elif args.decompress: + sym_tree_bits = args.sym_bits + dist_tree_bits = args.dist_bits - if out_file is None: - out_file = in_file + ".arc" + if not input_path.exists(): + raise FileNotFoundError(input_path) - ash_out = open(out_file, "wb") - ash_out.write(ash_decompressed) - ash_out.close() + ash_file = open(input_path, "rb") + ash_data = ash_file.read() + ash_file.close() + + ash_decompressed = libWiiPy.archive.decompress_ash(ash_data, sym_tree_bits=sym_tree_bits, + dist_tree_bits=dist_tree_bits) + + ash_out = open(output_path, "wb") + ash_out.write(ash_decompressed) + ash_out.close() + + print("ASH file decompressed!") diff --git a/modules/nus.py b/modules/nus.py index 9fd9cbd..39180c4 100644 --- a/modules/nus.py +++ b/modules/nus.py @@ -4,19 +4,21 @@ import libWiiPy -def download_title(title_id: str, title_version_input: str = None): +def handle_nus(args): title_version = None - if title_version_input is not None: + if args.version is not None: try: - title_version = int(title_version_input) + title_version = int(args.version) except ValueError: print("Enter a valid integer for the Title Version.") return - title = libWiiPy.title.download_title(title_id, title_version) + title = libWiiPy.title.download_title(args.tid, title_version) - file_name = title_id + "-v" + str(title.tmd.title_version) + ".wad" + file_name = args.tid + "-v" + str(title.tmd.title_version) + ".wad" wad_file = open(file_name, "wb") wad_file.write(title.dump_wad()) wad_file.close() + + print("Downloaded title with Title ID \"" + args.tid + "\"!") diff --git a/modules/u8.py b/modules/u8.py index fe89ea2..1a24cc1 100644 --- a/modules/u8.py +++ b/modules/u8.py @@ -1,29 +1,37 @@ # "u8.py" from WiiPy by NinjaCheetah # https://github.com/NinjaCheetah/WiiPy -import os +import pathlib import libWiiPy -def extract_u8_to_folder(in_file: str, out_folder: str): - if not os.path.isfile(in_file): - raise FileNotFoundError(in_file) +def handle_u8(args): + input_path = pathlib.Path(args.input) + output_path = pathlib.Path(args.output) - u8_data = open(in_file, "rb").read() + if args.pack: + try: + u8_data = libWiiPy.archive.pack_u8(input_path) + except ValueError: + print("Error: Specified input file/folder does not exist!") + return - try: - libWiiPy.archive.extract_u8(u8_data, out_folder) - except ValueError: - print("Specified output folder already exists!") + out_file = open(output_path, "wb") + out_file.write(u8_data) + out_file.close() + print("U8 archive packed!") -def pack_u8_from_folder(in_folder: str, out_file: str): - try: - u8_data = libWiiPy.archive.pack_u8(in_folder) - except ValueError: - print("Specified input file/folder does not exist!") - return + elif args.unpack: + if not input_path.exists(): + raise FileNotFoundError(args.input) - out_file = open(out_file, "wb") - out_file.write(u8_data) - out_file.close() + u8_data = open(input_path, "rb").read() + + if output_path.exists(): + print("Error: Specified output directory already exists!") + return + + libWiiPy.archive.extract_u8(u8_data, str(output_path)) + + print("U8 archive unpacked!") diff --git a/modules/wad.py b/modules/wad.py index 903967e..95c8f47 100644 --- a/modules/wad.py +++ b/modules/wad.py @@ -1,108 +1,100 @@ # "wad.py" from WiiPy by NinjaCheetah # https://github.com/NinjaCheetah/WiiPy -import os import pathlib import binascii import libWiiPy -def extract_wad_to_folder(in_file: str, out_folder: str): - if not os.path.isfile(in_file): - raise FileNotFoundError(in_file) - out_folder = pathlib.Path(out_folder) - if not out_folder.is_dir(): - out_folder.mkdir() +def handle_wad(args): + input_path = pathlib.Path(args.input) + output_path = pathlib.Path(args.output) - with open(in_file, "rb") as wad_file: - wad = libWiiPy.title.WAD() - wad.load(wad_file.read()) + if args.pack: + if not input_path.exists(): + raise FileNotFoundError(input_path) + if not input_path.is_dir(): + raise NotADirectoryError(input_path) - tmd = libWiiPy.title.TMD() - tmd.load(wad.get_tmd_data()) - ticket = libWiiPy.title.Ticket() - ticket.load(wad.get_ticket_data()) - content_region = libWiiPy.title.ContentRegion() - content_region.load(wad.get_content_data(), tmd.content_records) + tmd_file = list(input_path.glob("*.tmd"))[0] + if not tmd_file.exists(): + raise FileNotFoundError("Cannot find a TMD! Exiting...") - title_key = ticket.get_title_key() + ticket_file = list(input_path.glob("*.tik"))[0] + if not ticket_file.exists(): + raise FileNotFoundError("Cannot find a Ticket! Exiting...") - cert_name = tmd.title_id + ".cert" - cert_out = open(os.path.join(out_folder, cert_name), "wb") - cert_out.write(wad.get_cert_data()) - cert_out.close() + cert_file = list(input_path.glob("*.cert"))[0] + if not cert_file.exists(): + raise FileNotFoundError("Cannot find a cert! Exiting...") - tmd_name = tmd.title_id + ".tmd" - tmd_out = open(os.path.join(out_folder, tmd_name), "wb") - tmd_out.write(wad.get_tmd_data()) - tmd_out.close() + content_files = list(input_path.glob("*.app")) + if not content_files: + raise FileNotFoundError("Cannot find any contents! Exiting...") - ticket_name = tmd.title_id + ".tik" - ticket_out = open(os.path.join(out_folder, ticket_name), "wb") - ticket_out.write(wad.get_ticket_data()) - ticket_out.close() + with open(output_path, "wb") as output_path: + title = libWiiPy.title.Title() - meta_name = tmd.title_id + ".footer" - meta_out = open(os.path.join(out_folder, meta_name), "wb") - meta_out.write(wad.get_meta_data()) - meta_out.close() + 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()) + footer_file = list(input_path.glob("*.footer"))[0] + if footer_file.exists(): + title.wad.set_meta_data(open(footer_file, "rb").read()) + title.load_content_records() - for content_file in range(0, tmd.num_contents): - content_file_name = "000000" - content_file_name += str(binascii.hexlify(content_file.to_bytes()).decode()) + ".app" - content_out = open(os.path.join(out_folder, content_file_name), "wb") - content_out.write(content_region.get_content_by_index(content_file, title_key)) - content_out.close() + title_key = title.ticket.get_title_key() + content_list = list(input_path.glob("*.app")) + for index in range(len(title.content.content_records)): + for content in range(len(content_list)): + dec_content = open(content_list[content], "rb").read() + try: + # Attempt to load the content into the correct index. + title.content.load_content(dec_content, index, title_key) + break + except ValueError: + # Wasn't the right content, so try again. + pass -def pack_wad_from_folder(in_folder, out_file): - if not os.path.exists(in_folder): - raise FileNotFoundError(in_folder) - if not os.path.isdir(in_folder): - raise NotADirectoryError(in_folder) + output_path.write(title.dump_wad()) - out_file = pathlib.Path(out_file) - in_folder = pathlib.Path(in_folder) + print("WAD file packed!") - tmd_file = list(in_folder.glob("*.tmd"))[0] - if not os.path.exists(tmd_file): - raise FileNotFoundError("Cannot find a TMD! Exiting...") + elif args.unpack: + if not input_path.exists(): + raise FileNotFoundError(input_path) + if not output_path.is_dir(): + output_path.mkdir() - ticket_file = list(in_folder.glob("*.tik"))[0] - if not os.path.exists(ticket_file): - raise FileNotFoundError("Cannot find a Ticket! Exiting...") + with open(args.input, "rb") as wad_file: + title = libWiiPy.title.Title() + title.load_wad(wad_file.read()) - cert_file = list(in_folder.glob("*.cert"))[0] - if not os.path.exists(cert_file): - raise FileNotFoundError("Cannot find a cert! Exiting...") + cert_name = title.tmd.title_id + ".cert" + cert_out = open(output_path.joinpath(cert_name), "wb") + cert_out.write(title.wad.get_cert_data()) + cert_out.close() - content_files = list(in_folder.glob("*.app")) - if not content_files: - raise FileNotFoundError("Cannot find any contents! Exiting...") + tmd_name = title.tmd.title_id + ".tmd" + tmd_out = open(output_path.joinpath(tmd_name), "wb") + tmd_out.write(title.wad.get_tmd_data()) + tmd_out.close() - with open(out_file, "wb") as out_file: - title = libWiiPy.title.Title() + ticket_name = title.tmd.title_id + ".tik" + ticket_out = open(output_path.joinpath(ticket_name), "wb") + ticket_out.write(title.wad.get_ticket_data()) + ticket_out.close() - 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()) - footer_file = list(in_folder.glob("*.footer"))[0] - if os.path.exists(footer_file): - title.wad.set_meta_data(open(footer_file, "rb").read()) - title.load_content_records() + meta_name = title.tmd.title_id + ".footer" + meta_out = open(output_path.joinpath(meta_name), "wb") + meta_out.write(title.wad.get_meta_data()) + meta_out.close() - title_key = title.ticket.get_title_key() + for content_file in range(0, title.tmd.num_contents): + content_file_name = "000000" + str(binascii.hexlify(content_file.to_bytes()).decode()) + ".app" + content_out = open(output_path.joinpath(content_file_name), "wb") + content_out.write(title.get_content_by_index(content_file)) + content_out.close() - content_list = list(in_folder.glob("*.app")) - for index in range(len(title.content.content_records)): - for content in range(len(content_list)): - dec_content = open(content_list[content], "rb").read() - try: - # Attempt to load the content into the correct index. - title.content.load_content(dec_content, index, title_key) - break - except ValueError: - # Wasn't the right content, so try again. - pass - - out_file.write(title.dump_wad()) + print("WAD file unpacked!") diff --git a/wiipy.py b/wiipy.py index 3503493..19f76cf 100644 --- a/wiipy.py +++ b/wiipy.py @@ -1,61 +1,55 @@ # "wiipy.py" from WiiPy by NinjaCheetah # https://github.com/NinjaCheetah/WiiPy -import sys +import argparse + from modules.wad import * from modules.nus import * from modules.u8 import * from modules.ash import * -opts = [opt for opt in sys.argv[1:] if opt.startswith("-")] -args = [arg for arg in sys.argv[1:] if not arg.startswith("-")] - if __name__ == "__main__": - if "WAD" in args: - if "-u" in opts: - if len(args) == 3: - extract_wad_to_folder(args[1], args[2]) - print("Success!") - sys.exit(0) - if "-p" in opts: - if len(args) == 3: - pack_wad_from_folder(args[1], args[2]) - print("Success!") - sys.exit(0) - raise SystemExit(f"Usage: {sys.argv[0]} WAD (-u | -p) ") - elif "NUS" in args: - if "-d" in opts: - if len(args) == 2: - download_title(args[1]) - print("Success!") - sys.exit(0) - elif len(args) == 3: - download_title(args[1], args[2]) - print("Success!") - sys.exit(0) - raise SystemExit(f"Usage: {sys.argv[0]} NUS -d <Title Version (Optional)>") - elif "U8" in args: - if "-u" in opts: - if len(args) == 3: - extract_u8_to_folder(args[1], args[2]) - print("Success!") - sys.exit(0) - elif "-p" in opts: - if len(args) == 3: - pack_u8_from_folder(args[1], args[2]) - print("Success!") - sys.exit(0) - raise SystemExit(f"Usage: {sys.argv[0]} U8 (-u | -p) <input> <output>") - elif "ASH" in args: - if "-d" in opts: - if len(args) == 2: - decompress_ash(args[1]) - print("Success!") - sys.exit(0) - elif len(args) == 3: - decompress_ash(args[1], args[2]) - print("Success!") - sys.exit(0) - raise SystemExit(f"Usage: {sys.argv[0]} ASH -d <input> <output (Optional)>") - else: - raise SystemExit(f"Usage: {sys.argv[0]} WAD (-u | -p) <input> <output>") + parser = argparse.ArgumentParser(description="WiiPy is a simple CLI tool to manage file formats used by the Wii.") + subparsers = parser.add_subparsers(dest="subcommand", required=True) + + 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) + wad_group = wad_parser.add_mutually_exclusive_group(required=True) + wad_group.add_argument("-p", "--pack", help="pack a directory to a WAD file", action="store_true") + 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") + + 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") + + 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) + u8_group = u8_parser.add_mutually_exclusive_group(required=True) + u8_group.add_argument("-p", "--pack", help="pack a directory to a U8 archive", action="store_true") + u8_group.add_argument("-u", "--unpack", help="unpack a U8 archive to a directory", action="store_true") + u8_parser.add_argument("input", metavar="IN", type=str, help="input file") + u8_parser.add_argument("output", metavar="OUT", type=str, help="output file") + + 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) + ash_group = ash_parser.add_mutually_exclusive_group(required=True) + ash_group.add_argument("-c", "--compress", help="compress a file into an ASH file", action="store_true") + ash_group.add_argument("-d", "--decompress", help="decompress an ASH file", action="store_true") + ash_parser.add_argument("input", metavar="IN", type=str, help="input file") + ash_parser.add_argument("output", metavar="OUT", type=str, help="output file") + ash_parser.add_argument("--sym-bits", metavar="SYM_BITS", type=int, + help="number of bits in each symbol tree leaf (default: 9)", default=9) + ash_parser.add_argument("--dist-bits", metavar="DIST_BITS", type=int, + help="number of bits in each distance tree leaf (default: 11)", default=11) + + args = parser.parse_args() + + args.func(args)