diff --git a/src/libWiiPy/title/content.py b/src/libWiiPy/title/content.py index 92dbb6a..00ff87f 100644 --- a/src/libWiiPy/title/content.py +++ b/src/libWiiPy/title/content.py @@ -8,11 +8,19 @@ import io import hashlib from typing import List from dataclasses import dataclass as _dataclass +from enum import IntEnum as _IntEnum from ..types import _ContentRecord from ..shared import _pad_bytes, _align_value from .crypto import decrypt_content, encrypt_content +class ContentType(_IntEnum): + NORMAL = 1 + HASH_TREE = 3 + DLC = 16385 + SHARED = 32769 + + class ContentRegion: """ A ContentRegion object to parse the continuous content region of a WAD. Allows for retrieving content from the @@ -286,7 +294,7 @@ class ContentRegion: cid : int The Content ID to assign the new content in the content record. index : int - The index to place the new content at. + The index used when encrypting the new content. content_type : int The type of the new content. content_size : int @@ -304,10 +312,11 @@ class ContentRegion: self.content_list.append(enc_content) self.content_records.append(_ContentRecord(cid, index, content_type, content_size, content_hash)) - def add_content(self, dec_content: bytes, cid: int, index: int, content_type: int, title_key: bytes) -> None: + def add_content(self, dec_content: bytes, cid: int, content_type: int, title_key: bytes) -> None: """ - Adds a new decrypted content to the ContentRegion, and adds the provided Content ID, index, content type, - content size, and content hash to a new record in the ContentRecord list. + Adds a new decrypted content to the end of the ContentRegion, and adds the provided Content ID, content type, + content size, and content hash to a new record in the ContentRecord list. The index will be automatically + assigned by incrementing the current highest index in the records. This first gets the content hash and size from the provided data, and then encrypts the content with the provided Title Key before adding it to the ContentRegion. @@ -318,13 +327,16 @@ class ContentRegion: The new decrypted content to add. cid : int The Content ID to assign the new content in the content record. - index : int - The index to place the new content at. content_type : int The type of the new content. title_key : bytes The Title Key that matches the other content in the ContentRegion. """ + # Find the current highest content index and increment it for this content. + content_indices = [] + for record in self.content_records: + content_indices.append(record.index) + index = max(content_indices) + 1 content_size = len(dec_content) content_hash = str.encode(hashlib.sha1(dec_content).hexdigest()) enc_content = encrypt_content(dec_content, title_key, index) diff --git a/src/libWiiPy/title/title.py b/src/libWiiPy/title/title.py index f7f7a4d..29c16a8 100644 --- a/src/libWiiPy/title/title.py +++ b/src/libWiiPy/title/title.py @@ -76,7 +76,9 @@ class Title: if self.tmd.title_id == "0000000100000001": self.wad.wad_type = "ib" # Dump the TMD and set it in the WAD. + # This requires updating the content records and number of contents in the TMD first. self.tmd.content_records = self.content.content_records + self.tmd.num_contents = len(self.content.content_records) self.wad.set_tmd_data(self.tmd.dump()) # Dump the Ticket and set it in the WAD. self.wad.set_ticket_data(self.ticket.dump()) @@ -247,7 +249,7 @@ class Title: """ Sets the content at the provided content index to the provided new encrypted content. The provided hash and content size are set in the corresponding content record. A new Content ID or content type can also be - specified, but if it isn't than the current values are preserved. + specified, but if it isn't then the current values are preserved. This uses the content index, which is the value tied to each content and used as the IV for encryption, rather than the literal index in the array of content, because sometimes the contents end up out of order in a WAD @@ -279,7 +281,7 @@ class Title: """ Sets the content at the provided content index to the provided new decrypted content. The hash and content size of this content will be generated and then set in the corresponding content record. A new Content ID or content - type can also be specified, but if it isn't than the current values are preserved. + type can also be specified, but if it isn't then the current values are preserved. This also updates the content records in the TMD after the content is set. diff --git a/src/libWiiPy/title/tmd.py b/src/libWiiPy/title/tmd.py index 1be74dc..36eb0ae 100644 --- a/src/libWiiPy/title/tmd.py +++ b/src/libWiiPy/title/tmd.py @@ -8,7 +8,7 @@ import binascii import hashlib import struct from typing import List -from enum import IntEnum +from enum import IntEnum as _IntEnum from ..types import _ContentRecord from ..shared import _bitmask from .util import title_ver_dec_to_standard, title_ver_standard_to_dec @@ -390,7 +390,7 @@ class TMD: raise IndexError("Invalid content record! TMD lists '" + str(self.num_contents - 1) + "' contents but index was '" + str(record) + "'!") - class AccessFlags(IntEnum): + class AccessFlags(_IntEnum): AHB = 0 DVD_VIDEO = 1