mirror of
https://github.com/NinjaCheetah/WiiPy.git
synced 2025-04-26 13:21:01 -04:00
Optimized large amounts of file I/O code that was very long-winded before
This commit is contained in:
parent
31635a8015
commit
1612d2ecb9
GPG Key ID:
670C282B3291D63D
@ -1,7 +1,6 @@
|
||||
# "modules/title/nus.py" from WiiPy by NinjaCheetah
|
||||
# https://github.com/NinjaCheetah/WiiPy
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import pathlib
|
||||
import binascii
|
||||
@ -37,12 +36,10 @@ def handle_nus_title(args):
|
||||
if args.output is not None:
|
||||
output_dir = pathlib.Path(args.output)
|
||||
if output_dir.exists():
|
||||
if output_dir.is_dir() and next(os.scandir(output_dir), None):
|
||||
raise ValueError("Output folder is not empty!")
|
||||
elif output_dir.is_file():
|
||||
if output_dir.is_file():
|
||||
raise ValueError("A file already exists with the provided directory name!")
|
||||
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
|
||||
# provide verbose output.
|
||||
@ -62,9 +59,7 @@ def handle_nus_title(args):
|
||||
title_version = title.tmd.title_version
|
||||
# Write out the TMD to a file.
|
||||
if output_dir is not None:
|
||||
tmd_out = open(output_dir.joinpath("tmd." + str(title_version)), "wb")
|
||||
tmd_out.write(title.tmd.dump())
|
||||
tmd_out.close()
|
||||
output_dir.joinpath("tmd." + str(title_version)).write_bytes(title.tmd.dump())
|
||||
|
||||
# Download the ticket, if we can.
|
||||
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))
|
||||
can_decrypt = True
|
||||
if output_dir is not None:
|
||||
ticket_out = open(output_dir.joinpath("tik"), "wb")
|
||||
ticket_out.write(title.ticket.dump())
|
||||
ticket_out.close()
|
||||
output_dir.joinpath("tik").write_bytes(title.ticket.dump())
|
||||
except ValueError:
|
||||
# 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.
|
||||
@ -100,9 +93,7 @@ def handle_nus_title(args):
|
||||
print(" - Done!")
|
||||
# If we're supposed to be outputting to a folder, then write these files out.
|
||||
if output_dir is not None:
|
||||
enc_content_out = open(output_dir.joinpath(content_file_name), "wb")
|
||||
enc_content_out.write(content_list[content])
|
||||
enc_content_out.close()
|
||||
output_dir.joinpath(content_file_name).write_bytes(content_list[content])
|
||||
title.content.content_list = content_list
|
||||
|
||||
# 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) + ")...")
|
||||
dec_content = title.get_content_by_index(content)
|
||||
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")
|
||||
dec_content_out.write(dec_content)
|
||||
dec_content_out.close()
|
||||
output_dir.joinpath(content_file_name).write_bytes(dec_content)
|
||||
else:
|
||||
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":
|
||||
wad_file = wad_file.with_suffix(".wad")
|
||||
# Have libWiiPy dump the WAD, and write that data out.
|
||||
file = open(wad_file, "wb")
|
||||
file.write(title.dump_wad())
|
||||
file.close()
|
||||
pathlib.Path(wad_file).write_bytes(title.dump_wad())
|
||||
|
||||
print("Downloaded title with Title ID \"" + args.tid + "\"!")
|
||||
|
||||
@ -140,7 +127,6 @@ def handle_nus_content(args):
|
||||
tid = args.tid
|
||||
cid = args.cid
|
||||
version = args.version
|
||||
out = args.output
|
||||
if args.decrypt:
|
||||
decrypt_content = True
|
||||
else:
|
||||
@ -151,23 +137,21 @@ def handle_nus_content(args):
|
||||
try:
|
||||
content_id = int.from_bytes(binascii.unhexlify(cid))
|
||||
except binascii.Error:
|
||||
print("Invalid Content ID! Content ID must be in format \"000000xx\"!")
|
||||
return
|
||||
raise ValueError("Invalid Content ID! Content ID must be in format \"000000xx\"!")
|
||||
|
||||
# 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()
|
||||
output_path = pathlib.Path(content_file_name)
|
||||
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.
|
||||
print("Downloading content with Content ID " + cid + "...")
|
||||
try:
|
||||
content_data = libWiiPy.title.download_content(tid, content_id)
|
||||
except ValueError:
|
||||
print("The Title ID or Content ID you specified could not be found!")
|
||||
return
|
||||
raise ValueError("The Title ID or Content ID you specified could not be found!")
|
||||
|
||||
if decrypt_content is True:
|
||||
# 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"
|
||||
"Expected hash is: {}\n".format(content_hash) +
|
||||
"Actual hash is: {}".format(content_dec_hash))
|
||||
file = open(output_path, "wb")
|
||||
file.write(content_dec)
|
||||
file.close()
|
||||
output_path.write_bytes(content_dec)
|
||||
else:
|
||||
file = open(output_path, "wb")
|
||||
file.write(content_data)
|
||||
file.close()
|
||||
output_path.write_bytes(content_data)
|
||||
|
||||
print("Downloaded content with Content ID \"" + cid + "\"!")
|
||||
print(f"Downloaded content with Content ID \"{cid}\"!")
|
||||
|
||||
|
||||
def handle_nus_tmd(args):
|
||||
tid = args.tid
|
||||
version = args.version
|
||||
out = args.output
|
||||
|
||||
# Check if --version was passed, because it'll be None if it wasn't.
|
||||
if args.version is not None:
|
||||
try:
|
||||
title_version = int(args.version)
|
||||
version = int(args.version)
|
||||
except ValueError:
|
||||
print("Enter a valid integer for the Title Version.")
|
||||
return
|
||||
raise ValueError("Enter a valid integer for the TMD Version.")
|
||||
else:
|
||||
version = None
|
||||
|
||||
# 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")
|
||||
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.
|
||||
print("Downloading TMD for title " + tid + "...")
|
||||
print(f"Downloading TMD for title {tid}...")
|
||||
try:
|
||||
tmd_data = libWiiPy.title.download_tmd(tid, version)
|
||||
except ValueError:
|
||||
print("The Title ID or version you specified could not be found!")
|
||||
return
|
||||
raise ValueError("The Title ID or version you specified could not be found!")
|
||||
|
||||
file = open(output_path, "wb")
|
||||
file.write(tmd_data)
|
||||
file.close()
|
||||
output_path.write_bytes(tmd_data)
|
||||
|
||||
print("Downloaded TMD for title \"" + tid + "\"!")
|
||||
print(f"Downloaded TMD for title \"{tid}\"!")
|
||||
|
||||
@ -20,14 +20,9 @@ def handle_wad_add(args):
|
||||
if not content_path.exists():
|
||||
raise FileNotFoundError(content_path)
|
||||
|
||||
wad_file = open(input_path, 'rb')
|
||||
title = libWiiPy.title.Title()
|
||||
title.load_wad(wad_file.read())
|
||||
wad_file.close()
|
||||
|
||||
content_file = open(content_path, 'rb')
|
||||
content_data = content_file.read()
|
||||
content_file.close()
|
||||
title.load_wad(input_path.read_bytes())
|
||||
content_data = content_path.read_bytes()
|
||||
|
||||
# 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.
|
||||
@ -68,10 +63,7 @@ def handle_wad_add(args):
|
||||
|
||||
# Auto fakesign because we've edited the title.
|
||||
title.fakesign()
|
||||
|
||||
out_file = open(output_path, 'wb')
|
||||
out_file.write(title.dump_wad())
|
||||
out_file.close()
|
||||
output_path.write_bytes(title.dump_wad())
|
||||
|
||||
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.")
|
||||
elif len(tmd_list) == 0:
|
||||
raise FileNotFoundError("No TMD file found! Cannot pack WAD.")
|
||||
else:
|
||||
tmd_file = tmd_list[0]
|
||||
tmd_file = pathlib.Path(tmd_list[0])
|
||||
|
||||
# Repeat the same process as above for all .tik files.
|
||||
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.")
|
||||
elif len(ticket_list) == 0:
|
||||
raise FileNotFoundError("No Ticket file found! Cannot pack WAD.")
|
||||
else:
|
||||
ticket_file = ticket_list[0]
|
||||
ticket_file = pathlib.Path(ticket_list[0])
|
||||
|
||||
# And one more time for all .cert files.
|
||||
cert_list = list(input_path.glob('*.[cC][eE][rR][tT]'))
|
||||
@ -113,49 +103,39 @@ def handle_wad_pack(args):
|
||||
"WAD.")
|
||||
elif len(cert_list) == 0:
|
||||
raise FileNotFoundError("No certificate file found! Cannot pack WAD.")
|
||||
else:
|
||||
cert_file = cert_list[0]
|
||||
cert_file = pathlib.Path(cert_list[0])
|
||||
|
||||
# Make sure that there's at least one content to pack.
|
||||
content_files = list(input_path.glob("*.[aA][pP][pP]"))
|
||||
if not content_files:
|
||||
raise FileNotFoundError("No contents found! Cannot pack WAD.")
|
||||
|
||||
# Semi-hacky sorting method, but it works. Should maybe be changed eventually.
|
||||
content_files_ordered = []
|
||||
for index in range(len(content_files)):
|
||||
content_files_ordered.append(None)
|
||||
for content_file in content_files:
|
||||
content_index = int(content_file.stem, 16)
|
||||
content_files_ordered[content_index] = content_file
|
||||
content_files_ordered.append(pathlib.Path(content_files[index]))
|
||||
|
||||
# Open the output file, and load all the component files that we've now verified we have into a libWiiPy Title()
|
||||
# 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_ticket(ticket_file.read_bytes())
|
||||
title.wad.set_cert_data(cert_file.read_bytes())
|
||||
# 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.
|
||||
footer_file = pathlib.Path(list(input_path.glob("*.[fF][oO][oO][tT][eE][rR]"))[0])
|
||||
if footer_file.exists():
|
||||
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.
|
||||
title.load_content_records()
|
||||
# Iterate over every file in the content_files list, and set them in the Title().
|
||||
for index in range(title.content.num_contents):
|
||||
dec_content = content_files_ordered[index].read_bytes()
|
||||
title.set_content(dec_content, index)
|
||||
|
||||
title.load_tmd(open(tmd_file, "rb").read())
|
||||
title.load_ticket(open(ticket_file, "rb").read())
|
||||
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
|
||||
# the step where we'd pack it.
|
||||
footer_file = list(input_path.glob("*.[fF][oO][oO][tT][eE][rR]"))[0]
|
||||
if footer_file.exists():
|
||||
title.wad.set_meta_data(open(footer_file, "rb").read())
|
||||
# Method to ensure that the title's content records match between the TMD() and ContentRegion() objects.
|
||||
title.load_content_records()
|
||||
|
||||
# Iterate over every file in the content_files list, and set them in the Title().
|
||||
for record in title.content.content_records:
|
||||
index = title.content.content_records.index(record)
|
||||
dec_content = open(content_files_ordered[index], "rb").read()
|
||||
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+.
|
||||
if args.fakesign:
|
||||
title.fakesign()
|
||||
|
||||
output_path.write(title.dump_wad())
|
||||
# Fakesign the TMD and Ticket using the trucha bug, if enabled. This is built-in in libWiiPy v0.4.1+.
|
||||
if args.fakesign:
|
||||
title.fakesign()
|
||||
output_path.write_bytes(title.dump_wad())
|
||||
|
||||
print("WAD file packed!")
|
||||
|
||||
@ -170,10 +150,8 @@ def handle_wad_remove(args):
|
||||
if not input_path.exists():
|
||||
raise FileNotFoundError(input_path)
|
||||
|
||||
wad_file = open(input_path, 'rb')
|
||||
title = libWiiPy.title.Title()
|
||||
title.load_wad(wad_file.read())
|
||||
wad_file.close()
|
||||
title.load_wad(input_path.read_bytes())
|
||||
|
||||
if args.index is not None:
|
||||
# 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)
|
||||
# Auto fakesign because we've edited the title.
|
||||
title.fakesign()
|
||||
out_file = open(output_path, 'wb')
|
||||
out_file.write(title.dump_wad())
|
||||
out_file.close()
|
||||
output_path.write_bytes(title.dump_wad())
|
||||
print(f"Removed content at content index {args.index}!")
|
||||
|
||||
elif args.cid is not None:
|
||||
@ -203,9 +179,7 @@ def handle_wad_remove(args):
|
||||
title.content.remove_content_by_cid(target_cid)
|
||||
# Auto fakesign because we've edited the title.
|
||||
title.fakesign()
|
||||
out_file = open(output_path, 'wb')
|
||||
out_file.write(title.dump_wad())
|
||||
out_file.close()
|
||||
output_path.write_bytes(title.dump_wad())
|
||||
print(f"Removed content with Content ID \"{target_cid:08X}\"!")
|
||||
|
||||
|
||||
@ -223,9 +197,8 @@ def handle_wad_set(args):
|
||||
raise FileNotFoundError(content_path)
|
||||
|
||||
title = libWiiPy.title.Title()
|
||||
title.load_wad(open(input_path, "rb").read())
|
||||
|
||||
content_data = open(content_path, "rb").read()
|
||||
title.load_wad(input_path.read_bytes())
|
||||
content_data = content_path.read_bytes()
|
||||
|
||||
# Get the new type of the content, if one was specified.
|
||||
if args.type is not None:
|
||||
@ -254,7 +227,7 @@ def handle_wad_set(args):
|
||||
title.set_content(content_data, args.index)
|
||||
# Auto fakesign because we've edited the title.
|
||||
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}!")
|
||||
|
||||
|
||||
@ -275,7 +248,7 @@ def handle_wad_set(args):
|
||||
title.set_content(content_data, target_index)
|
||||
# Auto fakesign because we've edited the title.
|
||||
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}\"!")
|
||||
|
||||
|
||||
@ -287,49 +260,36 @@ def handle_wad_unpack(args):
|
||||
raise FileNotFoundError(input_path)
|
||||
# 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.is_dir() and next(os.scandir(output_path), None):
|
||||
# raise ValueError("Output folder is not empty!")
|
||||
if output_path.is_file():
|
||||
raise ValueError("A file already exists with the provided directory name!")
|
||||
else:
|
||||
os.mkdir(output_path)
|
||||
|
||||
# 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.load_wad(wad_file.read())
|
||||
title = libWiiPy.title.Title()
|
||||
title.load_wad(input_path.read_bytes())
|
||||
|
||||
cert_name = title.tmd.title_id + ".cert"
|
||||
cert_out = open(output_path.joinpath(cert_name), "wb")
|
||||
cert_out.write(title.wad.get_cert_data())
|
||||
cert_out.close()
|
||||
cert_name = title.tmd.title_id + ".cert"
|
||||
output_path.joinpath(cert_name).write_bytes(title.wad.get_cert_data())
|
||||
|
||||
tmd_name = title.tmd.title_id + ".tmd"
|
||||
tmd_out = open(output_path.joinpath(tmd_name), "wb")
|
||||
tmd_out.write(title.wad.get_tmd_data())
|
||||
tmd_out.close()
|
||||
tmd_name = title.tmd.title_id + ".tmd"
|
||||
output_path.joinpath(tmd_name).write_bytes(title.wad.get_tmd_data())
|
||||
|
||||
ticket_name = title.tmd.title_id + ".tik"
|
||||
ticket_out = open(output_path.joinpath(ticket_name), "wb")
|
||||
ticket_out.write(title.wad.get_ticket_data())
|
||||
ticket_out.close()
|
||||
ticket_name = title.tmd.title_id + ".tik"
|
||||
output_path.joinpath(ticket_name).write_bytes(title.wad.get_ticket_data())
|
||||
|
||||
meta_name = title.tmd.title_id + ".footer"
|
||||
meta_out = open(output_path.joinpath(meta_name), "wb")
|
||||
meta_out.write(title.wad.get_meta_data())
|
||||
meta_out.close()
|
||||
meta_name = title.tmd.title_id + ".footer"
|
||||
output_path.joinpath(meta_name).write_bytes(title.wad.get_meta_data())
|
||||
|
||||
# Skip validating hashes if -s/--skip-hash was passed.
|
||||
if args.skip_hash:
|
||||
skip_hash = True
|
||||
else:
|
||||
skip_hash = False
|
||||
# Skip validating hashes if -s/--skip-hash was passed.
|
||||
if args.skip_hash:
|
||||
skip_hash = True
|
||||
else:
|
||||
skip_hash = False
|
||||
|
||||
for content_file in range(0, title.tmd.num_contents):
|
||||
content_file_name = f"{content_file:08X}".lower() + ".app"
|
||||
content_out = open(output_path.joinpath(content_file_name), "wb")
|
||||
content_out.write(title.get_content_by_index(content_file, skip_hash))
|
||||
content_out.close()
|
||||
for content_file in range(0, title.tmd.num_contents):
|
||||
content_file_name = f"{content_file:08X}".lower() + ".app"
|
||||
output_path.joinpath(content_file_name).write_bytes(title.get_content_by_index(content_file, skip_hash))
|
||||
|
||||
print("WAD file unpacked!")
|
||||
|
||||
@ -355,7 +315,7 @@ def handle_wad_d2r(args):
|
||||
title.ticket.title_key_enc = title_key_retail
|
||||
title.tmd.signature_issuer = "Root-CA00000001-CP00000004" + title.tmd.signature_issuer[26:]
|
||||
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}\"!")
|
||||
|
||||
|
||||
@ -380,5 +340,5 @@ def handle_wad_r2d(args):
|
||||
title.ticket.title_key_enc = title_key_dev
|
||||
title.tmd.signature_issuer = "Root-CA00000002-CP00000007" + title.tmd.signature_issuer[26:]
|
||||
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}\"!")
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user