From 2755364472c459e4a93d6f25b4b2be525fe74d6f Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Wed, 5 Jun 2024 22:02:19 -0400 Subject: [PATCH] Rewrote U8 extraction code entirely, now handles U8 files in all cases --- src/libWiiPy/archive/u8.py | 75 +++++++++++++++++++++++++++++--------- 1 file changed, 57 insertions(+), 18 deletions(-) diff --git a/src/libWiiPy/archive/u8.py b/src/libWiiPy/archive/u8.py index c5443ae..314244d 100644 --- a/src/libWiiPy/archive/u8.py +++ b/src/libWiiPy/archive/u8.py @@ -5,6 +5,7 @@ import io import os +import pathlib from dataclasses import dataclass from typing import List from src.libWiiPy.shared import align_value @@ -159,30 +160,68 @@ class U8Archive: def extract_u8(u8_data, output_folder) -> None: - if os.path.isdir(output_folder): + """ + Extracts the provided U8 archive file data into the provided output folder path. Note that the folder must not + already exist to ensure that the output can correctly represent the file structure of the original U8 archive. + + Parameters + ---------- + u8_data : bytes + The data for the U8 file to extract. + output_folder : str + The path to a new folder to extract the archive to. + """ + output_folder = pathlib.Path(output_folder) + if pathlib.Path.is_dir(output_folder): raise ValueError("Output folder already exists!") os.mkdir(output_folder) # Create a new U8Archive object and load the provided U8 file data into it. u8_archive = U8Archive() u8_archive.load(u8_data) - # TODO: Comment this - # Also TODO: You can go more than two layers! Really should've checked that more before assuming it was the case. - current_dir = "" + # This variable stores the path of the directory we're currently processing. + current_dir = output_folder + # This variable stores the final nodes for every directory we've entered, and is used to handle the recursion of + # those directories to ensure that everything gets where it belongs. + directory_recursion = [0] + # Iterate over every node and extract the files and folders. for node in range(len(u8_archive.u8_node_list)): - if u8_archive.u8_node_list[node].name_offset != 0: - if u8_archive.u8_node_list[node].type == 256: - if u8_archive.u8_node_list[node].data_offset == 0: - os.mkdir(os.path.join(output_folder, u8_archive.file_name_list[node])) - current_dir = u8_archive.file_name_list[node] - elif u8_archive.u8_node_list[node].data_offset < node: - lower_path = os.path.join(output_folder, current_dir) - os.mkdir(os.path.join(lower_path, u8_archive.file_name_list[node])) - current_dir = os.path.join(current_dir, u8_archive.file_name_list[node]) - elif u8_archive.u8_node_list[node].type == 0: - lower_path = os.path.join(output_folder, current_dir) - output_file = open(os.path.join(lower_path, u8_archive.file_name_list[node]), "wb") - output_file.write(u8_archive.file_data_list[node]) - output_file.close() + # Code for a directory node. Second check just ensures we ignore the root node. + if u8_archive.u8_node_list[node].type == 256 and u8_archive.u8_node_list[node].name_offset != 0: + # The size value for a directory node is the position of the last node in this directory, with the root node + # counting as node 1. + # If the current node is below the end of the current directory, create this directory inside the previous + # current directory and make the current. + if node + 1 < max(directory_recursion): + current_dir = current_dir.joinpath(u8_archive.file_name_list[node]) + os.mkdir(current_dir) + # If the current node is beyond the end of the current directory, we've followed that path all the way down, + # so reset back to the root directory and put our new directory there. + elif node + 1 > max(directory_recursion): + current_dir = output_folder.joinpath(u8_archive.file_name_list[node]) + os.mkdir(current_dir) + # This check is here just in case a directory ever ends with an empty directory and not a file. + elif node + 1 == max(directory_recursion): + current_dir = current_dir.parent + directory_recursion.pop() + # If the last node for the directory we just processed is new (which is always should be), add it to the + # recursion array. + if u8_archive.u8_node_list[node].size not in directory_recursion: + directory_recursion.append(u8_archive.u8_node_list[node].size) + # Code for a file node. + elif u8_archive.u8_node_list[node].type == 0: + # Write out the file to the current directory. + output_file = open(current_dir.joinpath(u8_archive.file_name_list[node]), "wb") + output_file.write(u8_archive.file_data_list[node]) + output_file.close() + # If this file is the final node for the current directory, pop() the recursion array and set the current + # directory to the parent of the previous current. + if node + 1 in directory_recursion: + current_dir = current_dir.parent + directory_recursion.pop() + # Code for a totally unrecognized node type, which should not happen. + elif u8_archive.u8_node_list[node].type != 0 and u8_archive.u8_node_list[node].type != 256: + raise ValueError("A node with an invalid type (" + str(u8_archive.u8_node_list[node].type) + ") was" + "found!") def pack_u8(input_path) -> bytes: