diff --git a/src/libWiiPy/title/content.py b/src/libWiiPy/title/content.py index f3001e0..67853a5 100644 --- a/src/libWiiPy/title/content.py +++ b/src/libWiiPy/title/content.py @@ -158,16 +158,11 @@ class ContentRegion: bytes The encrypted content listed in the content record. """ - # Get a list of the current Content IDs, so we can make sure the target one exists. - content_ids = [] - for record in self.content_records: - content_ids.append(record.content_id) - if cid not in content_ids: - raise ValueError("You are trying to get a content with Content ID " + str(cid) + ", but no content with " - "that ID exists!") - # Get the content index associated with the CID we now know exists. - target_index = content_ids.index(cid) - content_index = self.content_records[target_index].index + try: + content_index = self.get_index_from_cid(cid) + except ValueError: + raise ValueError(f"You are trying to get a content with Content ID {cid}, " + f"but no content with that ID exists!") content_enc = self.get_enc_content_by_index(content_index) return content_enc @@ -246,16 +241,11 @@ class ContentRegion: bytes The decrypted content listed in the content record. """ - # Get a list of the current Content IDs, so we can make sure the target one exists. - content_ids = [] - for record in self.content_records: - content_ids.append(record.content_id) - if cid not in content_ids: - raise ValueError("You are trying to get a content with Content ID " + str(cid) + ", but no content with " - "that ID exists!") - # Get the content index associated with the CID we now know exists. - target_index = content_ids.index(cid) - content_index = self.content_records[target_index].index + try: + content_index = self.get_index_from_cid(cid) + except ValueError: + raise ValueError(f"You are trying to get a content with Content ID {cid}, " + f"but no content with that ID exists!") content_dec = self.get_content_by_index(content_index, title_key, skip_hash) return content_dec @@ -281,6 +271,32 @@ class ContentRegion: dec_contents.append(self.get_content_by_index(content, title_key, skip_hash)) return dec_contents + def get_index_from_cid(self, cid: int) -> int: + """ + Gets the content index of a content by its Content ID. The returned index 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 while still retaining the original indices. + + Parameters + ---------- + cid : int + The Content ID to get the index of. + + Returns + ------- + int + The content index. + """ + # Get a list of the current Content IDs, so we can make sure the target one exists. + content_ids = [] + for record in self.content_records: + content_ids.append(record.content_id) + if cid not in content_ids: + raise ValueError("The specified Content ID does not exist!") + literal_index = content_ids.index(cid) + target_index = self.content_records[literal_index].index + return target_index + def add_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int, content_hash: bytes) -> None: """ @@ -534,16 +550,11 @@ class ContentRegion: cid : int The Content ID of the content you want to remove. """ - # Get a list of the current Content IDs, so we can make sure the target one exists. - content_ids = [] - for record in self.content_records: - content_ids.append(record.content_id) - if cid not in content_ids: - raise ValueError("You are trying to remove content with Content ID " + str(cid) + ", but no content with " - "that ID exists!") - # Remove the content index associated with the CID we now know exists. - target_index = content_ids.index(cid) - content_index = self.content_records[target_index].index + try: + content_index = self.get_index_from_cid(cid) + except ValueError: + raise ValueError(f"You are trying to remove content with Content ID {cid}, " + f"but no content with that ID exists!") self.remove_content_by_index(content_index) diff --git a/src/libWiiPy/title/title.py b/src/libWiiPy/title/title.py index 29c16a8..e2702a0 100644 --- a/src/libWiiPy/title/title.py +++ b/src/libWiiPy/title/title.py @@ -244,6 +244,55 @@ class Title: blocks = math.ceil(title_size_bytes / 131072) return blocks + def add_enc_content(self, enc_content: bytes, cid: int, index: int, content_type: int, content_size: int, + content_hash: bytes) -> None: + """ + Adds a new encrypted 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. + + Parameters + ---------- + enc_content : bytes + The new encrypted content to add. + cid : int + The Content ID to assign the new content in the content record. + index : int + The index used when encrypting the new content. + content_type : int + The type of the new content. + content_size : int + The size of the new encrypted content when decrypted. + content_hash : bytes + The hash of the new encrypted content when decrypted. + """ + # Add the encrypted content. + self.content.add_enc_content(enc_content, cid, index, content_type, content_size, content_hash) + # Update the TMD to match. + self.tmd.content_records = self.content.content_records + + def add_content(self, dec_content: bytes, cid: int, content_type: int) -> None: + """ + 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 + Title Key before adding it to the ContentRegion. + + Parameters + ---------- + dec_content : bytes + The new decrypted content to add. + cid : int + The Content ID to assign the new content in the content record. + content_type : int + The type of the new content. + """ + # Add the decrypted content. + self.content.add_content(dec_content, cid, content_type, self.ticket.get_title_key()) + # Update the TMD to match. + self.tmd.content_records = self.content.content_records + def set_enc_content(self, enc_content: bytes, index: int, content_size: int, content_hash: bytes, cid: int = None, content_type: int = None) -> None: """