From 3e9f45288531389d34373b134e786f195d4729ae Mon Sep 17 00:00:00 2001
From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com>
Date: Sun, 1 Dec 2024 12:12:33 -0500
Subject: [PATCH] Add option to set the channel name for WADs containing
 channels

---
 commands/title/info.py | 18 ++++++++++++------
 commands/title/wad.py  | 34 ++++++++++++++++++++++++++++++++++
 wiipy.py               |  2 ++
 3 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/commands/title/info.py b/commands/title/info.py
index d89acdd..1e385f7 100644
--- a/commands/title/info.py
+++ b/commands/title/info.py
@@ -98,8 +98,11 @@ def _print_ticket_info(ticket: libWiiPy.title.Ticket):
         print(f"  Title ID: {ticket.title_id.decode().upper()}")
     # This type of version number really only applies to the System Menu and IOS.
     if ticket.title_id.decode().startswith("00000001"):
-        print(f"  Title Version: {ticket.title_version} "
-              f"({libWiiPy.title.title_ver_dec_to_standard(ticket.title_version, ticket.title_id.decode())})")
+        if ticket.title_id.decode() == "0000000100000001":
+            print(f"  Title Version: {ticket.title_version} (boot2v{ticket.title_version})")
+        else:
+            print(f"  Title Version: {ticket.title_version} "
+                  f"({libWiiPy.title.title_ver_dec_to_standard(ticket.title_version, ticket.title_id.decode())})")
     else:
         print(f"  Title Version: {ticket.title_version}")
     print(f"  Ticket Version: {ticket.ticket_version}")
@@ -132,10 +135,13 @@ def _print_wad_info(title: libWiiPy.title.Title):
     print(f"WAD Info")
     banner_data = title.get_content_by_index(0)
     banner_u8 = libWiiPy.archive.U8Archive()
-    banner_u8.load(banner_data)
-    if banner_u8.imet_header.magic != "":
-        channel_title = banner_u8.imet_header.get_channel_names(banner_u8.imet_header.LocalizedTitles.TITLE_ENGLISH)
-        print(f"  Channel Name: {channel_title}")
+    try:
+        banner_u8.load(banner_data)
+        if banner_u8.imet_header.magic != "":
+            channel_title = banner_u8.imet_header.get_channel_names(banner_u8.imet_header.LocalizedTitles.TITLE_ENGLISH)
+            print(f"  Channel Name: {channel_title}")
+    except TypeError:
+        pass
     match title.wad.wad_type:
         case "Is":
             print(f"  WAD Type: Standard Installable")
diff --git a/commands/title/wad.py b/commands/title/wad.py
index bd44c27..c475d2e 100644
--- a/commands/title/wad.py
+++ b/commands/title/wad.py
@@ -1,6 +1,7 @@
 # "commands/title/wad.py" from WiiPy by NinjaCheetah
 # https://github.com/NinjaCheetah/WiiPy
 
+import io
 import pathlib
 from random import randint
 import libWiiPy
@@ -162,6 +163,39 @@ def handle_wad_edit(args):
         new_tid = title_edit_type(title.tmd.title_id, args.type)
         title.set_title_id(new_tid)
         edits_made = True
+    if args.channel_name is not None:
+        # Assess if this is actually a channel, because a channel name can't be set otherwise.
+        banner_data = title.get_content_by_index(0)
+        with io.BytesIO(banner_data) as data:
+            data.seek(0x40)
+            magic = data.read(4)
+            if magic != b'\x49\x4D\x45\x54':
+                data.seek(0x80)
+                magic = data.read(4)
+                if magic != b'\x49\x4D\x45\x54':
+                    fatal_error(f"This WAD file doesn't contain a Channel, so a new Channel name cannot be set!")
+                target = 0x40
+            else:
+                target = 0x0
+            # Read out the IMET header data, load it, edit it, then dump it back to bytes and directly write it over
+            # the old header data, since libWiiPy doesn't offer a cleaner solution currently.
+            data.seek(target)
+            imet_data = data.read(0x600)
+            imet_header = libWiiPy.archive.IMETHeader()
+            imet_header.load(imet_data)
+            target_languages = list(imet_header.LocalizedTitles)
+            try:
+                for target_language in target_languages:
+                    imet_header.set_channel_names((target_language, args.channel_name))
+            except ValueError:
+                fatal_error(f"The specified Channel name is not valid! Channel names must be no longer than 40 "
+                            f"characters.")
+            imet_data = imet_header.dump()
+            data.seek(target)
+            data.write(imet_data)
+            data.seek(0x0)
+            title.set_content(data.read(), 0)
+        edits_made = True
 
     if not edits_made:
         fatal_error("You must specify at least one change to make!")
diff --git a/wiipy.py b/wiipy.py
index c7e3b7d..f76c3fb 100644
--- a/wiipy.py
+++ b/wiipy.py
@@ -314,6 +314,8 @@ if __name__ == "__main__":
     wad_edit_parser.add_argument("--type", metavar="TYPE", type=str,
                                  help="a new title type for this WAD (valid options: System, Channel, SystemChannel, "
                                       "GameChannel, DLC, HiddenChannel)")
+    wad_edit_parser.add_argument("--channel-name", metavar="CHANNEL", type=str,
+                                 help="a new Channel name for this WAD, if it contains a channel")
     wad_edit_parser.add_argument("-o", "--output", metavar="OUT", type=str,
                                  help="file to output the updated WAD to (optional)")
     # Pack WAD subcommand.