From 577d5a0efa0e387c4c8f02b4cefba4c14bc24ff8 Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Fri, 25 Apr 2025 14:45:38 -0400 Subject: [PATCH] Added rustii CLI command to replace content in a WAD Also added required library features to make this possible. Rust makes the whole "getting content's index from its CID" thing so much easier. --- src/bin/rustii/main.rs | 3 ++ src/bin/rustii/title/wad.rs | 82 ++++++++++++++++++++++++++++++++++++- src/title/content.rs | 14 +++++++ src/title/mod.rs | 4 +- 4 files changed, 100 insertions(+), 3 deletions(-) diff --git a/src/bin/rustii/main.rs b/src/bin/rustii/main.rs index fdd0519..f81c865 100644 --- a/src/bin/rustii/main.rs +++ b/src/bin/rustii/main.rs @@ -127,6 +127,9 @@ fn main() -> Result<()> { }, title::wad::Commands::Unpack { input, output } => { title::wad::unpack_wad(input, output)? + }, + title::wad::Commands::Set { input, content, output, identifier, r#type} => { + title::wad::set_wad(input, content, output, identifier, r#type)? } } }, diff --git a/src/bin/rustii/title/wad.rs b/src/bin/rustii/title/wad.rs index 2d86a68..fdcc323 100644 --- a/src/bin/rustii/title/wad.rs +++ b/src/bin/rustii/title/wad.rs @@ -37,6 +37,21 @@ pub enum Commands { input: String, /// The directory to extract the WAD to output: String + }, + /// 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 + content: String, + /// An optional output path; defaults to overwriting input WAD file + #[arg(short, long)] + output: Option, + /// An optional new type for the content, can be "Normal", "Shared", or "DLC" + #[arg(short, long)] + r#type: Option, + #[command(flatten)] + identifier: ContentIdentifier, } } @@ -55,6 +70,18 @@ pub struct ConvertTargets { vwii: bool, } +#[derive(Args)] +#[clap(next_help_heading = "Content Identifier")] +#[group(multiple = false, required = true)] +pub struct ContentIdentifier { + /// The index of the content to replace + #[arg(short, long)] + index: Option, + /// The Content ID of the content to replace + #[arg(short, long)] + cid: Option, +} + enum Target { Retail, Dev, @@ -94,7 +121,7 @@ pub fn convert_wad(input: &str, target: &ConvertTargets, output: &Option Target::Vwii => PathBuf::from(format!("{}_vWii.wad", in_path.file_stem().unwrap().to_str().unwrap())), } }; - let mut title = title::Title::from_bytes(fs::read(in_path)?.as_slice()).with_context(|| "The provided WAD file could not be parsed, and is likely invalid.")?; + 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.")?; // Bail if the WAD is already using the selected encryption. if matches!(target, Target::Dev) && title.ticket.is_dev() { bail!("This is already a development WAD!"); @@ -243,3 +270,56 @@ pub fn unpack_wad(input: &str, output: &str) -> Result<()> { println!("WAD file unpacked!"); Ok(()) } + +pub fn set_wad(input: &str, content: &str, output: &Option, identifier: &ContentIdentifier, 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()); + } + // 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 { + in_path.to_path_buf() + }; + // Load the WAD and parse the new type, if one was specified. + 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 mut target_type: Option = None; + if ctype.is_some() { + target_type = match ctype.clone().unwrap().to_ascii_lowercase().as_str() { + "normal" => Some(tmd::ContentType::Normal), + "shared" => Some(tmd::ContentType::Shared), + "dlc" => Some(tmd::ContentType::DLC), + _ => bail!("The specified content type \"{}\" is invalid!", ctype.clone().unwrap()), + }; + } + // Parse the identifier passed to choose how to do the find and replace. + if identifier.index.is_some() { + match title.set_content(&new_content, identifier.index.unwrap(), None, target_type) { + Err(title::TitleError::Content(content::ContentError::IndexOutOfRange { index, max })) => { + bail!("The specified index {} does not exist in this WAD! The maximum index is {}.", index, max) + }, + Err(e) => bail!("An unknown error occurred while setting the new content: {e}"), + Ok(_) => (), + } + 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 replaced 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.set_content(&new_content, index, None, target_type).with_context(|| "An unknown error occurred while setting the new content.")?; + 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 replaced content with Content ID \"{}\" ({}) in WAD file \"{}\".", identifier.cid.clone().unwrap(), cid, out_path.display()); + } + Ok(()) +} diff --git a/src/title/content.rs b/src/title/content.rs index 5e13bdc..43ad82b 100644 --- a/src/title/content.rs +++ b/src/title/content.rs @@ -110,6 +110,20 @@ impl ContentRegion { } Ok(buf) } + + /// Gets the index of content using its Content ID. + pub fn get_index_from_cid(&self, cid: u32) -> Result { + // Use fancy Rust find and map methods to find the index matching the provided CID. Take + // that libWiiPy! + let content_index = self.content_records.iter() + .find(|record| record.content_id == cid) + .map(|record| record.index); + if let Some(index) = content_index { + Ok(index as usize) + } else { + Err(ContentError::CIDNotFound(cid)) + } + } /// Gets the encrypted content file from the ContentRegion at the specified index. pub fn get_enc_content_by_index(&self, index: usize) -> Result, ContentError> { diff --git a/src/title/mod.rs b/src/title/mod.rs index 590e8e2..b36dbbd 100644 --- a/src/title/mod.rs +++ b/src/title/mod.rs @@ -135,8 +135,8 @@ impl Title { /// Sets the content at the specified index to the provided decrypted content. This content will /// have its size and hash saved into the matching record. Optionally, a new Content ID or /// content type can be provided, with the existing values being preserved by default. - pub fn set_content(&mut self, content: &[u8], index: usize) -> Result<(), TitleError> { - self.content.set_content(content, index, None, None, self.ticket.dec_title_key())?; + pub fn set_content(&mut self, content: &[u8], index: usize, cid: Option, content_type: Option) -> Result<(), TitleError> { + self.content.set_content(content, index, cid, content_type, self.ticket.dec_title_key())?; self.tmd.content_records = self.content.content_records.clone(); Ok(()) }