From 5257c3ae6882212ba3f2b1d751f85217f1c3ff99 Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Mon, 1 Apr 2024 22:18:50 -0400
Subject: [PATCH] Rewritten to support upcoming libWiiPy v0.2.0

---
 main.py        |  4 +++
 modules/wad.py | 98 +++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 81 insertions(+), 21 deletions(-)

diff --git a/main.py b/main.py
index ae51e55..52850c7 100644
--- a/main.py
+++ b/main.py
@@ -13,6 +13,10 @@ if __name__ == "__main__":
             if len(args) == 3:
                 extract_wad_to_folder(args[1], args[2])
                 exit(0)
+        if "-p" in opts:
+            if len(args) == 3:
+                pack_wad_from_folder(args[1], args[2])
+                exit(0)
         raise SystemExit(f"Usage: {sys.argv[0]} WAD (-u | -p) <input> <output>")
     else:
         raise SystemExit(f"Usage: {sys.argv[0]} WAD (-u | -p) <input> <output>")
diff --git a/modules/wad.py b/modules/wad.py
index 55c1617..c5eafca 100644
--- a/modules/wad.py
+++ b/modules/wad.py
@@ -4,48 +4,104 @@
 import os
 import pathlib
 import binascii
-from libWiiPy import wad, tmd, ticket, content
+import libWiiPy
 
 
-def extract_wad_to_folder(in_file_input: str, out_folder_input: str):
-    if not os.path.isfile(in_file_input):
-        raise FileNotFoundError(in_file_input)
-    out_folder = pathlib.Path(out_folder_input)
+def extract_wad_to_folder(in_file: str, out_folder: str):
+    if not os.path.isfile(in_file):
+        raise FileNotFoundError(in_file)
+    out_folder = pathlib.Path(out_folder)
     if not out_folder.is_dir():
         out_folder.mkdir()
 
-    with open(in_file_input, "rb") as wad_file:
-        wad_data = wad.WAD(wad_file.read())
+    with open(in_file, "rb") as wad_file:
+        wad = libWiiPy.WAD()
+        wad.load(wad_file.read())
 
-        tmd_data = tmd.TMD(wad_data.get_tmd_data())
-        ticket_data = ticket.Ticket(wad_data.get_ticket_data())
-        content_data = content.ContentRegion(wad_data.get_content_data(), tmd_data.content_records)
+        tmd = libWiiPy.TMD()
+        tmd.load(wad.get_tmd_data())
+        ticket = libWiiPy.Ticket()
+        ticket.load(wad.get_ticket_data())
+        content_region = libWiiPy.ContentRegion()
+        content_region.load(wad.get_content_data(), tmd.content_records)
 
-        title_key = ticket_data.get_title_key()
+        title_key = ticket.get_title_key()
 
-        cert_name = tmd_data.title_id + ".cert"
+        cert_name = tmd.title_id + ".cert"
         cert_out = open(os.path.join(out_folder, cert_name), "wb")
-        cert_out.write(wad_data.get_cert_data())
+        cert_out.write(wad.get_cert_data())
         cert_out.close()
 
-        tmd_name = tmd_data.title_id + ".tmd"
+        tmd_name = tmd.title_id + ".tmd"
         tmd_out = open(os.path.join(out_folder, tmd_name), "wb")
-        tmd_out.write(wad_data.get_tmd_data())
+        tmd_out.write(wad.get_tmd_data())
         tmd_out.close()
 
-        ticket_name = tmd_data.title_id + ".tik"
+        ticket_name = tmd.title_id + ".tik"
         ticket_out = open(os.path.join(out_folder, ticket_name), "wb")
-        ticket_out.write(wad_data.get_ticket_data())
+        ticket_out.write(wad.get_ticket_data())
         ticket_out.close()
 
-        meta_name = tmd_data.title_id + ".footer"
+        meta_name = tmd.title_id + ".footer"
         meta_out = open(os.path.join(out_folder, meta_name), "wb")
-        meta_out.write(wad_data.get_meta_data())
+        meta_out.write(wad.get_meta_data())
         meta_out.close()
 
-        for content_file in range(0, tmd_data.num_contents):
+        for content_file in range(0, tmd.num_contents):
             content_file_name = "000000"
             content_file_name += str(binascii.hexlify(content_file.to_bytes()).decode()) + ".app"
             content_out = open(os.path.join(out_folder, content_file_name), "wb")
-            content_out.write(content_data.get_content(content_file, title_key))
+            content_out.write(content_region.get_content_by_index(content_file, title_key))
             content_out.close()
+
+
+def pack_wad_from_folder(in_folder, out_file):
+    if not os.path.exists(in_folder):
+        raise FileNotFoundError(in_folder)
+    if not os.path.isdir(in_folder):
+        raise NotADirectoryError(in_folder)
+
+    out_file = pathlib.Path(out_file)
+    in_folder = pathlib.Path(in_folder)
+
+    tmd_file = list(in_folder.glob("*.tmd"))[0]
+    if not os.path.exists(tmd_file):
+        raise FileNotFoundError("Cannot find a TMD! Exiting...")
+
+    ticket_file = list(in_folder.glob("*.tik"))[0]
+    if not os.path.exists(ticket_file):
+        raise FileNotFoundError("Cannot find a Ticket! Exiting...")
+
+    cert_file = list(in_folder.glob("*.cert"))[0]
+    if not os.path.exists(cert_file):
+        raise FileNotFoundError("Cannot find a cert! Exiting...")
+
+    content_files = list(in_folder.glob("*.app"))
+    if not content_files:
+        raise FileNotFoundError("Cannot find any contents! Exiting...")
+
+    with open(out_file, "wb") as out_file:
+        title = libWiiPy.Title()
+
+        title.load_tmd(open(tmd_file, "rb").read())
+        title.load_ticket(open(ticket_file, "rb").read())
+        title.wad.set_cert_data(open(cert_file, "rb").read())
+        footer_file = list(in_folder.glob("*.footer"))[0]
+        if os.path.exists(footer_file):
+            title.wad.set_meta_data(open(footer_file, "rb").read())
+        title.load_content_records()
+
+        title_key = title.ticket.get_title_key()
+
+        content_list = list(in_folder.glob("*.app"))
+        for index in range(len(title.content.content_records)):
+            for content in range(len(content_list)):
+                dec_content = open(content_list[content], "rb").read()
+                try:
+                    title.content.load_content(dec_content, index, title_key)
+                    break
+                except ValueError:
+                    print("Match not found, trying again...")
+
+        out_file.write(title.dump_wad())
+