Redesigned how U8 archives are represented in memory

This replaces the old 1D array with an actual directory tree that can be used to make packing, unpacking, and editing U8 archives much much easier than the old libWiiPy implementation.

src/archive/u8.rs

@@ -3,13 +3,18 @@
 //
 // Implements the structures and methods required for parsing U8 archives.
 
+use std::cell::RefCell;
 use std::io::{Cursor, Read, Seek, SeekFrom, Write};
-use std::path::Path;
+use std::rc::{Rc, Weak};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 use thiserror::Error;
 
 #[derive(Debug, Error)]
 pub enum U8Error {
+    #[error("the requested item could not be found in this U8 archive")]
+    ItemNotFound(String),
+    #[error("found invalid node type {0} while processing node at index {1}")]
+    InvalidNodeType(u8, usize),
     #[error("invalid file name at offset {0}")]
     InvalidFileName(u64),
     #[error("this does not appear to be a U8 archive (missing magic number)")]
@@ -18,6 +23,84 @@ pub enum U8Error {
     IO(#[from] std::io::Error),
 }
 
+#[derive(Clone, Debug)]
+pub struct U8Directory {
+    pub name: String,
+    pub parent: Option<Weak<RefCell<U8Directory>>>,
+    pub dirs: Vec<Rc<RefCell<U8Directory>>>,
+    pub files: Vec<Rc<RefCell<U8File>>>,
+}
+
+#[derive(Clone, Debug)]
+pub struct U8File {
+    pub name: String,
+    pub data: Vec<u8>,
+    pub parent: Option<Weak<RefCell<U8Directory>>>,
+}
+
+impl U8Directory {
+    pub fn new(name: String) -> Rc<RefCell<Self>> {
+        Rc::new(RefCell::new(Self {
+            name,
+            parent: None,
+            dirs: Vec::new(),
+            files: Vec::new(),
+        }))
+    }
+
+    pub fn add_dir(parent: &Rc<RefCell<Self>>, child: Rc<RefCell<Self>>) {
+        child.borrow_mut().parent = Some(Rc::downgrade(parent));
+        parent.borrow_mut().dirs.push(child);
+    }
+
+    pub fn add_file(parent: &Rc<RefCell<Self>>, file: Rc<RefCell<U8File>>) {
+        file.borrow_mut().parent = Some(Rc::downgrade(parent));
+        parent.borrow_mut().files.push(file);
+    }
+
+    pub fn get_parent(&self) -> Option<Rc<RefCell<U8Directory>>> {
+        self.parent.as_ref()?.upgrade()
+    }
+
+    pub fn get_child_dir(parent: &Rc<RefCell<U8Directory>>, name: &str) -> Option<Rc<RefCell<U8Directory>>> {
+        parent.borrow().dirs.iter()
+            .find(|dir| dir.borrow().name == name)
+            .map(Rc::clone)
+    }
+    
+    fn count_recursive(dir: &Rc<RefCell<U8Directory>>, count: &mut usize) {
+        *count += dir.borrow().files.len();
+        for dir in dir.borrow().dirs.iter() {
+            *count += 1;
+            Self::count_recursive(dir, count);
+        }
+    }
+    
+    pub fn count(&self) -> usize {
+        let mut count: usize = 1;
+        count += self.files.len();
+        for dir in &self.dirs {
+            count += 1;
+            Self::count_recursive(dir, &mut count);
+        }
+        count
+    }
+}
+
+impl U8File {
+    pub fn new(name: String, data: Vec<u8>) -> Rc<RefCell<Self>> {
+        Rc::new(RefCell::new(Self {
+            name,
+            data,
+            parent: None,
+        }))
+    }
+
+    pub fn get_parent(&self) -> Option<Rc<RefCell<U8Directory>>> {
+        self.parent.as_ref()?.upgrade()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct U8Node {
     pub node_type: u8,
@@ -26,15 +109,9 @@ pub struct U8Node {
     pub size: u32,
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
 pub struct U8Archive {
-    pub u8_nodes: Vec<U8Node>,
-    pub file_names: Vec<String>,
-    pub file_data: Vec<Vec<u8>>,
-    root_node_offset: u32,
-    header_size: u32,
-    data_offset: u32,
-    padding: [u8; 16],
+    pub node_tree: Rc<RefCell<U8Directory>>,
 }
 
 impl U8Archive {
@@ -73,11 +150,12 @@ impl U8Archive {
                 }
             }
         }
-        let root_node_offset = buf.read_u32::<BigEndian>()?;
-        let header_size = buf.read_u32::<BigEndian>()?;
-        let data_offset = buf.read_u32::<BigEndian>()?;
-        let mut padding = [0u8; 16];
-        buf.read_exact(&mut padding)?;
+        // We're skipping the following values:
+        // root_node_offset (u32): constant value, always 0x20
+        // header_size (u32): we don't need this because we already know how long the string table is
+        // data_offset (u32): we don't need this because nodes provide the absolute offset to their data
+        // padding (u8 * 16): it's padding, I have nothing to say about it
+        buf.seek(SeekFrom::Start(buf.position() + 28))?;
         // Manually read the root node, since we need its size anyway to know how many nodes there
         // are total.
         let root_node_type = buf.read_u8()?;
@@ -127,36 +205,92 @@ impl U8Archive {
                 file_data.push(Vec::new());
             }
         }
+        // Now that we have all the data loaded out of the file, assemble the tree of U8Items that
+        // provides an actual map of the archive's data.
+        let node_tree = U8Directory::new(String::new());
+        let mut focused_node = Rc::clone(&node_tree);
+        // This is the order of directory nodes we've traversed down.
+        let mut parent_dirs: Vec<u32> = Vec::from([0]);
+        for i in 0..u8_nodes.len() {
+            match u8_nodes[i].node_type {
+                1 => {
+                    // Code for a directory node.
+                    if u8_nodes[i].name_offset != 0 {
+                        // If we're already at the correct level, push a new empty dir item to the
+                        // item we're currently working on.
+                        if u8_nodes[i].data_offset == *parent_dirs.last().unwrap() {
+                            parent_dirs.push(i as u32);
+                            U8Directory::add_dir(&focused_node, U8Directory::new(file_names[i].clone()));
+                            focused_node = U8Directory::get_child_dir(&focused_node, &file_names[i]).unwrap();
+                        }
+                        // Otherwise, go back up the path until we're at the correct level.
+                        else {
+                            while u8_nodes[i].data_offset != *parent_dirs.last().unwrap() {
+                                parent_dirs.pop();
+                                let parent = focused_node.as_ref().borrow().get_parent().unwrap();
+                                focused_node = parent;
+                            }
+                            parent_dirs.push(i as u32);
+                            // Rebuild current working directory, and make sure all directories in the
+                            // path exist.
+                            U8Directory::add_dir(&focused_node, U8Directory::new(file_names[i].clone()));
+                            focused_node = U8Directory::get_child_dir(&focused_node, &file_names[i]).unwrap()
+                        }
+                    }
+                },
+                0 => {
+                    // Code for a file node.
+                    U8Directory::add_file(&focused_node, U8File::new(file_names[i].clone(), file_data[i].clone()));
+                },
+                x => return Err(U8Error::InvalidNodeType(x, i))
+            }
+        }
         Ok(U8Archive {
-            u8_nodes,
-            file_names,
-            file_data,
-            root_node_offset,
-            header_size,
-            data_offset,
-            padding,
+            node_tree,
         })
     }
-
-    fn pack_dir() {
-        todo!();
-    }
-
-    pub fn from_dir(_input: &Path) -> Result<Self, U8Error> {
-        todo!();
+    
+    fn pack_dir_recursive(file_names: &mut Vec<String>, file_data: &mut Vec<Vec<u8>>, u8_nodes: &mut Vec<U8Node>, current_node: &Rc<RefCell<U8Directory>>) {
+        // For files, read their data into the file data list, add their name into the file name 
+        // list, then calculate the offset for their file name and create a new U8Node() for them. 
+        // 0 values for name/data offsets are temporary and are set later.
+        let parent_node = u8_nodes.len() - 1;
+        for file in &current_node.borrow().files {
+            file_names.push(file.borrow().name.clone());
+            file_data.push(file.borrow().data.clone());
+            u8_nodes.push(U8Node { node_type: 0, name_offset: 0, data_offset: 0, size: file_data[u8_nodes.len()].len() as u32});
+        }
+        // For directories, add their name to the file name list, add empty data to the file data 
+        // list, find the total number of files and directories inside the directory to calculate 
+        // the final node included in it, then recursively call this function again on that 
+        // directory to process it.
+        for dir in &current_node.borrow().dirs {
+            file_names.push(dir.borrow().name.clone());
+            file_data.push(Vec::new());
+            let max_node = u8_nodes.len() + current_node.borrow().count() + 1;
+            u8_nodes.push(U8Node { node_type: 1, name_offset: 0, data_offset: parent_node as u32, size: max_node as u32});
+            U8Archive::pack_dir_recursive(file_names, file_data, u8_nodes, dir)
+        }
     }
 
     /// Dumps the data in a U8Archive instance back into binary data that can be written to a file.
     pub fn to_bytes(&self) -> Result<Vec<u8>, U8Error> {
+        // We need to start by rebuilding a flat list of the nodes from the directory tree.
+        let mut file_names: Vec<String> = vec![String::new()];
+        let mut file_data: Vec<Vec<u8>> = vec![Vec::new()];
+        let mut u8_nodes: Vec<U8Node> = Vec::new();
+        u8_nodes.push(U8Node { node_type: 1, name_offset: 0, data_offset: 0, size: self.node_tree.borrow().count() as u32 });
+        let root_node = Rc::clone(&self.node_tree);
+        U8Archive::pack_dir_recursive(&mut file_names, &mut file_data, &mut u8_nodes, &root_node);
         // Header size starts at 0 because the header size starts with the nodes and does not
         // include the actual file header.
         let mut header_size: u32 = 0;
         // Add 12 bytes for each node, since that's how many bytes each one is made up of.
-        for _ in 0..self.u8_nodes.len() {
+        for _ in 0..u8_nodes.len() {
             header_size += 12;
         }
         // Add the number of bytes used for each file/folder name in the string table.
-        for file_name in &self.file_names {
+        for file_name in &file_names {
             header_size += file_name.len() as u32 + 1
         }
         // The initial data offset is equal to the file header (32 bytes) + node data aligned to
@@ -166,7 +300,6 @@ impl U8Archive {
         // Nintendo-made U8 archives don't necessarily do this?
         let mut current_data_offset = data_offset;
         let mut current_name_offset: u32 = 0;
-        let mut u8_nodes = self.u8_nodes.clone();
         for i in 0..u8_nodes.len() {
             if u8_nodes[i].node_type == 0 {
                 u8_nodes[i].data_offset = (current_data_offset + 31) & !31;
@@ -174,7 +307,7 @@ impl U8Archive {
             }
             // Calculate the name offsets, including the extra 1 for the NULL byte.
             u8_nodes[i].name_offset = current_name_offset;
-            current_name_offset += self.file_names[i].len() as u32 + 1
+            current_name_offset += file_names[i].len() as u32 + 1
         }
         // Begin writing file data.
         let mut buf: Vec<u8> = Vec::new();
@@ -182,7 +315,7 @@ impl U8Archive {
         buf.write_u32::<BigEndian>(0x20)?; // The root node offset is always 0x20.
         buf.write_u32::<BigEndian>(header_size)?;
         buf.write_u32::<BigEndian>(data_offset)?;
-        buf.write_all(&self.padding)?;
+        buf.write_all(&[0; 16])?;
         // Iterate over nodes and write them out.
         for node in &u8_nodes {
             buf.write_u8(node.node_type)?;
@@ -191,7 +324,7 @@ impl U8Archive {
             buf.write_u32::<BigEndian>(node.size)?;
         }
         // Iterate over file names with a null byte at the end.
-        for file_name in &self.file_names {
+        for file_name in &file_names {
             buf.write_all(file_name.as_bytes())?;
             buf.write_u8(b'\0')?;
         }
@@ -199,10 +332,26 @@ impl U8Archive {
         buf.resize((buf.len() + 63) & !63, 0);
         // Iterate over the file data and dump it. The file needs to be aligned to 32 bytes after
         // each write.
-        for data in &self.file_data {
+        for data in &file_data {
             buf.write_all(data)?;
             buf.resize((buf.len() + 31) & !31, 0);
         }
         Ok(buf)
     }
 }
+
+// pub fn print_full_tree(dir: &Rc<RefCell<U8Directory>>, indent: usize) {
+//     let prefix = "  ".repeat(indent);
+//     println!("{}D {}", prefix, dir.borrow().name);
+// 
+//     // Print subdirectories
+//     for subdir in &dir.borrow().dirs {
+//         print_full_tree(subdir, indent + 1);
+//     }
+// 
+//     // Print files
+//     for file in &dir.borrow().files {
+//         let file_name = &file.borrow().name;
+//         println!("{}  F {}", prefix, file_name);
+//     }
+// }

src/bin/playground/main.rs

@@ -14,31 +14,35 @@ fn main() {
     println!("num content records: {:?}", title.tmd.content_records.len());
     println!("first record data: {:?}", title.tmd.content_records.first().unwrap());
     println!("TMD is fakesigned: {:?}",title.tmd.is_fakesigned());
-
+    
     println!("title version from ticket is: {:?}", title.ticket.title_version);
     println!("title key (enc): {:?}", title.ticket.title_key);
     println!("title key (dec): {:?}", title.ticket.dec_title_key());
     println!("ticket is fakesigned: {:?}", title.ticket.is_fakesigned());
     
     println!("title is fakesigned: {:?}", title.is_fakesigned());
-
+    
     println!("wad header: {:?}", wad.header);
     
     let cert_chain = &title.cert_chain;
     println!("cert chain OK");
     let result = cert::verify_ca_cert(&cert_chain.ca_cert()).unwrap();
     println!("CA cert {} verified successfully: {}", cert_chain.ca_cert().child_cert_identity(), result);
-
+    
     let result = cert::verify_child_cert(&cert_chain.ca_cert(), &cert_chain.tmd_cert()).unwrap();
     println!("TMD cert {} verified successfully: {}", cert_chain.tmd_cert().child_cert_identity(), result);
     let result = cert::verify_tmd(&cert_chain.tmd_cert(), &title.tmd).unwrap();
     println!("TMD verified successfully: {}", result);
-
+    
     let result = cert::verify_child_cert(&cert_chain.ca_cert(), &cert_chain.ticket_cert()).unwrap();
     println!("Ticket cert {} verified successfully: {}", cert_chain.ticket_cert().child_cert_identity(), result);
     let result = cert::verify_ticket(&cert_chain.ticket_cert(), &title.ticket).unwrap();
     println!("Ticket verified successfully: {}", result);
-
+    
     let result = title.verify().unwrap();
     println!("full title verified successfully: {}", result);
+    // let mut u8_archive = u8::U8Archive::from_bytes(&fs::read("00000001.app").unwrap()).unwrap();
+    // println!("files and dirs counted: {}", u8_archive.node_tree.borrow().count());
+    // fs::write("outfile.arc", u8_archive.to_bytes().unwrap()).unwrap();
+    // println!("re-written");
 }

src/bin/rustii/archive/ash.rs

@@ -45,7 +45,7 @@ pub fn decompress_ash(input: &str, output: &Option<String>) -> Result<()> {
     let out_path = if output.is_some() {
         PathBuf::from(output.clone().unwrap())
     } else {
-        PathBuf::from(in_path).with_extension(format!("{}.out", in_path.extension().unwrap_or("".as_ref()).to_str().unwrap()))
+        PathBuf::from(in_path.file_name().unwrap()).with_extension(format!("{}.out", in_path.extension().unwrap_or("".as_ref()).to_str().unwrap()))
     };
     fs::write(out_path.clone(), decompressed)?;
     println!("Successfully decompressed ASH file to \"{}\"!", out_path.display());

src/bin/rustii/archive/lz77.rs

@@ -57,7 +57,7 @@ pub fn decompress_lz77(input: &str, output: &Option<String>) -> Result<()> {
     let out_path = if output.is_some() {
         PathBuf::from(output.clone().unwrap())
     } else {
-        PathBuf::from(in_path).with_extension(format!("{}.out", in_path.extension().unwrap_or("".as_ref()).to_str().unwrap()))
+        PathBuf::from(in_path.file_name().unwrap()).with_extension(format!("{}.out", in_path.extension().unwrap_or("".as_ref()).to_str().unwrap()))
     };
     fs::write(out_path.clone(), decompressed)?;
     println!("Successfully decompressed LZ77 file to \"{}\"!", out_path.display());

src/bin/rustii/archive/u8.rs

@@ -4,7 +4,9 @@
 // Code for the U8 packing/unpacking commands in the rustii CLI.
 
 use std::{str, fs};
+use std::cell::RefCell;
 use std::path::{Path, PathBuf};
+use std::rc::Rc;
 use anyhow::{bail, Context, Result};
 use clap::Subcommand;
 use rustii::archive::u8;
@@ -32,6 +34,20 @@ pub fn pack_u8_archive(_input: &str, _output: &str) -> Result<()> {
     todo!();
 }
 
+fn unpack_dir_recursive(dir: &Rc<RefCell<u8::U8Directory>>, out_path: PathBuf) -> Result<()> {
+    let out_path = out_path.join(&dir.borrow().name);
+    for file in &dir.borrow().files {
+        fs::write(out_path.join(&file.borrow().name), &file.borrow().data).with_context(|| format!("Failed to write output file \"{}\".", &file.borrow().name))?;
+    }
+    for dir in &dir.borrow().dirs {
+        if !out_path.join(&dir.borrow().name).exists() {
+            fs::create_dir(out_path.join(&dir.borrow().name)).with_context(|| format!("The output directory \"{}\" could not be created.", out_path.display()))?;
+        }
+        unpack_dir_recursive(dir, out_path.clone())?;
+    }
+    Ok(())
+}
+
 pub fn unpack_u8_archive(input: &str, output: &str) -> Result<()> {
     let in_path = Path::new(input);
     if !in_path.exists() {
@@ -45,51 +61,10 @@ pub fn unpack_u8_archive(input: &str, output: &str) -> Result<()> {
     } else {
         fs::create_dir(&out_path).with_context(|| format!("The output directory \"{}\" could not be created.", out_path.display()))?;
     }
-    let u8_archive = u8::U8Archive::from_bytes(&fs::read(in_path).with_context(|| format!("Failed to open U8 archive \"{}\" for reading.", in_path.display()))?)?;
-    // This stores the path we're actively writing files to.
-    let mut current_dir = out_path.clone();
-    // This is the order of directory nodes we've traversed down.
-    let mut parent_dirs: Vec<u32> = Vec::from([0]);
-    for i in 0..u8_archive.u8_nodes.len() {
-        match u8_archive.u8_nodes[i].node_type {
-            1 => {
-                // Code for a directory node.
-                if u8_archive.u8_nodes[i].name_offset != 0 {
-                    // If we're already at the correct level, make a new directory and push it to
-                    // the parent_dirs vec.
-                    if u8_archive.u8_nodes[i].data_offset == *parent_dirs.last().unwrap() {
-                        current_dir = current_dir.join(&u8_archive.file_names[i]);
-                        if !current_dir.exists() {
-                            fs::create_dir(&current_dir).with_context(|| format!("Failed to create directory \"{}\".", current_dir.display()))?;
-                        }
-                        parent_dirs.push(i as u32);
-                    }
-                    // Otherwise, go back up the path until we're at the correct level.
-                    else {
-                        while u8_archive.u8_nodes[i].data_offset != *parent_dirs.last().unwrap() {
-                            parent_dirs.pop();
-                        }
-                        parent_dirs.push(i as u32);
-                        current_dir = out_path.clone();
-                        // Rebuild current working directory, and make sure all directories in the 
-                        // path exist.
-                        for dir in &parent_dirs {
-                            current_dir = current_dir.join(&u8_archive.file_names[*dir as usize]);
-                            if !current_dir.exists() {
-                                fs::create_dir(&current_dir).with_context(|| format!("Failed to create directory \"{}\".", current_dir.display()))?;
-                            }
-                        }
-                    }
-                }
-            },
-            0 => {
-                // Code for a file node.
-                fs::write(current_dir.join(&u8_archive.file_names[i]), &u8_archive.file_data[i])
-                    .with_context(|| format!("Failed to write file \"{}\" in directory \"{}\".", u8_archive.file_names[i], current_dir.display()))?;
-            },
-            _ => bail!("Node at index {} has an invalid type! U8 archive cannot be unpacked.", i)
-        }
-    }
+    // Extract the files and directories in the root, and then recurse over each directory to
+    // extract the files and directories they contain.
+    let u8_archive = u8::U8Archive::from_bytes(&fs::read(in_path).with_context(|| format!("Input file \"{}\" could not be read.", in_path.display()))?)?;
+    unpack_dir_recursive(&u8_archive.node_tree, out_path.clone())?;
     println!("Successfully unpacked U8 archive to directory \"{}\"!", out_path.display());
     Ok(())
 }