Allow creating blank objects for WADs, TMDs, Tickets, ContentRegions, and Titles to make WAD packing possible

This commit is contained in:
Campbell 2024-03-31 23:38:52 -04:00
parent 142a121fa9
commit 640ca91716
Signed by: NinjaCheetah
GPG Key ID: B547958AF96ED344
5 changed files with 89 additions and 102 deletions

View File

@ -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:
"""

View File

@ -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.

View File

@ -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:

View File

@ -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.

View File

@ -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.