diff --git a/src/libWiiPy/content.py b/src/libWiiPy/content.py index 344f1a3..97851f5 100644 --- a/src/libWiiPy/content.py +++ b/src/libWiiPy/content.py @@ -157,14 +157,14 @@ class ContentRegion: self.content_records[index].content_size) # Hash the decrypted content and ensure that the hash matches the one in its Content Record. # If it does not, then something has gone wrong in the decryption, and an error will be thrown. - content_dec_hash = hashlib.sha1(content_dec) + content_dec_hash = hashlib.sha1(content_dec).hexdigest() content_record_hash = str(self.content_records[index].content_hash.decode()) # Compare the hash and throw a ValueError if the hash doesn't match. - if content_dec_hash.hexdigest() != content_record_hash: + if content_dec_hash != content_record_hash: raise ValueError("Content hash did not match the expected hash in its record! The incorrect Title Key may " "have been used!.\n" "Expected hash is: {}\n".format(content_record_hash) + - "Actual hash is: {}".format(content_dec_hash.hexdigest())) + "Actual hash is: {}".format(content_dec_hash)) return content_dec def get_content_by_cid(self, cid: int, title_key: bytes) -> bytes: @@ -282,10 +282,52 @@ class ContentRegion: # Pass values to set_enc_content() self.set_enc_content(enc_content, cid, index, content_type, dec_content_size, dec_content_hash) - def load_enc_content(self, enc_content: bytes, index: int) -> bytes: + def load_enc_content(self, enc_content: bytes, index: int) -> None: """Loads the provided encrypted content into the content region at the specified index, with the assumption that - it matches the record at that index. + it matches the record at that index. Not recommended for most use cases, use decrypted content and + load_content() instead. - :param index: - :return: + Parameters + ---------- + enc_content : bytes + The encrypted content to load. + index : int + The content index to load the content at. """ + if (index + 1) > len(self.content_records) or len(self.content_records) == 0: + raise IndexError("No content records have been loaded, or that index is higher than the highest entry in " + "the content records.") + if (index + 1) > len(self.content_list): + self.content_list.append(enc_content) + else: + self.content_list[index] = enc_content + + def load_content(self, dec_content: bytes, index: int, title_key: bytes) -> None: + """Loads the provided decrypted content into the content region at the specified index, but first checks to make + sure it matches the record at that index before loading. This content will be encrypted when loaded. + + Parameters + ---------- + dec_content : bytes + The decrypted content to load. + index : int + The content index to load the content at. + title_key: bytes + The Title Key that matches the decrypted content. + """ + # Make sure that content records exist and that the provided index exists in them. + if (index + 1) > len(self.content_records) or len(self.content_records) == 0: + raise IndexError("No content records have been loaded, or that index is higher than the highest entry in " + "the content records.") + # Check the hash of the content against the hash stored in the record to ensure it matches. + content_hash = hashlib.sha1(dec_content).hexdigest() + if content_hash != self.content_records[index].content_hash.decode(): + raise ValueError("The decrypted content provided does not match the record at the provided index. \n" + "Expected hash is: {}\n".format(self.content_records[index].content_hash.decode()) + + "Actual hash is: {}".format(content_hash)) + # If the hash matches, encrypt the content and set it where it belongs. + enc_content = encrypt_content(dec_content, title_key, index) + if (index + 1) > len(self.content_list): + self.content_list.append(enc_content) + else: + self.content_list[index] = enc_content diff --git a/src/libWiiPy/title.py b/src/libWiiPy/title.py index 05e0dc5..69e9a52 100644 --- a/src/libWiiPy/title.py +++ b/src/libWiiPy/title.py @@ -27,7 +27,7 @@ class Title: self.ticket: Ticket = Ticket() self.content: ContentRegion = ContentRegion() - def set_wad(self, wad: bytes) -> None: + def load_wad(self, wad: bytes) -> None: """Load existing WAD data into the title and create WAD, TMD, Ticket, and ContentRegion objects based off of it to allow you to modify that data. Note that this will overwrite any existing data for this title. @@ -54,7 +54,7 @@ class Title: raise ValueError("The Title IDs of the TMD and Ticket in this WAD do not match. This WAD appears to be " "invalid.") - def dump(self) -> bytes: + def dump_wad(self) -> bytes: """Dumps all title components (TMD, Ticket, and contents) back into the WAD object, and then dumps the WAD back into raw data and returns it. @@ -73,6 +73,38 @@ class Title: wad_data = self.wad.dump() return wad_data + def load_tmd(self, tmd: bytes) -> None: + """Load existing TMD data into the title. Note that this will overwrite any existing TMD data for this title. + + Parameters + ---------- + tmd : bytes + The data for the WAD you wish to load. + """ + # Load TMD. + self.tmd.load(tmd) + + def load_ticket(self, ticket: bytes) -> None: + """Load existing Ticket data into the title. Note that this will overwrite any existing Ticket data for this + title. + + Parameters + ---------- + ticket : bytes + The data for the WAD you wish to load. + """ + # Load Ticket. + self.ticket.load(ticket) + + def load_content_records(self) -> None: + """Load content records from the TMD into the ContentRegion to allow loading content files based on the records. + This requires that a TMD has already been loaded and will throw an exception if it isn't. + """ + if not self.tmd.content_records: + ValueError("No TMD appears to have been loaded, so content records cannot be read from it.") + # Load the content records into the ContentRegion object. + self.content.content_records = self.tmd.content_records + def set_title_id(self, title_id: str) -> None: """Sets the Title ID of the title in both the TMD and Ticket. @@ -170,3 +202,17 @@ class Title: self.content.set_content(dec_content, cid, index, content_type, self.ticket.get_title_key()) # Update the TMD to match. self.tmd.content_records = self.content.content_records + + def load_content(self, dec_content: bytes, index: int) -> None: + """Loads the provided decrypted content into the content region at the specified index, but first checks to make + sure it matches the record at that index before loading. This content will be encrypted when loaded. + + Parameters + ---------- + dec_content : bytes + The decrypted content to load. + index : int + The content index to load the content at. + """ + # Load the decrypted content. + self.content.load_content(dec_content, index, self.ticket.get_title_key()) diff --git a/src/libWiiPy/wad.py b/src/libWiiPy/wad.py index 5e3de1f..f98c323 100644 --- a/src/libWiiPy/wad.py +++ b/src/libWiiPy/wad.py @@ -14,8 +14,8 @@ class WAD: """ def __init__(self): self.wad_hdr_size: int = 64 - self.wad_type: str = "" - self.wad_version: bytes = b'' + self.wad_type: str = "Is" + self.wad_version: bytes = b'\x00\x00' # === Sizes === self.wad_cert_size: int = 0 self.wad_crl_size: int = 0