From 96ace7154607f17998abf051bfa84cdc74607d7e Mon Sep 17 00:00:00 2001 From: NinjaCheetah <58050615+NinjaCheetah@users.noreply.github.com> Date: Wed, 23 Apr 2025 18:21:26 -0400 Subject: [PATCH] Added set_content/set_enc_content to load modified content into a ContentRegion This means that the rustii CLI wad pack command is now actually useful, as it supports loading modified content. --- src/bin/rustii/title/wad.rs | 8 ++++--- src/title/content.rs | 47 ++++++++++++++++++++++++++++++------- src/title/mod.rs | 11 ++++++++- 3 files changed, 53 insertions(+), 13 deletions(-) diff --git a/src/bin/rustii/title/wad.rs b/src/bin/rustii/title/wad.rs index 3375f4d..2d86a68 100644 --- a/src/bin/rustii/title/wad.rs +++ b/src/bin/rustii/title/wad.rs @@ -156,7 +156,7 @@ pub fn pack_wad(input: &str, output: &str) -> Result<()> { } else if tmd_files.len() > 1 { bail!("More than one TMD file found in the source directory."); } - let tmd = tmd::TMD::from_bytes(&fs::read(&tmd_files[0]).with_context(|| "Could not open TMD file for reading.")?) + let mut tmd = tmd::TMD::from_bytes(&fs::read(&tmd_files[0]).with_context(|| "Could not open TMD file for reading.")?) .with_context(|| "The provided TMD file appears to be invalid.")?; // Read Ticket file (only accept one file). let ticket_files: Vec = glob(&format!("{}/*.tik", in_path.display()))? @@ -189,9 +189,11 @@ pub fn pack_wad(input: &str, output: &str) -> Result<()> { let mut content_region = content::ContentRegion::new(tmd.content_records.clone())?; for content in tmd.content_records.clone() { let data = fs::read(format!("{}/{:08X}.app", in_path.display(), content.index)).with_context(|| format!("Could not open content file \"{:08X}.app\" for reading.", content.index))?; - content_region.load_content(&data, content.index as usize, tik.dec_title_key()) - .expect("failed to load content into ContentRegion, this is probably because content was modified which isn't supported yet"); + content_region.set_content(&data, content.index as usize, None, None, tik.dec_title_key()) + .with_context(|| "Failed to load content into the ContentRegion.")?; } + // Ensure that the TMD is modified with our potentially updated content records. + tmd.content_records = content_region.content_records.clone(); let wad = wad::WAD::from_parts(&cert_chain, &[], &tik, &tmd, &content_region, &footer).with_context(|| "An unknown error occurred while building a WAD from the input files.")?; // Write out WAD file. let mut out_path = PathBuf::from(output); diff --git a/src/title/content.rs b/src/title/content.rs index 5e2b910..5e13bdc 100644 --- a/src/title/content.rs +++ b/src/title/content.rs @@ -7,8 +7,9 @@ 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; +use crate::title::tmd::{ContentRecord, ContentType}; use crate::title::crypto; +use crate::title::crypto::encrypt_content; #[derive(Debug, Error)] pub enum ContentError { @@ -49,13 +50,7 @@ impl ContentRegion { } Some(*offset) })).take(content_records.len()).collect(); // Trims the extra final entry. - let total_content_size: u64 = content_records.iter().map(|x| (x.content_size + 63) & !63).sum(); // Parse the content blob and create a vector of vectors from it. - // Check that the content blob matches the total size of all the contents in the records. - if content_region_size != total_content_size as u32 { - println!("Content region size mismatch."); - //return Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid content blob for content records")); - } let mut contents: Vec> = Vec::with_capacity(num_contents as usize); let mut buf = Cursor::new(data); for i in 0..num_contents { @@ -165,7 +160,27 @@ impl ContentRegion { if index >= self.content_records.len() { return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 }); } - self.contents[index] = Vec::from(content); + self.contents[index] = content.to_vec(); + Ok(()) + } + + /// Sets the content at the specified index to the provided encrypted content. This requires + /// the size and hash of the original decrypted content to be known so that the appropriate + /// values can be set in the corresponding content record. Optionally, a new Content ID or + /// content type can be provided, with the existing values being preserved by default. + pub fn set_enc_content(&mut self, content: &[u8], index: usize, content_size: u64, content_hash: [u8; 20], cid: Option, content_type: Option) -> Result<(), ContentError> { + if index >= self.content_records.len() { + return Err(ContentError::IndexOutOfRange { index, max: self.content_records.len() - 1 }); + } + self.content_records[index].content_size = content_size; + self.content_records[index].content_hash = content_hash; + if cid.is_some() { + self.content_records[index].content_id = cid.unwrap(); + } + if content_type.is_some() { + self.content_records[index].content_type = content_type.unwrap(); + } + self.contents[index] = content.to_vec(); Ok(()) } @@ -184,8 +199,22 @@ impl ContentRegion { if result[..] != self.content_records[index].content_hash { return Err(ContentError::BadHash { hash: hex::encode(result), expected: hex::encode(self.content_records[index].content_hash) }); } - let content_enc = crypto::encrypt_content(content, title_key, self.content_records[index].index, self.content_records[index].content_size); + let content_enc = encrypt_content(content, title_key, self.content_records[index].index, self.content_records[index].content_size); self.contents[index] = content_enc; Ok(()) } + + /// 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. The + /// Title Key will be used to encrypt this content before it is stored. + pub fn set_content(&mut self, content: &[u8], index: usize, cid: Option, content_type: Option, title_key: [u8; 16]) -> Result<(), ContentError> { + 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, index as u16, content_size); + self.set_enc_content(&content_enc, index, content_size, content_hash, cid, content_type)?; + Ok(()) + } } diff --git a/src/title/mod.rs b/src/title/mod.rs index 6071003..590e8e2 100644 --- a/src/title/mod.rs +++ b/src/title/mod.rs @@ -131,6 +131,15 @@ impl Title { let content = self.content.get_content_by_cid(cid, self.ticket.dec_title_key())?; Ok(content) } + + /// 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())?; + 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. @@ -197,7 +206,7 @@ impl Title { self.tmd = tmd; } - pub fn set_content(&mut self, content: content::ContentRegion) { + pub fn set_content_region(&mut self, content: content::ContentRegion) { self.content = content; }