Adjusted WAD command syntax, added subcommand to add content to an existing WAD

modules/title/wad.py

@@ -3,15 +3,80 @@
 
 import os
 import pathlib
+from random import randint
 import libWiiPy
 
 
-def handle_wad(args):
+def handle_wad_add(args):
+    input_path = pathlib.Path(args.input)
+    content_path = pathlib.Path(args.content)
+    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)
+    if not content_path.exists():
+        raise FileNotFoundError(content_path)
+
+    wad_file = open(input_path, 'rb')
+    title = libWiiPy.title.Title()
+    title.load_wad(wad_file.read())
+    wad_file.close()
+
+    content_file = open(content_path, 'rb')
+    content_data = content_file.read()
+    content_file.close()
+
+    # Prepare the CID so it's ready when we go to add this content to the WAD.
+    # We need to both validate that this is a real CID, and also that it isn't already taken by another content.
+    if args.cid is not None:
+        if len(args.cid) != 8:
+            raise ValueError("The provided Content ID is invalid!")
+        target_cid = int(args.cid, 16)
+        for record in title.content.content_records:
+            if target_cid == record.content_id:
+                raise ValueError("The provided Content ID is already being used by this title!")
+        print(f"Using provided Content ID \"{target_cid:08X}\".")
+    # If we weren't given a CID, then we need to randomly assign one, and ensure it isn't being used.
+    else:
+        used_cids = []
+        for record in title.content.content_records:
+            used_cids.append(record.content_id)
+        target_cid = randint(0, 0x000000FF)
+        while target_cid in used_cids:
+            target_cid = randint(0, 0x000000FF)
+        print(f"Using randomly assigned Content ID \"{target_cid:08X}\" since none were provided.")
+
+    # Get the type of the new content.
+    if args.type is not None:
+        match str.lower(args.type):
+            case "normal":
+                target_type = libWiiPy.title.ContentType.NORMAL
+            case "shared":
+                target_type = libWiiPy.title.ContentType.SHARED
+            case "dlc":
+                target_type = libWiiPy.title.ContentType.DLC
+            case _:
+                raise ValueError("The provided content type is invalid!")
+    else:
+        target_type = libWiiPy.title.ContentType.NORMAL
+
+    # Call add_content to add our new content with the set parameters.
+    title.content.add_content(content_data, target_cid, target_type, title.ticket.get_title_key())
+
+    out_file = open(output_path, 'wb')
+    out_file.write(title.dump_wad())
+    out_file.close()
+
+    print(f"Successfully added new content with Content ID \"{target_cid:08X}\" and type \"{target_type.name}\"!")
+
+
+def handle_wad_pack(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():
@@ -91,8 +156,10 @@ def handle_wad(args):
 
     print("WAD file packed!")
 
-    # Code for if the --unpack argument was passed.
+def handle_wad_unpack(args):
-    elif args.unpack:
+    input_path = pathlib.Path(args.input)
+    output_path = pathlib.Path(args.output)
+
     if not input_path.exists():
         raise FileNotFoundError(input_path)
     # Check if the output path already exists, and if it does, ensure that it is both a directory and empty.

wiipy.py

@@ -143,19 +143,40 @@ if __name__ == "__main__":
     # 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)
+    wad_subparsers = wad_parser.add_subparsers(dest="subcommand", required=True)
-    wad_group = wad_parser.add_mutually_exclusive_group(required=True)
+    # Add WAD subcommand.
-    wad_group.add_argument("-p", "--pack", help="pack a directory to a WAD file", action="store_true")
+    wad_add_parser = wad_subparsers.add_parser("add", help="add decrypted content to a WAD file",
-    wad_group.add_argument("-u", "--unpack", help="unpack a WAD file to a directory", action="store_true")
+                                               description="add decrypted content to a WAD file; by default, this "
-    wad_parser.add_argument("input", metavar="IN", type=str, help="input file")
+                                                        "will overwrite the input file unless an output is specified")
-    wad_parser.add_argument("output", metavar="OUT", type=str, help="output file")
+    wad_add_parser.set_defaults(func=handle_wad_add)
-    wad_pack_group = wad_parser.add_argument_group(title="packing options")
+    wad_add_parser.add_argument("input", metavar="IN", type=str, help="WAD file to add to")
-    wad_pack_group.add_argument("-f", "--fakesign", help="fakesign the TMD and Ticket (trucha bug)",
+    wad_add_parser.add_argument("content", metavar="CONTENT", type=str, help="decrypted content to add")
+    wad_add_parser.add_argument("-c", "--cid", metavar="CID", type=str,
+                                help="content ID to assign the new content (optional, will be randomly assigned if "
+                                     "not specified)")
+    wad_add_parser.add_argument("-t", "--type", metavar="TYPE", type=str,
+                                help="the type of the new content, can be \"Normal\", \"Shared\", or \"DLC\" "
+                                     "(optional, will default to \"Normal\" if not specified)")
+    wad_add_parser.add_argument("-o", "--output", metavar="OUT", type=str,
+                                help="file to output the new WAD to (optional)")
+    # Pack WAD subcommand.
+    wad_pack_parser = wad_subparsers.add_parser("pack", help="pack a directory to a WAD file",
+                                                 description="pack a directory to a WAD file")
+    wad_pack_parser.set_defaults(func=handle_wad_pack)
+    wad_pack_parser.add_argument("input", metavar="IN", type=str, help="input directory")
+    wad_pack_parser.add_argument("output", metavar="OUT", type=str, help="WAD file to pack")
+    wad_pack_parser.add_argument("-f", "--fakesign", help="fakesign the TMD and Ticket (trucha bug)",
                                  action="store_true")
-    wad_unpack_group = wad_parser.add_argument_group(title="unpacking options")
+    # Unpack WAD subcommand.
-    wad_unpack_group.add_argument("-s", "--skip-hash", help="skips validating the hashes of decrypted "
+    wad_unpack_parser = wad_subparsers.add_parser("unpack", help="unpack a WAD file to a directory",
+                                                  description="unpack a WAD file to a directory")
+    wad_unpack_parser.set_defaults(func=handle_wad_unpack)
+    wad_unpack_parser.add_argument("input", metavar="IN", type=str, help="WAD file to unpack")
+    wad_unpack_parser.add_argument("output", metavar="OUT", type=str, help="output directory")
+    wad_unpack_parser.add_argument("-s", "--skip-hash", help="skips validating the hashes of decrypted "
                                    "content", action="store_true")
 
+
     # 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)