From 7b95918000430088139af81b9aea2d540e25f98b Mon Sep 17 00:00:00 2001 From: Lucas Schwiderski Date: Fri, 22 Sep 2023 15:42:16 +0200 Subject: [PATCH] Refactor code for file injection I ended up wrapping the raw data in a `BundleFile` twice. I also made '--compile' the default, as it should be much less often that raw data needs to be inserted. Even files that are essentially raw binary blobs, like `.wwise_event`, still have some custom fields that need to be accounted for. --- crates/dtmt/src/cmd/bundle/inject.rs | 327 +++++++++++++++++++++------ lib/sdk/src/bundle/file.rs | 23 +- lib/sdk/src/bundle/mod.rs | 3 +- lib/sdk/src/lib.rs | 2 +- 4 files changed, 277 insertions(+), 78 deletions(-) diff --git a/crates/dtmt/src/cmd/bundle/inject.rs b/crates/dtmt/src/cmd/bundle/inject.rs index 2b86691..21f4a91 100644 --- a/crates/dtmt/src/cmd/bundle/inject.rs +++ b/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 color_eyre::eyre::{self, Context, Result}; +use clap::{value_parser, Arg, ArgAction, ArgMatches, Command}; +use color_eyre::eyre::{self, Context, OptionExt, Result}; use color_eyre::Help; -use sdk::Bundle; -use tokio::fs::{self, File}; -use tokio::io::AsyncReadExt; +use path_slash::PathBufExt as _; +use sdk::murmur::IdString64; +use sdk::{Bundle, BundleFile, BundleFileType}; +use tokio::fs; pub(crate) fn command_definition() -> Command { Command::new("inject") - .about("Inject a file into a bundle.") - .arg( - Arg::new("replace") - .help("The name of a file in the bundle whos content should be replaced.") - .short('r') - .long("replace"), - ) + .subcommand_required(true) + .about("Inject a file into a bundle.\n\ + Raw binary data can be used to directly replace the file's variant data blob without affecting the metadata.\n\ + Alternatively, a compiler format may be specified, and a complete bundle file is created.") .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_` suffix if you also use '--patch'.", ) .short('o') .long("output") .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)), + Arg::new("patch") + .help("Create a patch bundle. Optionally, a patch NUMBER may be specified as \ + '--patch=123'.\nThe maximum number is 999, the default is 1.\n\ + If `--output` is not specified, the `.patch_` 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") - .help("Path to the file to inject.") - .required(true) - .value_parser(value_parser!(PathBuf)), + Arg::new("type") + .help("Compile the new file as the given TYPE. If omitted, the file type is \ + is guessed from the file extension.") + .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 + std::fmt::Debug, + name: impl Into + std::fmt::Debug, + file_type: BundleFileType, +) -> Result { + 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::("bundle") .expect("required parameter not found"); - let file_path = matches - .get_one::("file") + let in_file_path = sub_matches + .get_one::("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::("patch") + .map(|num| format!("{:03}", num)); - let mut bundle = { - let binary = fs::read(bundle_path).await?; - let name = Bundle::get_name_from_path(&ctx, bundle_path); - Bundle::from_binary(&ctx, name, binary).wrap_err("Failed to open bundle file")? + let output_path = matches + .get_one::("output") + .cloned() + .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::("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::("replace") { - let mut file = File::open(&file_path) - .await - .wrap_err_with(|| format!("Failed to open '{}'", file_path.display()))?; + let file_type = if let Some(forced_type) = matches.get_one::("type") { + BundleFileType::from_str(forced_type.as_str()).wrap_err("Unknown file type")? + } else { + 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() - .filter(|file| file.matches_name(name.clone())) - // TODO: Handle file variants - .find_map(|file| file.variants_mut().next()) - { - let mut data = Vec::new(); - file.read_to_end(&mut data) - .await - .wrap_err("Failed to read input file")?; - variant.set_data(data); - } else { - let err = eyre::eyre!("No file '{}' in this bundle.", name) - .with_suggestion(|| { + { + let span = tracing::Span::current(); + if !span.is_disabled() { + span.record("bundle_path", bundle_path.display().to_string()); + span.record("in_file_path", in_file_path.display().to_string()); + span.record("output_path", output_path.display().to_string()); + span.record("raw", sub_matches.get_flag("raw")); + span.record("target_name", target_name.display().to_string()); + span.record("file_type", format!("{:?}", file_type)); + } + } + + let bundle_name = Bundle::get_name_from_path(&ctx, bundle_path); + 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::("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::("output").unwrap_or(bundle_path); - let data = bundle - .to_binary() - .wrap_err("Failed to write changed bundle to output")?; + let data = output_bundle + .to_binary() + .wrap_err("Failed to write changed bundle to output")?; - fs::write(out_path, &data) - .await - .wrap_err("Failed to write data to output file")?; + fs::write(&output_path, &data) + .await + .wrap_err_with(|| format!("Failed to write data to '{}'", output_path.display()))?; - Ok(()) - } else { - eyre::bail!("Currently, only the '--replace' operation is supported."); - } + tracing::info!("Modified bundle written to '{}'", output_path.display()); + + Ok(()) } diff --git a/lib/sdk/src/bundle/file.rs b/lib/sdk/src/bundle/file.rs index f387409..6d49821 100644 --- a/lib/sdk/src/bundle/file.rs +++ b/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, @@ -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) { + 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(ctx: &crate::Context, r: &mut R, props: Properties) -> Result where @@ -299,14 +315,13 @@ impl BundleFile { s } - pub fn matches_name(&self, name: impl Into) -> bool { - let name = name.into(); - if self.name == name { + pub fn matches_name(&self, name: &IdString64) -> bool { + 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 } diff --git a/lib/sdk/src/bundle/mod.rs b/lib/sdk/src/bundle/mod.rs index 075f4d2..edb71bb 100644 --- a/lib/sdk/src/bundle/mod.rs +++ b/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)] diff --git a/lib/sdk/src/lib.rs b/lib/sdk/src/lib.rs index a24b3bd..9b1806b 100644 --- a/lib/sdk/src/lib.rs +++ b/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};