diff --git a/src/libWiiPy/content.py b/src/libWiiPy/content.py index 64d2340..344f1a3 100644 --- a/src/libWiiPy/content.py +++ b/src/libWiiPy/content.py @@ -13,33 +13,27 @@ from .crypto import decrypt_content, encrypt_content class ContentRegion: """Creates a ContentRegion object to parse the continuous content region of a WAD. - - Parameters - ---------- - content_region : bytes - A bytes object containing the content region of a WAD file. - content_records : list[ContentRecord] - A list of ContentRecord objects detailing all contents contained in the region. """ - def __init__(self, content_region, content_records: List[ContentRecord]): - self.content_region = content_region - self.content_records = content_records + def __init__(self): + self.content_records: List[ContentRecord] = [] self.content_region_size: int = 0 # Size of the content region. self.num_contents: int = 0 # Number of contents in the content region. self.content_start_offsets: List[int] = [0] # The start offsets of each content in the content region. self.content_list: List[bytes] = [] - # Call load() to set all of the attributes from the raw content region provided. - self.load() - def load(self): + def load(self, content_region: bytes, content_records: List[ContentRecord]) -> None: """Loads the raw content region and builds a list of all the contents. - Returns - ------- - none + Parameters + ---------- + content_region : bytes + The raw data for the content region being loaded. + content_records : list[ContentRecord] + A list of ContentRecord objects detailing all contents contained in the region. """ - with io.BytesIO(self.content_region) as content_region_data: + self.content_records = content_records + with io.BytesIO(content_region) as content_region_data: # Get the total size of the content region. self.content_region_size = sys.getsizeof(content_region_data) self.num_contents = len(self.content_records) @@ -86,13 +80,9 @@ class ContentRegion: if padding_bytes > 0: content_region_data.write(b'\x00' * padding_bytes) content_region_data.seek(0x0) - self.content_region = content_region_data.read() - # Clear existing lists. - self.content_start_offsets = [0] - self.content_list = [] - # Reload object's attributes to ensure the raw data and object match. - self.load() - return self.content_region + content_region_raw = content_region_data.read() + # Return the raw ContentRegion for the data contained in the object. + return content_region_raw def get_enc_content_by_index(self, index: int) -> bytes: """Gets an individual content from the content region based on the provided index, in encrypted form. @@ -291,3 +281,11 @@ class ContentRegion: enc_content = encrypt_content(dec_content, title_key, index) # 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: + """Loads the provided encrypted content into the content region at the specified index, with the assumption that + it matches the record at that index. + + :param index: + :return: + """ diff --git a/src/libWiiPy/ticket.py b/src/libWiiPy/ticket.py index c61a2a5..39cd53c 100644 --- a/src/libWiiPy/ticket.py +++ b/src/libWiiPy/ticket.py @@ -11,12 +11,9 @@ from typing import List class Ticket: - """Creates a Ticket object to parse a Ticket file to retrieve the Title Key needed to decrypt it. - - Parameters - ---------- - ticket : bytes - A bytes object containing the contents of a ticket file. + """ + Creates a Ticket object that allows for either loading and editing an existing Ticket or creating one manually if + desired. Attributes ---------- @@ -35,8 +32,7 @@ class Ticket: common_key_index : int The index of the common key required to decrypt this ticket's Title Key. """ - def __init__(self, ticket): - self.ticket = ticket + def __init__(self): # Signature blob header self.signature_type: bytes = b'' # Type of signature, always 0x10001 for RSA-2048 self.signature: bytes = b'' # Actual signature data @@ -60,17 +56,17 @@ class Ticket: self.title_limits_list: List[TitleLimit] = [] # List of play limits applied to the title. # v1 ticket data # TODO: Write in v1 ticket attributes here. This code can currently only handle v0 tickets, and will reject v1. - # Call load() to set all of the attributes from the raw Ticket data provided. - self.load() - def load(self): - """Loads the raw Ticket data and sets all attributes of the Ticket object. + def load(self, ticket: bytes) -> None: + """Loads raw Ticket data and sets all attributes of the WAD object. This allows for manipulating an already + existing Ticket. - Returns - ------- - none + Parameters + ---------- + ticket : bytes + The data for the Ticket you wish to load. """ - with io.BytesIO(self.ticket) as ticket_data: + with io.BytesIO(ticket) as ticket_data: # ==================================================================================== # Parses each of the keys contained in the Ticket. # ==================================================================================== @@ -209,10 +205,9 @@ class Ticket: title_limit_data.close() # Set the Ticket attribute of the object to the new raw Ticket. ticket_data.seek(0x0) - self.ticket = ticket_data.read() - # Reload object's attributes to ensure the raw data and object match. - self.load() - return self.ticket + ticket_data_raw = ticket_data.read() + # Return the raw TMD for the data contained in the object. + return ticket_data_raw def get_title_id(self): """Gets the Title ID of the ticket's associated title. diff --git a/src/libWiiPy/title.py b/src/libWiiPy/title.py index 28ee20e..05e0dc5 100644 --- a/src/libWiiPy/title.py +++ b/src/libWiiPy/title.py @@ -12,11 +12,6 @@ from .wad import WAD class Title: """Creates a Title object that contains all components of a title, and allows altering them. - Parameters - ---------- - wad : WAD - A WAD object to load data from. - Attributes ---------- tmd : TMD @@ -26,18 +21,33 @@ class Title: content: ContentRegion A ContentRegion object containing the title's contents. """ - def __init__(self, wad: WAD): - self.wad = wad - self.tmd: TMD - self.ticket: Ticket - self.content: ContentRegion - # Load data from the WAD object, and generate all other objects from the data in it. + def __init__(self): + self.wad: WAD = WAD() + self.tmd: TMD = TMD() + self.ticket: Ticket = Ticket() + self.content: ContentRegion = ContentRegion() + + def set_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. + + Parameters + ---------- + wad : bytes + The data for the WAD you wish to load. + """ + # Create a new WAD object based on the WAD data provided. + self.wad = WAD() + self.wad.load(wad) # Load the TMD. - self.tmd = TMD(self.wad.get_tmd_data()) + self.tmd = TMD() + self.tmd.load(self.wad.get_tmd_data()) # Load the ticket. - self.ticket = Ticket(self.wad.get_ticket_data()) + self.ticket = Ticket() + self.ticket.load(self.wad.get_ticket_data()) # Load the content. - self.content = ContentRegion(self.wad.get_content_data(), self.tmd.content_records) + self.content = ContentRegion() + self.content.load(self.wad.get_content_data(), self.tmd.content_records) # Ensure that the Title IDs of the TMD and Ticket match before doing anything else. If they don't, throw an # error because clearly something strange has gone on with the WAD and editing it probably won't work. if self.tmd.title_id != self.ticket.title_id_str: diff --git a/src/libWiiPy/tmd.py b/src/libWiiPy/tmd.py index 5029309..7ead8c0 100644 --- a/src/libWiiPy/tmd.py +++ b/src/libWiiPy/tmd.py @@ -12,12 +12,7 @@ from .types import ContentRecord class TMD: """ - Creates a TMD object to parse a TMD file to retrieve information about a title. - - Parameters - ---------- - tmd : bytes - A bytes object containing the contents of a TMD file. + Creates a TMD object that allows for either loading and editing an existing TMD or creating one manually if desired. Attributes ---------- @@ -34,8 +29,7 @@ class TMD: num_contents : int The number of contents listed in the TMD. """ - def __init__(self, tmd): - self.tmd = tmd + def __init__(self): self.blob_header: bytes = b'' self.sig_type: int = 0 self.sig: bytes = b'' @@ -57,17 +51,17 @@ class TMD: self.num_contents: int = 0 # The number of contents contained in the associated title. self.boot_index: int = 0 self.content_records: List[ContentRecord] = [] - # 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. + def load(self, tmd: bytes) -> None: + """Loads raw TMD data and sets all attributes of the WAD object. This allows for manipulating an already + existing TMD. - Returns - ------- - none + Parameters + ---------- + tmd : bytes + The data for the TMD you wish to load. """ - with io.BytesIO(self.tmd) as tmd_data: + with io.BytesIO(tmd) as tmd_data: # ==================================================================================== # Parses each of the keys contained in the TMD. # ==================================================================================== @@ -214,12 +208,9 @@ class TMD: content_data.close() # Set the TMD attribute of the object to the new raw TMD. tmd_data.seek(0x0) - self.tmd = tmd_data.read() - # Clear existing lists. - self.content_records = [] - # Reload object's attributes to ensure the raw data and object match. - self.load() - return self.tmd + tmd_data_raw = tmd_data.read() + # Return the raw TMD for the data contained in the object. + return tmd_data_raw def get_title_region(self): """Gets the region of the TMD's associated title. diff --git a/src/libWiiPy/wad.py b/src/libWiiPy/wad.py index c605f27..5e3de1f 100644 --- a/src/libWiiPy/wad.py +++ b/src/libWiiPy/wad.py @@ -10,15 +10,9 @@ from .shared import align_value, pad_bytes_stream class WAD: """ - Creates a WAD object to parse the header of a WAD file and retrieve the data contained in it. - - Parameters - ---------- - wad : bytes - A bytes object containing the contents of a WAD file. + Creates a WAD object that allows for either loading and editing an existing WAD or creating a new WAD from raw data. """ - def __init__(self, wad): - self.wad = wad + def __init__(self): self.wad_hdr_size: int = 64 self.wad_type: str = "" self.wad_version: bytes = b'' @@ -37,17 +31,17 @@ class WAD: self.wad_tmd_data: bytes = b'' self.wad_content_data: bytes = b'' self.wad_meta_data: bytes = b'' - # 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. + def load(self, wad_data) -> None: + """Loads raw WAD data and sets all attributes of the WAD object. This allows for manipulating an already + existing WAD file. - Returns - ------- - none + Parameters + ---------- + wad_data : bytes + The data for the WAD you wish to load. """ - with io.BytesIO(self.wad) as wad_data: + with io.BytesIO(wad_data) 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. wad_data.seek(0x0) @@ -120,8 +114,8 @@ class WAD: self.wad_meta_data = wad_data.read(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. + """Dumps the WAD object into the raw WAD file. This allows for creating a WAD file from the data contained in + the WAD object. Returns ------- @@ -169,10 +163,9 @@ class WAD: 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 + wad_data_raw = wad_data.read() + # Return the raw WAD file for the data contained in the object. + return wad_data_raw def get_wad_type(self): """Gets the type of the WAD.