diff --git a/src/archive/u8.rs b/src/archive/u8.rs index 15a1e26..4073944 100644 --- a/src/archive/u8.rs +++ b/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>>, + pub dirs: Vec>>, + pub files: Vec>>, +} + +#[derive(Clone, Debug)] +pub struct U8File { + pub name: String, + pub data: Vec, + pub parent: Option>>, +} + +impl U8Directory { + pub fn new(name: String) -> Rc> { + Rc::new(RefCell::new(Self { + name, + parent: None, + dirs: Vec::new(), + files: Vec::new(), + })) + } + + pub fn add_dir(parent: &Rc>, child: Rc>) { + child.borrow_mut().parent = Some(Rc::downgrade(parent)); + parent.borrow_mut().dirs.push(child); + } + + pub fn add_file(parent: &Rc>, file: Rc>) { + file.borrow_mut().parent = Some(Rc::downgrade(parent)); + parent.borrow_mut().files.push(file); + } + + pub fn get_parent(&self) -> Option>> { + self.parent.as_ref()?.upgrade() + } + + pub fn get_child_dir(parent: &Rc>, name: &str) -> Option>> { + parent.borrow().dirs.iter() + .find(|dir| dir.borrow().name == name) + .map(Rc::clone) + } + + fn count_recursive(dir: &Rc>, 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) -> Rc> { + Rc::new(RefCell::new(Self { + name, + data, + parent: None, + })) + } + + pub fn get_parent(&self) -> Option>> { + 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, - pub file_names: Vec, - pub file_data: Vec>, - root_node_offset: u32, - header_size: u32, - data_offset: u32, - padding: [u8; 16], + pub node_tree: Rc>, } impl U8Archive { @@ -73,11 +150,12 @@ impl U8Archive { } } } - let root_node_offset = buf.read_u32::()?; - let header_size = buf.read_u32::()?; - let data_offset = buf.read_u32::()?; - 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 = 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 { - todo!(); + + fn pack_dir_recursive(file_names: &mut Vec, file_data: &mut Vec>, u8_nodes: &mut Vec, current_node: &Rc>) { + // 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 ¤t_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 ¤t_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, U8Error> { + // We need to start by rebuilding a flat list of the nodes from the directory tree. + let mut file_names: Vec = vec![String::new()]; + let mut file_data: Vec> = vec![Vec::new()]; + let mut u8_nodes: Vec = 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 = Vec::new(); @@ -182,7 +315,7 @@ impl U8Archive { buf.write_u32::(0x20)?; // The root node offset is always 0x20. buf.write_u32::(header_size)?; buf.write_u32::(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::(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>, 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); +// } +// } diff --git a/src/bin/playground/main.rs b/src/bin/playground/main.rs index 73ddc10..042cf00 100644 --- a/src/bin/playground/main.rs +++ b/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"); } diff --git a/src/bin/rustii/archive/ash.rs b/src/bin/rustii/archive/ash.rs index 630b182..9f474fa 100644 --- a/src/bin/rustii/archive/ash.rs +++ b/src/bin/rustii/archive/ash.rs @@ -45,7 +45,7 @@ pub fn decompress_ash(input: &str, output: &Option) -> 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()); diff --git a/src/bin/rustii/archive/lz77.rs b/src/bin/rustii/archive/lz77.rs index 9e49818..b4a3f5e 100644 --- a/src/bin/rustii/archive/lz77.rs +++ b/src/bin/rustii/archive/lz77.rs @@ -57,7 +57,7 @@ pub fn decompress_lz77(input: &str, output: &Option) -> 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()); diff --git a/src/bin/rustii/archive/u8.rs b/src/bin/rustii/archive/u8.rs index bfa9188..110ed0e 100644 --- a/src/bin/rustii/archive/u8.rs +++ b/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>, 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 = 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(¤t_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(¤t_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(()) }