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