New command to edit TMDs, currently just supports removing content records

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}\"!")

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()

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")