From 97bc77b337cce2e1d1cb125d81cdcae53d1308f6 Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Thu, 19 Sep 2024 14:41:23 -0400
Subject: [PATCH] Added wad set command to replace content in a WAD

---
 modules/title/wad.py | 66 ++++++++++++++++++++++++++++++++++++++++++++
 wiipy.py             | 18 ++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/modules/title/wad.py b/modules/title/wad.py
index 0d3fbdf..be35d6a 100644
--- a/modules/title/wad.py
+++ b/modules/title/wad.py
@@ -202,6 +202,72 @@ def handle_wad_remove(args):
         print(f"Removed content with Content ID \"{target_cid:08X}\"!")
 
 
+def handle_wad_set(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)
+
+    title = libWiiPy.title.Title()
+    title.load_wad(open(input_path, "rb").read())
+
+    content_data = open(content_path, "rb").read()
+
+    # Get the new type of the content, if one was specified.
+    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 = None
+
+    if args.index is not None:
+        # If we're replacing based on the index, then make sure the specified index exists.
+        existing_indices = []
+        for record in title.content.content_records:
+            existing_indices.append(record.index)
+        if args.index not in existing_indices:
+            raise ValueError("The provided index could not be found in this title!")
+        if target_type:
+            title.set_content(content_data, args.index, content_type=target_type)
+        else:
+            title.set_content(content_data, args.index)
+        open(output_path, "wb").write(title.dump_wad())
+        print(f"Replaced content at content index {args.index}!")
+
+
+    elif args.cid is not None:
+        # If we're replacing based on the CID, then make sure the specified CID is valid and exists.
+        if len(args.cid) != 8:
+            raise ValueError("The provided Content ID is invalid!")
+        target_cid = int(args.cid, 16)
+        existing_cids = []
+        for record in title.content.content_records:
+            existing_cids.append(record.content_id)
+        if target_cid not in existing_cids:
+            raise ValueError("The provided Content ID could not be found in this title!")
+        target_index = title.content.get_index_from_cid(target_cid)
+        if target_type:
+            title.set_content(content_data, target_index, content_type=target_type)
+        else:
+            title.set_content(content_data, target_index)
+        open(output_path, "wb").write(title.dump_wad())
+        print(f"Replaced content with Content ID \"{target_cid:08X}\"!")
+
+
 def handle_wad_unpack(args):
     input_path = pathlib.Path(args.input)
     output_path = pathlib.Path(args.output)
diff --git a/wiipy.py b/wiipy.py
index de5b625..d727292 100644
--- a/wiipy.py
+++ b/wiipy.py
@@ -198,6 +198,24 @@ if __name__ == "__main__":
                                     help="Content ID of the content to remove")
     wad_remove_parser.add_argument("-o", "--output", metavar="OUT", type=str,
                                    help="file to output the updated WAD to (optional)")
+    # Set WAD subcommand.
+    wad_set_parser = wad_subparsers.add_parser("set", help="set content in a WAD file",
+                                               description="replace existing content in a WAD file with new decrypted "
+                                                           "data; by default, this will overwrite the input file "
+                                                           "unless an output is specified")
+    wad_set_parser.set_defaults(func=handle_wad_set)
+    wad_set_parser.add_argument("input", metavar="IN", type=str, help="WAD file to replace content in")
+    wad_set_parser.add_argument("content", metavar="CONTENT", type=str, help="new decrypted content")
+    wad_set_targets = wad_set_parser.add_mutually_exclusive_group(required=True)
+    wad_set_targets.add_argument("-i", "--index", metavar="INDEX", type=int,
+                                 help="index of the content to replace")
+    wad_set_targets.add_argument("-c", "--cid", metavar="CID", type=str,
+                                 help="Content ID of the content to replace")
+    wad_set_parser.add_argument("-o", "--output", metavar="OUT", type=str,
+                                help="file to output the updated WAD to (optional)")
+    wad_set_parser.add_argument("-t", "--type", metavar="TYPE", type=str,
+                                help="specifies a new type for the content, can be \"Normal\", \"Shared\", or \"DLC\" "
+                                     "(optional)")
     # Unpack WAD subcommand.
     wad_unpack_parser = wad_subparsers.add_parser("unpack", help="unpack a WAD file to a directory",
                                                   description="unpack a WAD file to a directory")