wad.py and tmd.py now support dumping the object back into raw data

src/libWiiPy/shared.py

@@ -20,3 +20,23 @@ def align_value(value, alignment=64):
         aligned_value = value + (alignment - (value % alignment))
         return aligned_value
     return value
+
+
+def pad_bytes_stream(data, alignment=64):
+    """Pads the provided bytes stream to the provided alignment (defaults to 64).
+
+    Parameters
+    ----------
+    data : BytesIO
+        The data to align.
+    alignment : int
+        The number to align to. Defaults to 64.
+
+    Returns
+    -------
+    BytesIO
+        The aligned data.
+    """
+    while (data.getbuffer().nbytes % alignment) != 0:
+        data.write(b'\x00')
+    return data

src/libWiiPy/tmd.py

@@ -36,87 +36,103 @@ class TMD:
     """
     def __init__(self, tmd):
         self.tmd = tmd
-        self.sig_type: int
+        self.blob_header: bytes = b''
-        self.sig: bytearray
+        self.sig_type: int = 0
-        self.issuer: bytearray  # Follows the format "Root-CA%08x-CP%08x"
+        self.sig: bytes = b''
-        self.tmd_version: int  # This seems to always be 0 no matter what?
+        self.issuer: bytes = b''  # Follows the format "Root-CA%08x-CP%08x"
-        self.ca_crl_version: int
+        self.tmd_version: int = 0  # This seems to always be 0 no matter what?
-        self.signer_crl_version: int
+        self.ca_crl_version: int = 0
-        self.vwii: int  # Whether the title is for the vWii. 0 = No, 1 = Yes
+        self.signer_crl_version: int = 0
-        self.ios_tid: str  # The Title ID of the IOS version the associated title runs on.
+        self.vwii: int = 0  # Whether the title is for the vWii. 0 = No, 1 = Yes
-        self.ios_version: int  # The IOS version the associated title runs on.
+        self.ios_tid: str = ""  # The Title ID of the IOS version the associated title runs on.
-        self.title_id: str  # The Title ID of the associated title.
+        self.ios_version: int = 0  # The IOS version the associated title runs on.
-        self.content_type: str  # The type of content contained within the associated title.
+        self.title_id: str = ""  # The Title ID of the associated title.
-        self.group_id: int  # The ID of the publisher of the associated title.
+        self.content_type: str = ""  # The type of content contained within the associated title.
-        self.region: int  # The ID of the region of the associated title.
+        self.group_id: int = 0  # The ID of the publisher of the associated title.
-        self.ratings: int
+        self.region: int = 0  # The ID of the region of the associated title.
-        self.access_rights: int
+        self.ratings: bytes = b''
-        self.title_version: int  # The version of the associated title.
+        self.ipc_mask: bytes = b''
-        self.num_contents: int  # The number of contents contained in the associated title.
+        self.access_rights: bytes = b''
-        self.boot_index: int
+        self.title_version: int = 0  # The version of the associated title.
+        self.num_contents: int = 0  # The number of contents contained in the associated title.
+        self.boot_index: int = 0
         self.content_records: List[ContentRecord] = []
-        # Load data from TMD file
+        # Call load() to set all of the attributes from the raw TMD data provided.
+        self.load()
+
+    def load(self):
+        """Loads the raw TMD data and sets all attributes of the TMD object.
+
+        Returns
+        -------
+        none
+        """
         with io.BytesIO(self.tmd) as tmd_data:
             # ====================================================================================
             # Parses each of the keys contained in the TMD.
             # ====================================================================================
-            # Signing certificate issuer
+            tmd_data.seek(0x0)
+            self.blob_header = tmd_data.read(320)
+            # Signing certificate issuer.
             tmd_data.seek(0x140)
             self.issuer = tmd_data.read(64)
-            # TMD version, seems to usually be 0, but I've seen references to other numbers
+            # TMD version, seems to usually be 0, but I've seen references to other numbers.
             tmd_data.seek(0x180)
             self.tmd_version = int.from_bytes(tmd_data.read(1))
-            # TODO: label
+            # Root certificate crl version.
             tmd_data.seek(0x181)
-            self.ca_crl_version = tmd_data.read(1)
+            self.ca_crl_version = int.from_bytes(tmd_data.read(1))
-            # TODO: label
+            # Signer crl version.
             tmd_data.seek(0x182)
-            self.signer_crl_version = tmd_data.read(1)
+            self.signer_crl_version = int.from_bytes(tmd_data.read(1))
-            # If this is a vWii title or not
+            # If this is a vWii title or not.
             tmd_data.seek(0x183)
             self.vwii = int.from_bytes(tmd_data.read(1))
-            # TID of the IOS to use for the title, set to 0 if this title is the IOS, set to boot2 version if boot2
+            # TID of the IOS to use for the title, set to 0 if this title is the IOS, set to boot2 version if boot2.
             tmd_data.seek(0x184)
             ios_version_bin = tmd_data.read(8)
             ios_version_hex = binascii.hexlify(ios_version_bin)
             self.ios_tid = str(ios_version_hex.decode())
-            # Get IOS version based on TID
+            # Get IOS version based on TID.
             self.ios_version = int(self.ios_tid[-2:], 16)
-            # Title ID of the title
+            # Title ID of the title.
             tmd_data.seek(0x18C)
             title_id_bin = tmd_data.read(8)
             title_id_hex = binascii.hexlify(title_id_bin)
             self.title_id = str(title_id_hex.decode())
-            # Type of content
+            # Type of content.
             tmd_data.seek(0x194)
             content_type_bin = tmd_data.read(4)
             content_type_hex = binascii.hexlify(content_type_bin)
             self.content_type = str(content_type_hex.decode())
-            # Publisher of the title
+            # Publisher of the title.
             tmd_data.seek(0x198)
-            self.group_id = tmd_data.read(2)
+            self.group_id = int.from_bytes(tmd_data.read(2))
-            # Region of the title, 0 = JAP, 1 = USA, 2 = EUR, 3 = NONE, 4 = KOR
+            # Region of the title, 0 = JAP, 1 = USA, 2 = EUR, 3 = NONE, 4 = KOR.
             tmd_data.seek(0x19C)
             region_hex = tmd_data.read(2)
             self.region = int.from_bytes(region_hex)
-            # TODO: figure this one out
+            # Likely the localized content rating for the title. (ESRB, CERO, PEGI, etc.)
             tmd_data.seek(0x19E)
             self.ratings = tmd_data.read(16)
-            # Access rights of the title; DVD-video access and AHBPROT
+            # IPC mask.
+            tmd_data.seek(0x1BA)
+            self.ipc_mask = tmd_data.read(12)
+            # Access rights of the title; DVD-video access and AHBPROT.
             tmd_data.seek(0x1D8)
             self.access_rights = tmd_data.read(4)
-            # Calculate the version number by multiplying 0x1DC by 256 and adding 0x1DD
+            # Calculate the version number by multiplying 0x1DC by 256 and adding 0x1DD.
             tmd_data.seek(0x1DC)
             title_version_high = int.from_bytes(tmd_data.read(1)) * 256
             tmd_data.seek(0x1DD)
             title_version_low = int.from_bytes(tmd_data.read(1))
             self.title_version = title_version_high + title_version_low
-            # The number of contents listed in the TMD
+            # The number of contents listed in the TMD.
             tmd_data.seek(0x1DE)
             self.num_contents = int.from_bytes(tmd_data.read(2))
-            # Content index in content list that contains the boot file
+            # Content index in content list that contains the boot file.
             tmd_data.seek(0x1E0)
-            self.boot_index = tmd_data.read(2)
+            self.boot_index = int.from_bytes(tmd_data.read(2))
             # Get content records for the number of contents in num_contents.
             for content in range(0, self.num_contents):
                 tmd_data.seek(0x1E4 + (36 * content))
@@ -126,6 +142,82 @@ class TMD:
                                   int(content_record_hdr[2]), int.from_bytes(content_record_hdr[3]),
                                   binascii.hexlify(content_record_hdr[4])))
 
+    def dump(self) -> bytes:
+        """Dumps the TMD object back into bytes. This also sets the raw TMD attribute of TMD object to the dumped data,
+        and triggers load() again to ensure that the raw data and object match.
+
+        Returns
+        -------
+        bytes
+            The full TMD file as bytes.
+        """
+        # Open the stream and begin writing to it.
+        with io.BytesIO() as tmd_data:
+            # Signed blob header.
+            tmd_data.write(self.blob_header)
+            # Signing certificate issuer.
+            tmd_data.write(self.issuer)
+            # TMD version.
+            tmd_data.write(int.to_bytes(self.tmd_version, 1))
+            # Root certificate crl version.
+            tmd_data.write(int.to_bytes(self.ca_crl_version, 1))
+            # Signer crl version.
+            tmd_data.write(int.to_bytes(self.signer_crl_version, 1))
+            # If this is a vWii title or not.
+            tmd_data.write(int.to_bytes(self.vwii, 1))
+            # IOS Title ID.
+            tmd_data.write(binascii.unhexlify(self.ios_tid))
+            # Title's Title ID.
+            tmd_data.write(binascii.unhexlify(self.title_id))
+            # Content type.
+            tmd_data.write(binascii.unhexlify(self.content_type))
+            # Group ID.
+            tmd_data.write(int.to_bytes(self.group_id, 2))
+            # 2 bytes of zero for reasons.
+            tmd_data.write(b'\x00\x00')
+            # Region.
+            tmd_data.write(int.to_bytes(self.region, 2))
+            # Ratings.
+            tmd_data.write(self.ratings)
+            # Reserved (All \x00).
+            tmd_data.write(b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
+            # IPC mask.
+            tmd_data.write(self.ipc_mask)
+            # Reserved (ALl \x00).
+            tmd_data.write(b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
+            # Access rights.
+            tmd_data.write(self.access_rights)
+            # Title version.
+            title_version_high = round(self.title_version / 256)
+            tmd_data.write(int.to_bytes(title_version_high, 1))
+            title_version_low = self.title_version % 256
+            tmd_data.write(int.to_bytes(title_version_low, 1))
+            # Number of contents.
+            tmd_data.write(int.to_bytes(self.num_contents, 2))
+            # Boot index.
+            tmd_data.write(int.to_bytes(self.boot_index, 2))
+            # Minor version. Unused so write \x00.
+            tmd_data.write(b'\x00\x00')
+            # Iterate over content records, write them back into raw data, then add them to the TMD.
+            for content_record in range(self.num_contents):
+                content_data = io.BytesIO()
+                # Write all fields from the content record.
+                content_data.write(int.to_bytes(self.content_records[content_record].content_id, 4))
+                content_data.write(int.to_bytes(self.content_records[content_record].index, 2))
+                content_data.write(int.to_bytes(self.content_records[content_record].content_type, 2))
+                content_data.write(int.to_bytes(self.content_records[content_record].content_size, 8))
+                content_data.write(binascii.unhexlify(self.content_records[content_record].content_hash))
+                # Seek to the start and write the record to the TMD.
+                content_data.seek(0x0)
+                tmd_data.write(content_data.read())
+                content_data.close()
+
+            tmd_data.seek(0x0)
+            self.tmd = tmd_data.read()
+        # Reload object's attributes to ensure the raw data and object match.
+        self.load()
+        return self.tmd
+
     def get_title_region(self):
         """Gets the region of the TMD's associated title.
 

src/libWiiPy/wad.py

@@ -5,7 +5,7 @@
 
 import io
 import binascii
-from .shared import align_value
+from .shared import align_value, pad_bytes_stream
 
 
 class WAD:
@@ -19,25 +19,34 @@ class WAD:
     """
     def __init__(self, wad):
         self.wad = wad
-        self.wad_hdr_size: int
+        self.wad_hdr_size: int = 64
-        self.wad_type: str
+        self.wad_type: str = ""
-        self.wad_version: int
+        self.wad_version: bytes = b''
         # === Sizes ===
-        self.wad_cert_size: int
+        self.wad_cert_size: int = 0
-        self.wad_crl_size: int
+        self.wad_crl_size: int = 0
-        self.wad_tik_size: int
+        self.wad_tik_size: int = 0
-        self.wad_tmd_size: int
+        self.wad_tmd_size: int = 0
         # This is the size of the content region, which contains all app files combined.
-        self.wad_content_size: int
+        self.wad_content_size: int = 0
-        self.wad_meta_size: int
+        self.wad_meta_size: int = 0
         # === Offsets ===
-        self.wad_cert_offset: int
+        self.wad_cert_offset: int = 0
-        self.wad_crl_offset: int
+        self.wad_crl_offset: int = 0
-        self.wad_tik_offset: int
+        self.wad_tik_offset: int = 0
-        self.wad_tmd_offset: int
+        self.wad_tmd_offset: int = 0
-        self.wad_content_offset: int
+        self.wad_content_offset: int = 0
-        self.wad_meta_offset: int
+        self.wad_meta_offset: int = 0
-        # Load header data from WAD stream
+        # Call load() to set all of the attributes from the raw WAD data provided.
+        self.load()
+
+    def load(self):
+        """Loads the raw WAD data and sets all attributes of the WAD object.
+
+        Returns
+        -------
+        none
+        """
         with io.BytesIO(self.wad) as wad_data:
             # Read the first 8 bytes of the file to ensure that it's a WAD. This will currently reject boot2 WADs, but
             # this tool cannot handle them correctly right now anyway.
@@ -89,6 +98,61 @@ class WAD:
             self.wad_meta_offset = align_value(self.wad_tmd_offset + self.wad_tmd_size)
             self.wad_content_offset = align_value(self.wad_meta_offset + self.wad_meta_size)
 
+    def dump(self) -> bytes:
+        """Dumps the WAD object back into bytes. This also sets the raw WAD attribute of WAD object to the dumped data,
+        and triggers load() again to ensure that the raw data and object match.
+
+        Returns
+        -------
+        bytes
+            The full WAD file as bytes.
+        """
+        # Open the stream and begin writing data to it.
+        with io.BytesIO() as wad_data:
+            # Lead-in data.
+            wad_data.write(b'\x00\x00\x00\x20')
+            # WAD type.
+            wad_data.write(str.encode(self.wad_type))
+            # WAD version.
+            wad_data.write(self.wad_version)
+            # WAD cert size.
+            wad_data.write(int.to_bytes(self.wad_cert_size, 4))
+            # WAD crl size.
+            wad_data.write(int.to_bytes(self.wad_crl_size, 4))
+            # WAD ticket size.
+            wad_data.write(int.to_bytes(self.wad_tik_size, 4))
+            # WAD TMD size.
+            wad_data.write(int.to_bytes(self.wad_tmd_size, 4))
+            # WAD content size.
+            wad_data.write(int.to_bytes(self.wad_content_size, 4))
+            # WAD meta size.
+            wad_data.write(int.to_bytes(self.wad_meta_size, 4))
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the cert data and write it out.
+            wad_data.write(self.get_cert_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the crl data and write it out.
+            wad_data.write(self.get_crl_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the ticket data and write it out.
+            wad_data.write(self.get_ticket_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the TMD data and write it out.
+            wad_data.write(self.get_tmd_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the meta/footer data and write it out.
+            wad_data.write(self.get_meta_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Retrieve the content data and write it out.
+            wad_data.write(self.get_content_data())
+            wad_data = pad_bytes_stream(wad_data)
+            # Seek to the beginning and save this as the WAD data for the object.
+            wad_data.seek(0x0)
+            self.wad = wad_data.read()
+        # Reload object's attributes to ensure the raw data and object match.
+        self.load()
+        return self.wad
+
     def get_cert_region(self):
         """Gets the offset and size of the certificate data.