diff --git a/modules/title/tmd.py b/modules/title/tmd.py new file mode 100644 index 0000000..e3be086 --- /dev/null +++ b/modules/title/tmd.py @@ -0,0 +1,47 @@ +# "modules/title/tmd.py" from WiiPy by NinjaCheetah +# https://github.com/NinjaCheetah/WiiPy + +import pathlib +import libWiiPy + + +def handle_tmd_remove(args): + input_path = pathlib.Path(args.input) + if args.output is not None: + output_path = pathlib.Path(args.output) + else: + output_path = pathlib.Path(args.input) + + if not input_path.exists(): + raise FileNotFoundError(input_path) + + tmd = libWiiPy.title.TMD() + tmd.load(input_path.read_bytes()) + + if args.index is not None: + # Make sure the target index exists, then remove it from the TMD. + if args.index >= len(tmd.content_records): + raise ValueError("The provided index could not be found in this TMD!") + tmd.content_records.pop(args.index) + tmd.num_contents -= 1 + # Auto fakesign because we've edited the TMD. + tmd.fakesign() + output_path.write_bytes(tmd.dump()) + print(f"Removed content record at index {args.index}!") + + elif args.cid is not None: + if len(args.cid) != 8: + raise ValueError("The provided Content ID is invalid!") + target_cid = int(args.cid, 16) + # List Contents IDs in the title, and ensure that the target Content ID exists. + valid_ids = [] + for record in tmd.content_records: + valid_ids.append(record.content_id) + if target_cid not in valid_ids: + raise ValueError("The provided Content ID could not be found in this TMD!") + tmd.content_records.pop(valid_ids.index(target_cid)) + tmd.num_contents -= 1 + # Auto fakesign because we've edited the TMD. + tmd.fakesign() + output_path.write_bytes(tmd.dump()) + print(f"Removed content record with Content ID \"{target_cid:08X}\"!") diff --git a/modules/title/wad.py b/modules/title/wad.py index 394a6e5..99feee5 100644 --- a/modules/title/wad.py +++ b/modules/title/wad.py @@ -1,7 +1,6 @@ # "modules/title/wad.py" from WiiPy by NinjaCheetah # https://github.com/NinjaCheetah/WiiPy -import os import pathlib from random import randint import libWiiPy @@ -228,6 +227,7 @@ def handle_wad_remove(args): title = libWiiPy.title.Title() title.load_wad(input_path.read_bytes()) + # TODO: see if this implementation is problematic now if args.index is not None: # List indices in the title, and ensure that the target content index exists. valid_indices = [] @@ -338,7 +338,7 @@ def handle_wad_unpack(args): if output_path.is_file(): raise ValueError("A file already exists with the provided directory name!") else: - os.mkdir(output_path) + output_path.mkdir() # Step through each component of a WAD and dump it to a file. title = libWiiPy.title.Title() diff --git a/wiipy.py b/wiipy.py index 782c38e..f13868a 100644 --- a/wiipy.py +++ b/wiipy.py @@ -13,6 +13,7 @@ from modules.title.fakesign import * from modules.title.info import * from modules.title.iospatcher import * from modules.title.nus import * +from modules.title.tmd import * from modules.title.wad import * wiipy_ver = "1.4.0" @@ -195,6 +196,25 @@ if __name__ == "__main__": setting_gen_parser.add_argument("region", metavar="REGION", type=str, help="region of the console these settings are for (USA, EUR, JPN, or KOR)") + # Argument parser for the TMD subcommand. + tmd_parser = subparsers.add_parser("tmd", help="edit a TMD file", + description="edit a TMD file") + tmd_subparsers = tmd_parser.add_subparsers(dest="subcommand", required=True) + # Remove TMD subcommand. + tmd_remove_parser = tmd_subparsers.add_parser("remove", help="remove a content record from a TMD file", + description="remove a content record from a TMD file, either by its " + "CID or by its index; by default, this will overwrite " + "the input file unless an output is specified") + tmd_remove_parser.set_defaults(func=handle_tmd_remove) + tmd_remove_parser.add_argument("input", metavar="IN", type=str, help="TMD file to remove a content record from") + tmd_remove_targets = tmd_remove_parser.add_mutually_exclusive_group(required=True) + tmd_remove_targets.add_argument("-i", "--index", metavar="INDEX", type=int, + help="index of the content record to remove") + tmd_remove_targets.add_argument("-c", "--cid", metavar="CID", type=str, + help="Content ID of the content record to remove") + tmd_remove_parser.add_argument("-o", "--output", metavar="OUT", type=str, + help="file to output the updated TMD to (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")