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

@@ -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<PathBuf> = 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())
+        content_region.set_content(&data, content.index as usize, None, None, tik.dec_title_key())
-            .expect("failed to load content into ContentRegion, this is probably because content was modified which isn't supported yet");
+            .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);

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<u8>> = 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<u32>, content_type: Option<ContentType>) -> 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<u32>, content_type: Option<ContentType>, 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(())
+    }
 }

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;
     }