Added rustii CLI commands to add and remove content from a WAD

Also added required library features to make this possible, again.

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_chacha 0.3.1",
- "rand_core",
+ "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]]

Cargo.toml

@@ -36,3 +36,4 @@ clap = { version = "4", features = ["derive"] }
 anyhow = "1"
 thiserror = "2"
 reqwest = { version = "0", features = ["blocking"] }
+rand = "0"

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::Commands::Remove { input, output, identifier } => {
-                    title::wad::unpack_wad(input, output)?
+                    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 */}

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<String>,
+        /// An optional Content ID for the new content; defaults to being randomly assigned
+        #[arg(short, long)]
+        cid: Option<String>,
+        /// An optional type for the new content, can be "Normal", "Shared", or "DLC"; defaults to
+        /// "Normal"
+        #[arg(short, long)]
+        r#type: Option<String>,
+    },
     /// 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
+    /// Remove content from a WAD file
-    Unpack {
+    Remove {
-        /// The path to the WAD to unpack
+        /// The path to the WAD file to modify
         input: String,
-        /// The directory to extract the WAD to
+        /// An optional output path; defaults to overwriting input WAD file
-        output: String
+        #[arg(short, long)]
+        output: Option<String>,
+        #[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<String>,
         #[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<usize>,
-    /// The Content ID of the content to replace
+    /// The Content ID of the target content
     #[arg(short, long)]
     cid: Option<String>,
 }
@@ -98,6 +126,60 @@ impl fmt::Display for Target {
     }
 }
 
+pub fn add_wad(input: &str, content: &str, output: &Option<String>, cid: &Option<String>, ctype: &Option<String>) -> 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<String>) -> 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<String>, 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 out_path = if output.is_some() {
-    let title = title::Title::from_bytes(&wad_file).with_context(|| format!("The provided WAD file \"{}\" appears to be invalid.", in_path.display()))?;
+        PathBuf::from(output.clone().unwrap()).with_extension("wad")
-    let tid = hex::encode(title.tmd.title_id);
+    } else {
-    // Create output directory if it doesn't exist.
+        in_path.to_path_buf()
-    let out_path = Path::new(output);
+    };
-    if !out_path.exists() {
+    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.")?;
-        fs::create_dir(out_path).with_context(|| format!("The output directory \"{}\" could not be created.", out_path.display()))?;
+    // 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<String>, 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<String>, 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(())
+}

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<Vec<u8>>, content_records: Vec<ContentRecord>) -> Result<Self, ContentError> {
         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(())
+    }
 }

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.