From 83d55ec1890321ebd5de5cd1b3e95dd568c61589 Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Wed, 12 Jun 2024 14:22:17 -0400
Subject: [PATCH] Add basic support for packing/extracting U8 archives

---
 main.py       | 11 +++++++++++
 modules/u8.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+)
 create mode 100644 modules/u8.py

diff --git a/main.py b/main.py
index 7eafd5d..0b196f0 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@
 import sys
 from modules.wad import *
 from modules.nus import *
+from modules.u8 import *
 
 opts = [opt for opt in sys.argv[1:] if opt.startswith("-")]
 args = [arg for arg in sys.argv[1:] if not arg.startswith("-")]
@@ -28,5 +29,15 @@ if __name__ == "__main__":
                 download_title(args[1], args[2])
                 exit(0)
         raise SystemExit(f"Usage: {sys.argv[0]} NUS -d <Title ID> <Title Version (Optional)>")
+    elif "U8" in args:
+        if "-u" in opts:
+            if len(args) == 3:
+                extract_u8_to_folder(args[1], args[2])
+                exit(0)
+        elif "-p" in opts:
+            if len(args) == 3:
+                pack_u8_from_folder(args[1], args[2])
+                exit(0)
+        raise SystemExit(f"Usage: {sys.argv[0]} U8 (-u | -p) <input> <output>")
     else:
         raise SystemExit(f"Usage: {sys.argv[0]} WAD (-u | -p) <input> <output>")
diff --git a/modules/u8.py b/modules/u8.py
new file mode 100644
index 0000000..a30dad7
--- /dev/null
+++ b/modules/u8.py
@@ -0,0 +1,29 @@
+# "u8.py" from libWiiPy-cli by NinjaCheetah
+# https://github.com/NinjaCheetah/libWiiPy-cli
+
+import os
+import libWiiPy
+
+
+def extract_u8_to_folder(in_file: str, out_folder: str):
+    if not os.path.isfile(in_file):
+        raise FileNotFoundError(in_file)
+
+    u8_data = open(in_file, "rb").read()
+
+    try:
+        libWiiPy.archive.extract_u8(u8_data, out_folder)
+    except ValueError:
+        print("Specified output folder already exists!")
+
+
+def pack_u8_from_folder(in_folder: str, out_file: str):
+    try:
+        u8_data = libWiiPy.archive.pack_u8(in_folder)
+    except ValueError:
+        print("Specified input file/folder does not exist!")
+        return
+
+    out_file = open(out_file, "wb")
+    out_file.write(u8_data)
+    out_file.close()