NinjaCheetah/WiiPy
1
0
Fork 0
mirror of https://github.com/NinjaCheetah/WiiPy.git synced 2025-04-27 05:41:01 -04:00
Code Issues Packages Projects Releases Wiki Activity

Optimized large amounts of file I/O code that was very long-winded before

Browse Source
This commit is contained in:
Campbell 2024-10-20 21:50:48 -04:00
parent 31635a8015
commit 1612d2ecb9
Signed by: NinjaCheetah
GPG Key ID: 670C282B3291D63D
2 changed files with 77 additions and 141 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
modules/title/nus.py
View File

@ -1,7 +1,6 @@
# "modules/title/nus.py" from WiiPy by NinjaCheetah # "modules/title/nus.py" from WiiPy by NinjaCheetah
# https://github.com/NinjaCheetah/WiiPy # https://github.com/NinjaCheetah/WiiPy
import os
import hashlib import hashlib
import pathlib import pathlib
import binascii import binascii
@ -37,12 +36,10 @@ def handle_nus_title(args):
if args.output is not None: if args.output is not None:
output_dir = pathlib.Path(args.output) output_dir = pathlib.Path(args.output)
if output_dir.exists(): if output_dir.exists():
if output_dir.is_dir() and next(os.scandir(output_dir), None): if output_dir.is_file():
raise ValueError("Output folder is not empty!")
elif output_dir.is_file():
raise ValueError("A file already exists with the provided directory name!") raise ValueError("A file already exists with the provided directory name!")
else: else:
os.mkdir(output_dir) output_dir.mkdir()
# Download the title from the NUS. This is done "manually" (as opposed to using download_title()) so that we can # Download the title from the NUS. This is done "manually" (as opposed to using download_title()) so that we can
# provide verbose output. # provide verbose output.
@ -62,9 +59,7 @@ def handle_nus_title(args):
title_version = title.tmd.title_version title_version = title.tmd.title_version
# Write out the TMD to a file. # Write out the TMD to a file.
if output_dir is not None: if output_dir is not None:
tmd_out = open(output_dir.joinpath("tmd." + str(title_version)), "wb") output_dir.joinpath("tmd." + str(title_version)).write_bytes(title.tmd.dump())
tmd_out.write(title.tmd.dump())
tmd_out.close()
# Download the ticket, if we can. # Download the ticket, if we can.
print(" - Downloading and parsing Ticket...") print(" - Downloading and parsing Ticket...")
@ -72,9 +67,7 @@ def handle_nus_title(args):
title.load_ticket(libWiiPy.title.download_ticket(tid, wiiu_endpoint=wiiu_nus_enabled)) title.load_ticket(libWiiPy.title.download_ticket(tid, wiiu_endpoint=wiiu_nus_enabled))
can_decrypt = True can_decrypt = True
if output_dir is not None: if output_dir is not None:
ticket_out = open(output_dir.joinpath("tik"), "wb") output_dir.joinpath("tik").write_bytes(title.ticket.dump())
ticket_out.write(title.ticket.dump())
ticket_out.close()
except ValueError: except ValueError:
# If libWiiPy returns an error, then no ticket is available. Log this, and disable options requiring a # If libWiiPy returns an error, then no ticket is available. Log this, and disable options requiring a
# ticket so that they aren't attempted later. # ticket so that they aren't attempted later.
@ -100,9 +93,7 @@ def handle_nus_title(args):
print(" - Done!") print(" - Done!")
# If we're supposed to be outputting to a folder, then write these files out. # If we're supposed to be outputting to a folder, then write these files out.
if output_dir is not None: if output_dir is not None:
enc_content_out = open(output_dir.joinpath(content_file_name), "wb") output_dir.joinpath(content_file_name).write_bytes(content_list[content])
enc_content_out.write(content_list[content])
enc_content_out.close()
title.content.content_list = content_list title.content.content_list = content_list
# Try to decrypt the contents for this title if a ticket was available. # Try to decrypt the contents for this title if a ticket was available.
@ -113,9 +104,7 @@ def handle_nus_title(args):
" (Content ID: " + str(title.tmd.content_records[content].content_id) + ")...") " (Content ID: " + str(title.tmd.content_records[content].content_id) + ")...")
dec_content = title.get_content_by_index(content) dec_content = title.get_content_by_index(content)
content_file_name = f"{title.tmd.content_records[content].content_id:08X}".lower() + ".app" content_file_name = f"{title.tmd.content_records[content].content_id:08X}".lower() + ".app"
dec_content_out = open(output_dir.joinpath(content_file_name), "wb") output_dir.joinpath(content_file_name).write_bytes(dec_content)
dec_content_out.write(dec_content)
dec_content_out.close()
else: else:
print("Title has no Ticket, so content will not be decrypted!") print("Title has no Ticket, so content will not be decrypted!")
@ -129,9 +118,7 @@ def handle_nus_title(args):
if wad_file.suffix != ".wad": if wad_file.suffix != ".wad":
wad_file = wad_file.with_suffix(".wad") wad_file = wad_file.with_suffix(".wad")
# Have libWiiPy dump the WAD, and write that data out. # Have libWiiPy dump the WAD, and write that data out.
file = open(wad_file, "wb") pathlib.Path(wad_file).write_bytes(title.dump_wad())
file.write(title.dump_wad())
file.close()
print("Downloaded title with Title ID \"" + args.tid + "\"!") print("Downloaded title with Title ID \"" + args.tid + "\"!")
@ -140,7 +127,6 @@ def handle_nus_content(args):
tid = args.tid tid = args.tid
cid = args.cid cid = args.cid
version = args.version version = args.version
out = args.output
if args.decrypt: if args.decrypt:
decrypt_content = True decrypt_content = True
else: else:
@ -151,23 +137,21 @@ def handle_nus_content(args):
try: try:
content_id = int.from_bytes(binascii.unhexlify(cid)) content_id = int.from_bytes(binascii.unhexlify(cid))
except binascii.Error: except binascii.Error:
print("Invalid Content ID! Content ID must be in format \"000000xx\"!") raise ValueError("Invalid Content ID! Content ID must be in format \"000000xx\"!")
return
# Use the supplied output path if one was specified, otherwise generate one using the Content ID. # Use the supplied output path if one was specified, otherwise generate one using the Content ID.
if out is None: if args.output is None:
content_file_name = f"{content_id:08X}".lower() content_file_name = f"{content_id:08X}".lower()
output_path = pathlib.Path(content_file_name) output_path = pathlib.Path(content_file_name)
else: else:
output_path = pathlib.Path(out) output_path = pathlib.Path(args.output)
# Try to download the content, and catch the ValueError libWiiPy will throw if it can't be found. # Try to download the content, and catch the ValueError libWiiPy will throw if it can't be found.
print("Downloading content with Content ID " + cid + "...") print("Downloading content with Content ID " + cid + "...")
try: try:
content_data = libWiiPy.title.download_content(tid, content_id) content_data = libWiiPy.title.download_content(tid, content_id)
except ValueError: except ValueError:
print("The Title ID or Content ID you specified could not be found!") raise ValueError("The Title ID or Content ID you specified could not be found!")
return
if decrypt_content is True: if decrypt_content is True:
# Ensure that a version was supplied, because we need the matching TMD for decryption to work. # Ensure that a version was supplied, because we need the matching TMD for decryption to work.
@ -209,46 +193,38 @@ def handle_nus_content(args):
raise ValueError("The decrypted content provided does not match the record at the provided index. \n" raise ValueError("The decrypted content provided does not match the record at the provided index. \n"
"Expected hash is: {}\n".format(content_hash) + "Expected hash is: {}\n".format(content_hash) +
"Actual hash is: {}".format(content_dec_hash)) "Actual hash is: {}".format(content_dec_hash))
file = open(output_path, "wb") output_path.write_bytes(content_dec)
file.write(content_dec)
file.close()
else: else:
file = open(output_path, "wb") output_path.write_bytes(content_data)
file.write(content_data)
file.close()
print("Downloaded content with Content ID \"" + cid + "\"!") print(f"Downloaded content with Content ID \"{cid}\"!")
def handle_nus_tmd(args): def handle_nus_tmd(args):
tid = args.tid tid = args.tid
version = args.version
out = args.output
# Check if --version was passed, because it'll be None if it wasn't. # Check if --version was passed, because it'll be None if it wasn't.
if args.version is not None: if args.version is not None:
try: try:
title_version = int(args.version) version = int(args.version)
except ValueError: except ValueError:
print("Enter a valid integer for the Title Version.") raise ValueError("Enter a valid integer for the TMD Version.")
return else:
version = None
# Use the supplied output path if one was specified, otherwise generate one using the Title ID. # Use the supplied output path if one was specified, otherwise generate one using the Title ID.
if out is None: if args.output is None:
output_path = pathlib.Path(tid + ".tmd") output_path = pathlib.Path(tid + ".tmd")
else: else:
output_path = pathlib.Path(out) output_path = pathlib.Path(args.output)
# Try to download the TMD, and catch the ValueError libWiiPy will throw if it can't be found. # Try to download the TMD, and catch the ValueError libWiiPy will throw if it can't be found.
print("Downloading TMD for title " + tid + "...") print(f"Downloading TMD for title {tid}...")
try: try:
tmd_data = libWiiPy.title.download_tmd(tid, version) tmd_data = libWiiPy.title.download_tmd(tid, version)
except ValueError: except ValueError:
print("The Title ID or version you specified could not be found!") raise ValueError("The Title ID or version you specified could not be found!")
return
file = open(output_path, "wb") output_path.write_bytes(tmd_data)
file.write(tmd_data)
file.close()
print("Downloaded TMD for title \"" + tid + "\"!") print(f"Downloaded TMD for title \"{tid}\"!")

100
modules/title/wad.py
View File

@ -20,14 +20,9 @@ def handle_wad_add(args):
if not content_path.exists(): if not content_path.exists():
raise FileNotFoundError(content_path) raise FileNotFoundError(content_path)
wad_file = open(input_path, 'rb')
title = libWiiPy.title.Title() title = libWiiPy.title.Title()
title.load_wad(wad_file.read()) title.load_wad(input_path.read_bytes())
wad_file.close() content_data = content_path.read_bytes()
content_file = open(content_path, 'rb')
content_data = content_file.read()
content_file.close()
# Prepare the CID so it's ready when we go to add this content to the WAD. # Prepare the CID so it's ready when we go to add this content to the WAD.
# We need to both validate that this is a real CID, and also that it isn't already taken by another content. # We need to both validate that this is a real CID, and also that it isn't already taken by another content.
@ -68,10 +63,7 @@ def handle_wad_add(args):
# Auto fakesign because we've edited the title. # Auto fakesign because we've edited the title.
title.fakesign() title.fakesign()
output_path.write_bytes(title.dump_wad())
out_file = open(output_path, 'wb')
out_file.write(title.dump_wad())
out_file.close()
print(f"Successfully added new content with Content ID \"{target_cid:08X}\" and type \"{target_type.name}\"!") print(f"Successfully added new content with Content ID \"{target_cid:08X}\" and type \"{target_type.name}\"!")
@ -94,8 +86,7 @@ def handle_wad_pack(args):
raise FileExistsError("More than one TMD file was found! Only one TMD can be packed into a WAD.") raise FileExistsError("More than one TMD file was found! Only one TMD can be packed into a WAD.")
elif len(tmd_list) == 0: elif len(tmd_list) == 0:
raise FileNotFoundError("No TMD file found! Cannot pack WAD.") raise FileNotFoundError("No TMD file found! Cannot pack WAD.")
else: tmd_file = pathlib.Path(tmd_list[0])
tmd_file = tmd_list[0]
# Repeat the same process as above for all .tik files. # Repeat the same process as above for all .tik files.
ticket_list = list(input_path.glob('*.[tT][iI][kK]')) ticket_list = list(input_path.glob('*.[tT][iI][kK]'))
@ -103,8 +94,7 @@ def handle_wad_pack(args):
raise FileExistsError("More than one Ticket file was found! Only one Ticket can be packed into a WAD.") raise FileExistsError("More than one Ticket file was found! Only one Ticket can be packed into a WAD.")
elif len(ticket_list) == 0: elif len(ticket_list) == 0:
raise FileNotFoundError("No Ticket file found! Cannot pack WAD.") raise FileNotFoundError("No Ticket file found! Cannot pack WAD.")
else: ticket_file = pathlib.Path(ticket_list[0])
ticket_file = ticket_list[0]
# And one more time for all .cert files. # And one more time for all .cert files.
cert_list = list(input_path.glob('*.[cC][eE][rR][tT]')) cert_list = list(input_path.glob('*.[cC][eE][rR][tT]'))
@ -113,49 +103,39 @@ def handle_wad_pack(args):
"WAD.") "WAD.")
elif len(cert_list) == 0: elif len(cert_list) == 0:
raise FileNotFoundError("No certificate file found! Cannot pack WAD.") raise FileNotFoundError("No certificate file found! Cannot pack WAD.")
else: cert_file = pathlib.Path(cert_list[0])
cert_file = cert_list[0]
# Make sure that there's at least one content to pack. # Make sure that there's at least one content to pack.
content_files = list(input_path.glob("*.[aA][pP][pP]")) content_files = list(input_path.glob("*.[aA][pP][pP]"))
if not content_files: if not content_files:
raise FileNotFoundError("No contents found! Cannot pack WAD.") raise FileNotFoundError("No contents found! Cannot pack WAD.")
# Semi-hacky sorting method, but it works. Should maybe be changed eventually. # Semi-hacky sorting method, but it works. Should maybe be changed eventually.
content_files_ordered = [] content_files_ordered = []
for index in range(len(content_files)): for index in range(len(content_files)):
content_files_ordered.append(None) content_files_ordered.append(pathlib.Path(content_files[index]))
for content_file in content_files:
content_index = int(content_file.stem, 16)
content_files_ordered[content_index] = content_file
# Open the output file, and load all the component files that we've now verified we have into a libWiiPy Title() # Open the output file, and load all the component files that we've now verified we have into a libWiiPy Title()
# object. # object.
with open(output_path, "wb") as output_path:
title = libWiiPy.title.Title() title = libWiiPy.title.Title()
title.load_tmd(tmd_file.read_bytes())
title.load_tmd(open(tmd_file, "rb").read()) title.load_ticket(ticket_file.read_bytes())
title.load_ticket(open(ticket_file, "rb").read()) title.wad.set_cert_data(cert_file.read_bytes())
title.wad.set_cert_data(open(cert_file, "rb").read())
# Footers are not super common and are not required, so we don't care about one existing until we get to # Footers are not super common and are not required, so we don't care about one existing until we get to
# the step where we'd pack it. # the step where we'd pack it.
footer_file = list(input_path.glob("*.[fF][oO][oO][tT][eE][rR]"))[0] footer_file = pathlib.Path(list(input_path.glob("*.[fF][oO][oO][tT][eE][rR]"))[0])
if footer_file.exists(): if footer_file.exists():
title.wad.set_meta_data(open(footer_file, "rb").read()) title.wad.set_meta_data(footer_file.read_bytes())
# Method to ensure that the title's content records match between the TMD() and ContentRegion() objects. # Method to ensure that the title's content records match between the TMD() and ContentRegion() objects.
title.load_content_records() title.load_content_records()
# Iterate over every file in the content_files list, and set them in the Title(). # Iterate over every file in the content_files list, and set them in the Title().
for record in title.content.content_records: for index in range(title.content.num_contents):
index = title.content.content_records.index(record) dec_content = content_files_ordered[index].read_bytes()
dec_content = open(content_files_ordered[index], "rb").read()
title.set_content(dec_content, index) title.set_content(dec_content, index)
# Fakesign the TMD and Ticket using the trucha bug, if enabled. This is built-in in libWiiPy v0.4.1+. # Fakesign the TMD and Ticket using the trucha bug, if enabled. This is built-in in libWiiPy v0.4.1+.
if args.fakesign: if args.fakesign:
title.fakesign() title.fakesign()
output_path.write_bytes(title.dump_wad())
output_path.write(title.dump_wad())
print("WAD file packed!") print("WAD file packed!")
@ -170,10 +150,8 @@ def handle_wad_remove(args):
if not input_path.exists(): if not input_path.exists():
raise FileNotFoundError(input_path) raise FileNotFoundError(input_path)
wad_file = open(input_path, 'rb')
title = libWiiPy.title.Title() title = libWiiPy.title.Title()
title.load_wad(wad_file.read()) title.load_wad(input_path.read_bytes())
wad_file.close()
if args.index is not None: if args.index is not None:
# List indices in the title, and ensure that the target content index exists. # List indices in the title, and ensure that the target content index exists.
@ -185,9 +163,7 @@ def handle_wad_remove(args):
title.content.remove_content_by_index(args.index) title.content.remove_content_by_index(args.index)
# Auto fakesign because we've edited the title. # Auto fakesign because we've edited the title.
title.fakesign() title.fakesign()
out_file = open(output_path, 'wb') output_path.write_bytes(title.dump_wad())
out_file.write(title.dump_wad())
out_file.close()
print(f"Removed content at content index {args.index}!") print(f"Removed content at content index {args.index}!")
elif args.cid is not None: elif args.cid is not None:
@ -203,9 +179,7 @@ def handle_wad_remove(args):
title.content.remove_content_by_cid(target_cid) title.content.remove_content_by_cid(target_cid)
# Auto fakesign because we've edited the title. # Auto fakesign because we've edited the title.
title.fakesign() title.fakesign()
out_file = open(output_path, 'wb') output_path.write_bytes(title.dump_wad())
out_file.write(title.dump_wad())
out_file.close()
print(f"Removed content with Content ID \"{target_cid:08X}\"!") print(f"Removed content with Content ID \"{target_cid:08X}\"!")
@ -223,9 +197,8 @@ def handle_wad_set(args):
raise FileNotFoundError(content_path) raise FileNotFoundError(content_path)
title = libWiiPy.title.Title() title = libWiiPy.title.Title()
title.load_wad(open(input_path, "rb").read()) title.load_wad(input_path.read_bytes())
content_data = content_path.read_bytes()
content_data = open(content_path, "rb").read()
# Get the new type of the content, if one was specified. # Get the new type of the content, if one was specified.
if args.type is not None: if args.type is not None:
@ -254,7 +227,7 @@ def handle_wad_set(args):
title.set_content(content_data, args.index) title.set_content(content_data, args.index)
# Auto fakesign because we've edited the title. # Auto fakesign because we've edited the title.
title.fakesign() title.fakesign()
open(output_path, "wb").write(title.dump_wad()) output_path.write_bytes(title.dump_wad())
print(f"Replaced content at content index {args.index}!") print(f"Replaced content at content index {args.index}!")
@ -275,7 +248,7 @@ def handle_wad_set(args):
title.set_content(content_data, target_index) title.set_content(content_data, target_index)
# Auto fakesign because we've edited the title. # Auto fakesign because we've edited the title.
title.fakesign() title.fakesign()
open(output_path, "wb").write(title.dump_wad()) output_path.write_bytes(title.dump_wad())
print(f"Replaced content with Content ID \"{target_cid:08X}\"!") print(f"Replaced content with Content ID \"{target_cid:08X}\"!")
@ -287,37 +260,26 @@ def handle_wad_unpack(args):
raise FileNotFoundError(input_path) raise FileNotFoundError(input_path)
# Check if the output path already exists, and if it does, ensure that it is both a directory and empty. # Check if the output path already exists, and if it does, ensure that it is both a directory and empty.
if output_path.exists(): if output_path.exists():
# if output_path.is_dir() and next(os.scandir(output_path), None):
# raise ValueError("Output folder is not empty!")
if output_path.is_file(): if output_path.is_file():
raise ValueError("A file already exists with the provided directory name!") raise ValueError("A file already exists with the provided directory name!")
else: else:
os.mkdir(output_path) os.mkdir(output_path)
# Step through each component of a WAD and dump it to a file. # Step through each component of a WAD and dump it to a file.
with open(args.input, "rb") as wad_file:
title = libWiiPy.title.Title() title = libWiiPy.title.Title()
title.load_wad(wad_file.read()) title.load_wad(input_path.read_bytes())
cert_name = title.tmd.title_id + ".cert" cert_name = title.tmd.title_id + ".cert"
cert_out = open(output_path.joinpath(cert_name), "wb") output_path.joinpath(cert_name).write_bytes(title.wad.get_cert_data())
cert_out.write(title.wad.get_cert_data())
cert_out.close()
tmd_name = title.tmd.title_id + ".tmd" tmd_name = title.tmd.title_id + ".tmd"
tmd_out = open(output_path.joinpath(tmd_name), "wb") output_path.joinpath(tmd_name).write_bytes(title.wad.get_tmd_data())
tmd_out.write(title.wad.get_tmd_data())
tmd_out.close()
ticket_name = title.tmd.title_id + ".tik" ticket_name = title.tmd.title_id + ".tik"
ticket_out = open(output_path.joinpath(ticket_name), "wb") output_path.joinpath(ticket_name).write_bytes(title.wad.get_ticket_data())
ticket_out.write(title.wad.get_ticket_data())
ticket_out.close()
meta_name = title.tmd.title_id + ".footer" meta_name = title.tmd.title_id + ".footer"
meta_out = open(output_path.joinpath(meta_name), "wb") output_path.joinpath(meta_name).write_bytes(title.wad.get_meta_data())
meta_out.write(title.wad.get_meta_data())
meta_out.close()
# Skip validating hashes if -s/--skip-hash was passed. # Skip validating hashes if -s/--skip-hash was passed.
if args.skip_hash: if args.skip_hash:
@ -327,9 +289,7 @@ def handle_wad_unpack(args):
for content_file in range(0, title.tmd.num_contents): for content_file in range(0, title.tmd.num_contents):
content_file_name = f"{content_file:08X}".lower() + ".app" content_file_name = f"{content_file:08X}".lower() + ".app"
content_out = open(output_path.joinpath(content_file_name), "wb") output_path.joinpath(content_file_name).write_bytes(title.get_content_by_index(content_file, skip_hash))
content_out.write(title.get_content_by_index(content_file, skip_hash))
content_out.close()
print("WAD file unpacked!") print("WAD file unpacked!")
@ -355,7 +315,7 @@ def handle_wad_d2r(args):
title.ticket.title_key_enc = title_key_retail title.ticket.title_key_enc = title_key_retail
title.tmd.signature_issuer = "Root-CA00000001-CP00000004" + title.tmd.signature_issuer[26:] title.tmd.signature_issuer = "Root-CA00000001-CP00000004" + title.tmd.signature_issuer[26:]
title.fakesign() title.fakesign()
open(output_path, "wb").write(title.dump_wad()) output_path.write_bytes(title.dump_wad())
print(f"Successfully converted development WAD to retail WAD \"{output_path.name}\"!") print(f"Successfully converted development WAD to retail WAD \"{output_path.name}\"!")
@ -380,5 +340,5 @@ def handle_wad_r2d(args):
title.ticket.title_key_enc = title_key_dev title.ticket.title_key_enc = title_key_dev
title.tmd.signature_issuer = "Root-CA00000002-CP00000007" + title.tmd.signature_issuer[26:] title.tmd.signature_issuer = "Root-CA00000002-CP00000007" + title.tmd.signature_issuer[26:]
title.fakesign() title.fakesign()
open(output_path, "wb").write(title.dump_wad()) output_path.write_bytes(title.dump_wad())
print(f"Successfully converted retail WAD to development WAD \"{output_path.name}\"!") print(f"Successfully converted retail WAD to development WAD \"{output_path.name}\"!")