crates/dtmt/src/cmd/bundle/inject.rs

@@ -1,112 +1,297 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
+use std::str::FromStr as _;
 
-use clap::{value_parser, Arg, ArgMatches, Command};
+use clap::{value_parser, Arg, ArgAction, ArgMatches, Command};
-use color_eyre::eyre::{self, Context, Result};
+use color_eyre::eyre::{self, Context, OptionExt, Result};
 use color_eyre::Help;
-use sdk::Bundle;
+use path_slash::PathBufExt as _;
-use tokio::fs::{self, File};
+use sdk::murmur::IdString64;
-use tokio::io::AsyncReadExt;
+use sdk::{Bundle, BundleFile, BundleFileType};
+use tokio::fs;
 
 pub(crate) fn command_definition() -> Command {
     Command::new("inject")
-        .about("Inject a file into a bundle.")
+        .subcommand_required(true)
-        .arg(
+        .about("Inject a file into a bundle.\n\
-            Arg::new("replace")
+            Raw binary data can be used to directly replace the file's variant data blob without affecting the metadata.\n\
-                .help("The name of a file in the bundle whos content should be replaced.")
+            Alternatively, a compiler format may be specified, and a complete bundle file is created.")
-                .short('r')
-                .long("replace"),
-        )
         .arg(
             Arg::new("output")
                 .help(
                     "The path to write the changed bundle to. \
-                    If omitted, the input bundle will be overwritten.",
+                    If omitted, the input bundle will be overwritten.\n\
+                    Remember to add a `.patch_<NUMBER>` suffix if you also use '--patch'.",
                 )
                 .short('o')
                 .long("output")
                 .value_parser(value_parser!(PathBuf)),
         )
         .arg(
-            Arg::new("bundle")
+            Arg::new("patch")
-                .help("Path to the bundle to inject the file into.")
+                .help("Create a patch bundle. Optionally, a patch NUMBER may be specified as \
-                .required(true)
+                    '--patch=123'.\nThe maximum number is 999, the default is 1.\n\
-                .value_parser(value_parser!(PathBuf)),
+                    If `--output` is not specified, the `.patch_<NUMBER>` suffix is added to \
+                    the given bundle name.")
+                .short('p')
+                .long("patch")
+                .num_args(0..=1)
+                .require_equals(true)
+                .default_missing_value("1")
+                .value_name("NUMBER")
+                .value_parser(value_parser!(u16))
         )
         .arg(
-            Arg::new("file")
+            Arg::new("type")
-                .help("Path to the file to inject.")
+                .help("Compile the new file as the given TYPE. If omitted, the file type is \
-                .required(true)
+                    is guessed from the file extension.")
-                .value_parser(value_parser!(PathBuf)),
+                .value_name("TYPE")
         )
+        .subcommand(
+            Command::new("replace")
+                .about("Replace an existing file in the bundle")
+                .arg(
+                    Arg::new("variant")
+                        .help("In combination with '--raw', specify the variant index to replace.")
+                        .long("variant")
+                        .default_value("0")
+                        .value_parser(value_parser!(u8))
+                )
+                .arg(
+                    Arg::new("raw")
+                        .help("Insert the given file as raw binary data.\n\
+                            Cannot be used with '--patch'.")
+                        .long("raw")
+                        .action(ArgAction::SetTrue)
+                )
+                .arg(
+                    Arg::new("bundle")
+                        .help("Path to the bundle to inject the file into.")
+                        .required(true)
+                        .value_parser(value_parser!(PathBuf)),
+                )
+                .arg(
+                    Arg::new("bundle-file")
+                        .help("The name of a file in the bundle whose content should be replaced.")
+                        .required(true),
+                )
+                .arg(
+                    Arg::new("new-file")
+                        .help("Path to the file to inject.")
+                        .required(true)
+                        .value_parser(value_parser!(PathBuf)),
+                ),
+        )
+    // .subcommand(
+    //     Command::new("add")
+    //         .about("Add a new file to the bundle")
+    //         .arg(
+    //             Arg::new("new-file")
+    //                 .help("Path to the file to inject.")
+    //                 .required(true)
+    //                 .value_parser(value_parser!(PathBuf)),
+    //         )
+    //         .arg(
+    //             Arg::new("bundle")
+    //                 .help("Path to the bundle to inject the file into.")
+    //                 .required(true)
+    //                 .value_parser(value_parser!(PathBuf)),
+    //         ),
+    // )
 }
 
-#[tracing::instrument(skip_all)]
+#[tracing::instrument]
+async fn compile_file(
+    path: impl AsRef<Path> + std::fmt::Debug,
+    name: impl Into<IdString64> + std::fmt::Debug,
+    file_type: BundleFileType,
+) -> Result<BundleFile> {
+    let path = path.as_ref();
+
+    let file_data = fs::read(&path)
+        .await
+        .wrap_err_with(|| format!("Failed to read file '{}'", path.display()))?;
+    let _sjson = String::from_utf8(file_data)
+        .wrap_err_with(|| format!("Invalid UTF8 data in '{}'", path.display()))?;
+
+    let _root = path.parent().ok_or_eyre("File path has no parent")?;
+
+    eyre::bail!(
+        "Compilation for type '{}' is not implemented, yet",
+        file_type
+    )
+}
+
+#[tracing::instrument(
+    skip_all,
+    fields(
+        bundle_path = tracing::field::Empty,
+        in_file_path = tracing::field::Empty,
+        output_path = tracing::field::Empty,
+        target_name = tracing::field::Empty,
+        file_type = tracing::field::Empty,
+        raw = tracing::field::Empty,
+    )
+)]
 pub(crate) async fn run(ctx: sdk::Context, matches: &ArgMatches) -> Result<()> {
-    let bundle_path = matches
+    let Some((op, sub_matches)) = matches.subcommand() else {
+        unreachable!("clap is configured to require a subcommand, and they're all handled above");
+    };
+
+    let bundle_path = sub_matches
         .get_one::<PathBuf>("bundle")
         .expect("required parameter not found");
 
-    let file_path = matches
+    let in_file_path = sub_matches
-        .get_one::<PathBuf>("file")
+        .get_one::<PathBuf>("new-file")
         .expect("required parameter not found");
 
-    tracing::trace!(bundle_path = %bundle_path.display(), file_path = %file_path.display());
+    let patch_number = matches
+        .get_one::<u16>("patch")
+        .map(|num| format!("{:03}", num));
 
-    let mut bundle = {
+    let output_path = matches
-        let binary = fs::read(bundle_path).await?;
+        .get_one::<PathBuf>("output")
-        let name = Bundle::get_name_from_path(&ctx, bundle_path);
+        .cloned()
-        Bundle::from_binary(&ctx, name, binary).wrap_err("Failed to open bundle file")?
+        .unwrap_or_else(|| {
+            let mut output_path = bundle_path.clone();
+
+            if let Some(patch_number) = patch_number.as_ref() {
+                output_path.set_extension(format!("patch_{:03}", patch_number));
+            }
+
+            output_path
+        });
+
+    let target_name = if op == "replace" {
+        sub_matches
+            .get_one::<String>("bundle-file")
+            .map(|name| match u64::from_str_radix(name, 16) {
+                Ok(id) => IdString64::from(id),
+                Err(_) => IdString64::String(name.clone()),
+            })
+            .expect("argument is required")
+    } else {
+        let mut path = PathBuf::from(in_file_path);
+        path.set_extension("");
+        IdString64::from(path.to_slash_lossy().to_string())
     };
 
-    if let Some(name) = matches.get_one::<String>("replace") {
+    let file_type = if let Some(forced_type) = matches.get_one::<String>("type") {
-        let mut file = File::open(&file_path)
+        BundleFileType::from_str(forced_type.as_str()).wrap_err("Unknown file type")?
-            .await
+    } else {
-            .wrap_err_with(|| format!("Failed to open '{}'", file_path.display()))?;
+        in_file_path
+            .extension()
+            .and_then(|s| s.to_str())
+            .ok_or_eyre("File extension missing")
+            .and_then(BundleFileType::from_str)
+            .wrap_err("Unknown file type")
+            .with_suggestion(|| "Use '--type TYPE' to specify the file type")?
+    };
 
-        if let Some(variant) = bundle
+    {
-            .files_mut()
+        let span = tracing::Span::current();
-            .filter(|file| file.matches_name(name.clone()))
+        if !span.is_disabled() {
-            // TODO: Handle file variants
+            span.record("bundle_path", bundle_path.display().to_string());
-            .find_map(|file| file.variants_mut().next())
+            span.record("in_file_path", in_file_path.display().to_string());
-        {
+            span.record("output_path", output_path.display().to_string());
-            let mut data = Vec::new();
+            span.record("raw", sub_matches.get_flag("raw"));
-            file.read_to_end(&mut data)
+            span.record("target_name", target_name.display().to_string());
-                .await
+            span.record("file_type", format!("{:?}", file_type));
-                .wrap_err("Failed to read input file")?;
+        }
-            variant.set_data(data);
+    }
-        } else {
+
-            let err = eyre::eyre!("No file '{}' in this bundle.", name)
+    let bundle_name = Bundle::get_name_from_path(&ctx, bundle_path);
-                .with_suggestion(|| {
+    let mut bundle = {
+        let binary = fs::read(bundle_path).await?;
+        Bundle::from_binary(&ctx, bundle_name.clone(), binary)
+            .wrap_err_with(|| format!("Failed to open bundle '{}'", bundle_path.display()))?
+    };
+
+    if op == "copy" {
+        unimplemented!("Implement copying a file from one bundle to the other.");
+    }
+
+    let output_bundle = match op {
+        "replace" => {
+            let Some(file) = bundle
+                .files_mut()
+                .find(|file| *file.base_name() == target_name)
+            else {
+                let err = eyre::eyre!(
+                    "No file with name '{}' in bundle '{}'",
+                    target_name.display(),
+                    bundle_path.display()
+                );
+
+                return Err(err).with_suggestion(|| {
                     format!(
-                        "Run '{} bundle list {}' to list the files in this bundle.",
+                        "Run '{} bundle list \"{}\"' to list the files in this bundle.",
                         clap::crate_name!(),
                         bundle_path.display()
                     )
-                })
-                .with_suggestion(|| {
-                    format!(
-                        "Use '{} bundle inject --add {} {} {}' to add it as a new file",
-                        clap::crate_name!(),
-                        name,
-                        bundle_path.display(),
-                        file_path.display()
-                    )
                 });
+            };
 
-            return Err(err);
+            if sub_matches.get_flag("raw") {
+                let variant_index = sub_matches
+                    .get_one::<u8>("variant")
+                    .expect("argument with default missing");
+
+                let Some(variant) = file.variants_mut().nth(*variant_index as usize) else {
+                    let err = eyre::eyre!(
+                        "Variant index '{}' does not exist in '{}'",
+                        variant_index,
+                        target_name.display()
+                    );
+
+                    return Err(err).with_suggestion(|| {
+                        format!(
+                            "See '{} bundle inject add --help' if you want to add it as a new file",
+                            clap::crate_name!(),
+                        )
+                    });
+                };
+
+                let data = tokio::fs::read(&in_file_path).await.wrap_err_with(|| {
+                    format!("Failed to read file '{}'", in_file_path.display())
+                })?;
+                variant.set_data(data);
+                file.set_modded(true);
+                bundle
+            } else {
+                let mut bundle_file = compile_file(in_file_path, target_name.clone(), file_type)
+                    .await
+                    .wrap_err("Failed to compile")?;
+
+                bundle_file.set_modded(true);
+
+                if patch_number.is_some() {
+                    let mut output_bundle = Bundle::new(bundle_name);
+                    output_bundle.add_file(bundle_file);
+                    output_bundle
+                } else {
+                    *file = bundle_file;
+
+                    dbg!(&file);
+                    bundle
+                }
+            }
         }
+        "add" => {
+            unimplemented!("Implement adding a new file to the bundle.");
+        }
+        _ => unreachable!("no other operations exist"),
+    };
 
-        let out_path = matches.get_one::<PathBuf>("output").unwrap_or(bundle_path);
+    let data = output_bundle
-        let data = bundle
+        .to_binary()
-            .to_binary()
+        .wrap_err("Failed to write changed bundle to output")?;
-            .wrap_err("Failed to write changed bundle to output")?;
 
-        fs::write(out_path, &data)
+    fs::write(&output_path, &data)
-            .await
+        .await
-            .wrap_err("Failed to write data to output file")?;
+        .wrap_err_with(|| format!("Failed to write data to '{}'", output_path.display()))?;
 
-        Ok(())
+    tracing::info!("Modified bundle written to '{}'", output_path.display());
-    } else {
+
-        eyre::bail!("Currently, only the '--replace' operation is supported.");
+    Ok(())
-    }
 }

lib/sdk/src/bundle/file.rs

@@ -20,6 +20,7 @@ struct BundleFileHeader {
     len_data_file_name: usize,
 }
 
+#[derive(Clone, Debug)]
 pub struct BundleFileVariant {
     property: u32,
     data: Vec<u8>,
@@ -109,9 +110,12 @@ bitflags! {
     #[derive(Default, Clone, Copy, Debug)]
     pub struct Properties: u32 {
         const DATA = 0b100;
+        // A custom flag used by DTMT to signify a file altered by mods.
+        const MODDED = 1 << 31;
     }
 }
 
+#[derive(Clone, Debug)]
 pub struct BundleFile {
     file_type: BundleFileType,
     name: IdString64,
@@ -133,6 +137,18 @@ impl BundleFile {
         self.variants.push(variant)
     }
 
+    pub fn set_variants(&mut self, variants: Vec<BundleFileVariant>) {
+        self.variants = variants;
+    }
+
+    pub fn set_props(&mut self, props: Properties) {
+        self.props = props;
+    }
+
+    pub fn set_modded(&mut self, is_modded: bool) {
+        self.props.set(Properties::MODDED, is_modded);
+    }
+
     #[tracing::instrument(name = "File::read", skip(ctx, r))]
     pub fn from_reader<R>(ctx: &crate::Context, r: &mut R, props: Properties) -> Result<Self>
     where
@@ -299,14 +315,13 @@ impl BundleFile {
         s
     }
 
-    pub fn matches_name(&self, name: impl Into<IdString64>) -> bool {
+    pub fn matches_name(&self, name: &IdString64) -> bool {
-        let name = name.into();
+        if self.name == *name {
-        if self.name == name {
             return true;
         }
 
         if let IdString64::String(name) = name {
-            self.name(false, None) == name || self.name(true, None) == name
+            self.name(false, None) == *name || self.name(true, None) == *name
         } else {
             false
         }

lib/sdk/src/bundle/mod.rs

@@ -7,14 +7,13 @@ use color_eyre::{Help, Report, SectionExt};
 use oodle::{OodleLZ_CheckCRC, OodleLZ_FuzzSafe, CHUNK_SIZE};
 
 use crate::binary::sync::*;
-use crate::bundle::file::Properties;
 use crate::murmur::{HashGroup, IdString64, Murmur64};
 
 pub(crate) mod database;
 pub(crate) mod file;
 pub(crate) mod filetype;
 
-pub use file::{BundleFile, BundleFileVariant};
+pub use file::{BundleFile, BundleFileVariant, Properties};
 pub use filetype::BundleFileType;
 
 #[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]

lib/sdk/src/lib.rs

@@ -9,5 +9,5 @@ pub mod murmur;
 pub use binary::{FromBinary, ToBinary};
 pub use bundle::database::BundleDatabase;
 pub use bundle::decompress;
-pub use bundle::{Bundle, BundleFile, BundleFileType, BundleFileVariant};
+pub use bundle::{Bundle, BundleFile, BundleFileType, BundleFileVariant, Properties};
 pub use context::{CmdLine, Context};