Refactor code for file injection

This commit is contained in:
Lucas Schwiderski 2024-07-26 15:00:58 +02:00
parent bcbc005df7
commit 6f6df14bfc
Signed by: lucas
GPG key ID: AA12679AAA6DF4D8
2 changed files with 221 additions and 81 deletions

View file

@ -1,40 +1,68 @@
use std::path::PathBuf;
use clap::{value_parser, Arg, ArgMatches, Command};
use color_eyre::eyre::{self, Context, Result};
use color_eyre::eyre::{self, Context, OptionExt as _, Result};
use color_eyre::Help;
use path_slash::PathBufExt as _;
use sdk::filetype::texture;
use sdk::murmur::IdString64;
use sdk::Bundle;
use tokio::fs::{self, File};
use tokio::io::AsyncReadExt;
use tokio::fs;
pub(crate) fn command_definition() -> Command {
Command::new("inject")
.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("replace")
.help("The name of a file in the bundle whose content should be replaced.")
.short('r')
.long("replace"),
)
.arg(
Arg::new("compile")
.help("Compile the file with the given data format before injecting.")
.long("compile")
.short('c')
)
.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("patch")
.help("Create a patch bundle. Optionally, a patch NUMBER may be specified as \
'--patch=123'.\nThe maximum number is 1.")
.short('p')
.long("patch")
.num_args(0..=1)
.require_equals(true)
.default_missing_value("1")
.value_name("NUMBER")
.value_parser(value_parser!(u16))
)
.subcommand(
Command::new("replace")
.about("Replace an existing file in the bundle")
.arg(
Arg::new("compile")
.help("Compile the file as TYPE before injecting. \
Metadata will also be overwritten")
.long("compile")
.short('c')
.value_name("TYPE")
.value_parser(["texture"])
)
.arg(
Arg::new("variant")
.help("Sepcify the variant index to replace.")
.long("variant")
.default_value("0")
.value_parser(value_parser!(u8))
)
.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.")
@ -42,85 +70,198 @@ pub(crate) fn command_definition() -> Command {
.value_parser(value_parser!(PathBuf)),
)
.arg(
Arg::new("file")
.help("Path to the file to inject.")
.required(true)
.value_parser(value_parser!(PathBuf)),
Arg::new("bundle-file")
.help("The name of a file in the bundle whose content should be replaced.")
.required(true),
),
)
// .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(
skip_all,
fields(
bundle_path = tracing::field::Empty,
in_file_path = tracing::field::Empty,
compile = tracing::field::Empty,
output_path = tracing::field::Empty,
file_name = 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 compile = matches.get_one::<String>("compile");
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")
.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 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::<PathBuf>("output")
.cloned()
.unwrap_or_else(|| {
let mut output_path = bundle_path.clone();
let name = match matches.get_one::<String>("replace") {
Some(name) => match u64::from_str_radix(name, 16) {
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()),
},
None => eyre::bail!("Currently, only the '--replace' operation is supported."),
})
.expect("argument is required")
} else {
let mut path = PathBuf::from(in_file_path);
path.set_extension("");
IdString64::from(path.to_slash_lossy().to_string())
};
let mut file = File::open(&file_path)
.await
.wrap_err_with(|| format!("Failed to open '{}'", file_path.display()))?;
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);
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("compile", compile);
span.record("file_name", target_name.display().to_string());
}
}
let mut bundle = {
let name = Bundle::get_name_from_path(&ctx, bundle_path);
if patch_number.is_some() {
unimplemented!("Create patch bundle");
} else {
let err =
eyre::eyre!("No file '{}' in this bundle.", name.display()).with_suggestion(|| {
let binary = fs::read(bundle_path).await?;
Bundle::from_binary(&ctx, name, 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 file_data = tokio::fs::read(&in_file_path)
.await
.wrap_err_with(|| format!("Failed to read file '{}'", in_file_path.display()))?;
let data = if let Some(compile_type) = compile {
match compile_type.as_str() {
"texture" => {
let sjson = String::from_utf8(file_data).wrap_err_with(|| {
format!("Invalid UTF8 data in '{}'", in_file_path.display())
})?;
let root = in_file_path
.parent()
.ok_or_eyre("File path has no parent")?;
let texture = texture::compile(target_name.clone(), sjson, root)
.await
.wrap_err_with(|| {
format!(
"Run '{} bundle list {}' to list the files in this bundle.",
"Failed to compile file as texture: {}",
in_file_path.display()
)
})?;
texture.to_binary()?
}
_ => unreachable!("clap only allows file types that we implement"),
}
} else {
file_data
};
match op {
"replace" => {
let Some(file) = bundle
.files_mut()
.find(|file| file.matches_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.",
clap::crate_name!(),
bundle_path.display()
)
// Not yet supported.
// })
// .with_suggestion(|| {
// format!(
// "Use '{} bundle inject --add {} {} {}' to add it as a new file",
// clap::crate_name!(),
// name.display(),
// bundle_path.display(),
// file_path.display()
// )
});
};
return Err(err);
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!(),
)
});
};
variant.set_data(data);
}
"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 = bundle
.to_binary()
.wrap_err("Failed to write changed bundle to output")?;
fs::write(out_path, &data)
fs::write(&output_path, &data)
.await
.wrap_err("Failed to write data to output file")?;
.wrap_err_with(|| format!("Failed to write data to '{}'", output_path.display()))?;
Ok(())
}

View file

@ -335,14 +335,13 @@ impl BundleFile {
s
}
pub fn matches_name(&self, name: impl Into<IdString64>) -> 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
}