From 277c5d64397827bebdee61dbb1a8411be36cf6d2 Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Sun, 27 Apr 2025 15:25:47 -0400 Subject: [PATCH] Added rustii CLI commands to add and remove content from a WAD Also added required library features to make this possible, again. --- Cargo.lock | 42 ++++++-- Cargo.toml | 1 + src/bin/rustii/main.rs | 12 ++- src/bin/rustii/title/wad.rs | 193 +++++++++++++++++++++++++++++------- src/title/content.rs | 59 ++++++++++- src/title/mod.rs | 10 ++ 6 files changed, 271 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fc913f2..058c2bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -921,7 +921,7 @@ dependencies = [ "num-integer", "num-iter", "num-traits", - "rand", + "rand 0.8.5", "smallvec", "zeroize", ] @@ -1108,8 +1108,18 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.3", ] [[package]] @@ -1119,7 +1129,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.3", ] [[package]] @@ -1131,6 +1151,15 @@ dependencies = [ "getrandom 0.2.15", ] +[[package]] +name = "rand_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +dependencies = [ + "getrandom 0.3.2", +] + [[package]] name = "regex" version = "1.11.1" @@ -1232,7 +1261,7 @@ dependencies = [ "num-traits", "pkcs1", "pkcs8", - "rand_core", + "rand_core 0.6.4", "sha2", "signature", "spki", @@ -1257,6 +1286,7 @@ dependencies = [ "clap", "glob", "hex", + "rand 0.9.1", "regex", "reqwest", "rsa", @@ -1439,7 +1469,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ "digest", - "rand_core", + "rand_core 0.6.4", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index dacbde6..aea40e7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,3 +36,4 @@ clap = { version = "4", features = ["derive"] } anyhow = "1" thiserror = "2" reqwest = { version = "0", features = ["blocking"] } +rand = "0" diff --git a/src/bin/rustii/main.rs b/src/bin/rustii/main.rs index f81c865..b5d6309 100644 --- a/src/bin/rustii/main.rs +++ b/src/bin/rustii/main.rs @@ -119,18 +119,24 @@ fn main() -> Result<()> { }, Some(Commands::Wad { command }) => { match command { + title::wad::Commands::Add { input, content, output, cid, r#type } => { + title::wad::add_wad(input, content, output, cid, r#type)? + }, title::wad::Commands::Convert { input, target, output } => { title::wad::convert_wad(input, target, output)? }, title::wad::Commands::Pack { input, output} => { title::wad::pack_wad(input, output)? }, - title::wad::Commands::Unpack { input, output } => { - title::wad::unpack_wad(input, output)? + title::wad::Commands::Remove { input, output, identifier } => { + title::wad::remove_wad(input, output, identifier)? }, title::wad::Commands::Set { input, content, output, identifier, r#type} => { title::wad::set_wad(input, content, output, identifier, r#type)? - } + }, + title::wad::Commands::Unpack { input, output } => { + title::wad::unpack_wad(input, output)? + }, } }, None => { /* Clap handles no passed command by itself */} diff --git a/src/bin/rustii/title/wad.rs b/src/bin/rustii/title/wad.rs index fdcc323..e6996c8 100644 --- a/src/bin/rustii/title/wad.rs +++ b/src/bin/rustii/title/wad.rs @@ -8,12 +8,30 @@ use std::path::{Path, PathBuf}; use anyhow::{bail, Context, Result}; use clap::{Subcommand, Args}; use glob::glob; +use rand::prelude::*; use rustii::title::{cert, crypto, tmd, ticket, content, wad}; use rustii::title; #[derive(Subcommand)] #[command(arg_required_else_help = true)] pub enum Commands { + /// Add new content to a WAD file + Add { + /// The path to the WAD file to modify + input: String, + /// The path to the new content to add + content: String, + /// An optional output path; defaults to overwriting input WAD file + #[arg(short, long)] + output: Option, + /// An optional Content ID for the new content; defaults to being randomly assigned + #[arg(short, long)] + cid: Option, + /// An optional type for the new content, can be "Normal", "Shared", or "DLC"; defaults to + /// "Normal" + #[arg(short, long)] + r#type: Option, + }, /// Re-encrypt a WAD file with a different key Convert { /// The path to the WAD to convert @@ -31,18 +49,21 @@ pub enum Commands { /// The name of the packed WAD file output: String }, - /// Unpack a WAD file into a directory - Unpack { - /// The path to the WAD to unpack + /// Remove content from a WAD file + Remove { + /// The path to the WAD file to modify input: String, - /// The directory to extract the WAD to - output: String + /// An optional output path; defaults to overwriting input WAD file + #[arg(short, long)] + output: Option, + #[command(flatten)] + identifier: ContentIdentifier, }, /// Replace existing content in a WAD file with new data Set { /// The path to the WAD file to modify input: String, - /// The new WAD content + /// The path to the new content to set content: String, /// An optional output path; defaults to overwriting input WAD file #[arg(short, long)] @@ -52,7 +73,14 @@ pub enum Commands { r#type: Option, #[command(flatten)] identifier: ContentIdentifier, - } + }, + /// Unpack a WAD file into a directory + Unpack { + /// The path to the WAD to unpack + input: String, + /// The directory to extract the WAD to + output: String + }, } #[derive(Args)] @@ -74,10 +102,10 @@ pub struct ConvertTargets { #[clap(next_help_heading = "Content Identifier")] #[group(multiple = false, required = true)] pub struct ContentIdentifier { - /// The index of the content to replace + /// The index of the target content #[arg(short, long)] index: Option, - /// The Content ID of the content to replace + /// The Content ID of the target content #[arg(short, long)] cid: Option, } @@ -98,6 +126,60 @@ impl fmt::Display for Target { } } +pub fn add_wad(input: &str, content: &str, output: &Option, cid: &Option, ctype: &Option) -> Result<()> { + let in_path = Path::new(input); + if !in_path.exists() { + bail!("Source WAD \"{}\" could not be found.", in_path.display()); + } + let content_path = Path::new(content); + if !content_path.exists() { + bail!("New content \"{}\" could not be found.", content_path.display()); + } + let out_path = if output.is_some() { + PathBuf::from(output.clone().unwrap()).with_extension("wad") + } else { + in_path.to_path_buf() + }; + // Load the WAD and parse the target type and Content ID. + let mut title = title::Title::from_bytes(&fs::read(in_path)?).with_context(|| "The provided WAD file could not be parsed, and is likely invalid.")?; + let new_content = fs::read(content_path)?; + let target_type = if ctype.is_some() { + match ctype.clone().unwrap().to_ascii_lowercase().as_str() { + "normal" => tmd::ContentType::Normal, + "shared" => tmd::ContentType::Shared, + "dlc" => tmd::ContentType::DLC, + _ => bail!("The specified content type \"{}\" is invalid!", ctype.clone().unwrap()), + } + } else { + tmd::ContentType::Normal + }; + let target_cid = if cid.is_some() { + let cid = u32::from_str_radix(cid.clone().unwrap().as_str(), 16).with_context(|| "The specified Content ID is invalid!")?; + if title.content.content_records.iter().any(|record| record.content_id == cid) { + bail!("The specified Content ID \"{:08X}\" is already being used in this WAD!", cid); + } + cid + } else { + // Generate a random CID if one wasn't specified, and ensure that it isn't already in use. + let mut rng = rand::rng(); + let mut cid: u32; + loop { + cid = rng.random_range(0..=0xFF); + if !title.content.content_records.iter().any(|record| record.content_id == cid) { + break; + } + } + cid + }; + title.add_content(&new_content, target_cid, target_type.clone()).with_context(|| "An unknown error occurred while setting the new content.")?; + title.tmd.content_records = title.content.content_records.clone(); + title.tmd.num_contents = title.content.num_contents; + title.fakesign().with_context(|| "An unknown error occurred while fakesigning the modified WAD.")?; + fs::write(&out_path, title.to_wad()?.to_bytes()?).with_context(|| "Could not open output file for writing.")?; + println!("Successfully added new content with Content ID \"{:08X}\" ({}) and type \"{}\" to WAD file \"{}\".", target_cid, target_cid, target_type, out_path.display()); + Ok(()) +} + pub fn convert_wad(input: &str, target: &ConvertTargets, output: &Option) -> Result<()> { let in_path = Path::new(input); if !in_path.exists() { @@ -239,35 +321,45 @@ pub fn pack_wad(input: &str, output: &str) -> Result<()> { Ok(()) } -pub fn unpack_wad(input: &str, output: &str) -> Result<()> { +pub fn remove_wad(input: &str, output: &Option, identifier: &ContentIdentifier) -> Result<()> { let in_path = Path::new(input); if !in_path.exists() { - bail!("Source WAD \"{}\" could not be found.", input); + bail!("Source WAD \"{}\" could not be found.", in_path.display()); } - let wad_file = fs::read(in_path).with_context(|| format!("Failed to open WAD file \"{}\" for reading.", in_path.display()))?; - let title = title::Title::from_bytes(&wad_file).with_context(|| format!("The provided WAD file \"{}\" appears to be invalid.", in_path.display()))?; - let tid = hex::encode(title.tmd.title_id); - // Create output directory if it doesn't exist. - let out_path = Path::new(output); - if !out_path.exists() { - fs::create_dir(out_path).with_context(|| format!("The output directory \"{}\" could not be created.", out_path.display()))?; + let out_path = if output.is_some() { + PathBuf::from(output.clone().unwrap()).with_extension("wad") + } else { + in_path.to_path_buf() + }; + let mut title = title::Title::from_bytes(&fs::read(in_path)?).with_context(|| "The provided WAD file could not be parsed, and is likely invalid.")?; + // Parse the identifier passed to choose how to find and remove the target. + // ...maybe don't take the above comment out of context + if identifier.index.is_some() { + title.content.remove_content(identifier.index.unwrap()).with_context(|| "The specified index does not exist in the provided WAD!")?; + // Sync the content records in the TMD with the modified ones in the ContentRegion. The fact + // that this is required is probably bad and should be addressed on the library side at some + // point. + title.tmd.content_records = title.content.content_records.clone(); + title.tmd.num_contents = title.content.num_contents; + println!("{:?}", title.tmd); + title.fakesign().with_context(|| "An unknown error occurred while fakesigning the modified WAD.")?; + fs::write(&out_path, title.to_wad()?.to_bytes()?).with_context(|| "Could not open output file for writing.")?; + println!("Successfully removed content at index {} in WAD file \"{}\".", identifier.index.unwrap(), out_path.display()); + } else if identifier.cid.is_some() { + let cid = u32::from_str_radix(identifier.cid.clone().unwrap().as_str(), 16).with_context(|| "The specified Content ID is invalid!")?; + let index = match title.content.get_index_from_cid(cid) { + Ok(index) => index, + Err(_) => bail!("The specified Content ID \"{}\" ({}) does not exist in this WAD!", identifier.cid.clone().unwrap(), cid), + }; + title.content.remove_content(index).with_context(|| "An unknown error occurred while removing content from the WAD.")?; + // Ditto. + title.tmd.content_records = title.content.content_records.clone(); + title.tmd.num_contents = title.content.num_contents; + println!("{:?}", title.tmd); + title.fakesign().with_context(|| "An unknown error occurred while fakesigning the modified WAD.")?; + fs::write(&out_path, title.to_wad()?.to_bytes()?).with_context(|| "Could not open output file for writing.")?; + println!("Successfully removed content with Content ID \"{}\" ({}) in WAD file \"{}\".", identifier.cid.clone().unwrap(), cid, out_path.display()); } - // Write out all WAD components. - let tmd_file_name = format!("{}.tmd", tid); - fs::write(Path::join(out_path, tmd_file_name.clone()), title.tmd.to_bytes()?).with_context(|| format!("Failed to open TMD file \"{}\" for writing.", tmd_file_name))?; - let ticket_file_name = format!("{}.tik", tid); - fs::write(Path::join(out_path, ticket_file_name.clone()), title.ticket.to_bytes()?).with_context(|| format!("Failed to open Ticket file \"{}\" for writing.", ticket_file_name))?; - let cert_file_name = format!("{}.cert", tid); - fs::write(Path::join(out_path, cert_file_name.clone()), title.cert_chain.to_bytes()?).with_context(|| format!("Failed to open certificate chain file \"{}\" for writing.", cert_file_name))?; - let meta_file_name = format!("{}.footer", tid); - fs::write(Path::join(out_path, meta_file_name.clone()), title.meta()).with_context(|| format!("Failed to open footer file \"{}\" for writing.", meta_file_name))?; - // Iterate over contents, decrypt them, and write them out. - for i in 0..title.tmd.num_contents { - let content_file_name = format!("{:08X}.app", title.content.content_records[i as usize].index); - let dec_content = title.get_content_by_index(i as usize).with_context(|| format!("Failed to unpack content with Content ID {:08X}.", title.content.content_records[i as usize].content_id))?; - fs::write(Path::join(out_path, content_file_name), dec_content).with_context(|| format!("Failed to open content file \"{:08X}.app\" for writing.", title.content.content_records[i as usize].content_id))?; - } - println!("WAD file unpacked!"); Ok(()) } @@ -280,7 +372,6 @@ pub fn set_wad(input: &str, content: &str, output: &Option, identifier: if !content_path.exists() { bail!("New content \"{}\" could not be found.", content_path.display()); } - // Get the output name now that we know the target, if one wasn't passed. let out_path = if output.is_some() { PathBuf::from(output.clone().unwrap()).with_extension("wad") } else { @@ -323,3 +414,35 @@ pub fn set_wad(input: &str, content: &str, output: &Option, identifier: } Ok(()) } + +pub fn unpack_wad(input: &str, output: &str) -> Result<()> { + let in_path = Path::new(input); + if !in_path.exists() { + bail!("Source WAD \"{}\" could not be found.", input); + } + let wad_file = fs::read(in_path).with_context(|| format!("Failed to open WAD file \"{}\" for reading.", in_path.display()))?; + let title = title::Title::from_bytes(&wad_file).with_context(|| format!("The provided WAD file \"{}\" appears to be invalid.", in_path.display()))?; + let tid = hex::encode(title.tmd.title_id); + // Create output directory if it doesn't exist. + let out_path = Path::new(output); + if !out_path.exists() { + fs::create_dir(out_path).with_context(|| format!("The output directory \"{}\" could not be created.", out_path.display()))?; + } + // Write out all WAD components. + let tmd_file_name = format!("{}.tmd", tid); + fs::write(Path::join(out_path, tmd_file_name.clone()), title.tmd.to_bytes()?).with_context(|| format!("Failed to open TMD file \"{}\" for writing.", tmd_file_name))?; + let ticket_file_name = format!("{}.tik", tid); + fs::write(Path::join(out_path, ticket_file_name.clone()), title.ticket.to_bytes()?).with_context(|| format!("Failed to open Ticket file \"{}\" for writing.", ticket_file_name))?; + let cert_file_name = format!("{}.cert", tid); + fs::write(Path::join(out_path, cert_file_name.clone()), title.cert_chain.to_bytes()?).with_context(|| format!("Failed to open certificate chain file \"{}\" for writing.", cert_file_name))?; + let meta_file_name = format!("{}.footer", tid); + fs::write(Path::join(out_path, meta_file_name.clone()), title.meta()).with_context(|| format!("Failed to open footer file \"{}\" for writing.", meta_file_name))?; + // Iterate over contents, decrypt them, and write them out. + for i in 0..title.tmd.num_contents { + let content_file_name = format!("{:08X}.app", title.content.content_records[i as usize].index); + let dec_content = title.get_content_by_index(i as usize).with_context(|| format!("Failed to unpack content with Content ID {:08X}.", title.content.content_records[i as usize].content_id))?; + fs::write(Path::join(out_path, content_file_name), dec_content).with_context(|| format!("Failed to open content file \"{:08X}.app\" for writing.", title.content.content_records[i as usize].content_id))?; + } + println!("WAD file unpacked!"); + Ok(()) +} diff --git a/src/title/content.rs b/src/title/content.rs index 43ad82b..531d6e4 100644 --- a/src/title/content.rs +++ b/src/title/content.rs @@ -6,7 +6,6 @@ use std::io::{Cursor, Read, Seek, SeekFrom, Write}; use sha1::{Sha1, Digest}; use thiserror::Error; -use crate::title::content::ContentError::MissingContents; use crate::title::tmd::{ContentRecord, ContentType}; use crate::title::crypto; use crate::title::crypto::encrypt_content; @@ -19,6 +18,10 @@ pub enum ContentError { MissingContents { required: usize, found: usize }, #[error("content with requested Content ID {0} could not be found")] CIDNotFound(u32), + #[error("the specified index {0} already exists in the content records")] + IndexAlreadyExists(u16), + #[error("the specified Content ID {0} already exists in the content records")] + CIDAlreadyExists(u32), #[error("content's hash did not match the expected value (was {hash}, expected {expected})")] BadHash { hash: String, expected: String }, #[error("content data is not in a valid format")] @@ -73,7 +76,7 @@ impl ContentRegion { /// digital Wii title from a vector of contents and the ContentRecords from a TMD. pub fn from_contents(contents: Vec>, content_records: Vec) -> Result { if contents.len() != content_records.len() { - return Err(MissingContents { required: content_records.len(), found: contents.len()}); + return Err(ContentError::MissingContents { required: content_records.len(), found: contents.len()}); } let mut content_region = Self::new(content_records)?; for i in 0..contents.len() { @@ -189,6 +192,10 @@ impl ContentRegion { self.content_records[index].content_size = content_size; self.content_records[index].content_hash = content_hash; if cid.is_some() { + // Make sure that the new CID isn't already in use. + if self.content_records.iter().any(|record| record.content_id == cid.unwrap()) { + return Err(ContentError::CIDAlreadyExists(cid.unwrap())); + } self.content_records[index].content_id = cid.unwrap(); } if content_type.is_some() { @@ -231,4 +238,52 @@ impl ContentRegion { self.set_enc_content(&content_enc, index, content_size, content_hash, cid, content_type)?; Ok(()) } + + /// Removes the content at the specified index from the content list and content records. This + /// may leave a gap in the indexes recorded in the content records, but this should not cause + /// issues on the Wii or with correctly implemented WAD parsers. + pub fn remove_content(&mut self, index: usize) -> Result<(), ContentError> { + if self.contents.get(index).is_none() || self.content_records.get(index).is_none() { + return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 }); + } + self.contents.remove(index); + self.content_records.remove(index); + self.num_contents -= 1; + Ok(()) + } + + /// Adds new encrypted content to the end of the content list and content records. The provided + /// Content ID, type, index, and decrypted hash will be added to the record. + pub fn add_enc_content(&mut self, content: &[u8], index: u16, cid: u32, content_type: ContentType, content_size: u64, content_hash: [u8; 20]) -> Result<(), ContentError> { + // Return an error if the specified index or CID already exist in the records. + if self.content_records.iter().any(|record| record.index == index) { + return Err(ContentError::IndexAlreadyExists(index)); + } + if self.content_records.iter().any(|record| record.content_id == cid) { + return Err(ContentError::CIDAlreadyExists(cid)); + } + self.contents.push(content.to_vec()); + self.content_records.push(ContentRecord { content_id: cid, index, content_type, content_size, content_hash }); + self.num_contents += 1; + Ok(()) + } + + /// Adds new decrypted content to the end of the content list and content records. The provided + /// Content ID and type will be added to the record alongside a hash of the decrypted data. An + /// index will be automatically assigned based on the highest index currently recorded in the + /// content records. + pub fn add_content(&mut self, content: &[u8], cid: u32, content_type: ContentType, title_key: [u8; 16]) -> Result<(), ContentError> { + let max_index = self.content_records.iter() + .max_by_key(|record| record.index) + .map(|record| record.index) + .unwrap_or(0); // This should be impossible, but I guess 0 is a safe value just in case? + let new_index = max_index + 1; + let content_size = content.len() as u64; + let mut hasher = Sha1::new(); + hasher.update(content); + let content_hash: [u8; 20] = hasher.finalize().into(); + let content_enc = encrypt_content(content, title_key, new_index, content_size); + self.add_enc_content(&content_enc, new_index, cid, content_type, content_size, content_hash)?; + Ok(()) + } } diff --git a/src/title/mod.rs b/src/title/mod.rs index b36dbbd..1bdd4df 100644 --- a/src/title/mod.rs +++ b/src/title/mod.rs @@ -140,6 +140,16 @@ impl Title { self.tmd.content_records = self.content.content_records.clone(); Ok(()) } + + /// Adds new decrypted content to the end of the content list and content records. The provided + /// Content ID and type will be added to the record alongside a hash of the decrypted data. An + /// index will be automatically assigned based on the highest index currently recorded in the + /// content records. + pub fn add_content(&mut self, content: &[u8], cid: u32, content_type: tmd::ContentType) -> Result<(), TitleError> { + self.content.add_content(content, cid, content_type, self.ticket.dec_title_key())?; + self.tmd.content_records = self.content.content_records.clone(); + Ok(()) + } /// Gets the installed size of the title, in bytes. Use the optional parameter "absolute" to set /// whether shared content should be included in this total or not.