diff --git a/modules/title/nus.py b/modules/title/nus.py index e7a2dd1..d13a04a 100644 --- a/modules/title/nus.py +++ b/modules/title/nus.py @@ -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}\"!") diff --git a/modules/title/wad.py b/modules/title/wad.py index 7027e8e..211c9b9 100644 --- a/modules/title/wad.py +++ b/modules/title/wad.py @@ -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}\"!")