From d500b01709451dfca11ae1ef1cdaed6bd83c836e Mon Sep 17 00:00:00 2001 From: Lucas Schwiderski Date: Wed, 16 Nov 2022 09:36:46 +0100 Subject: [PATCH] feat: Implement bundle writing and file injecting --- .gitignore | 1 + src/bin/cmd/bundle/inject.rs | 112 +++++++++++++++++++++++++++++++++++ src/bin/cmd/bundle/mod.rs | 3 + src/bin/cmd/util.rs | 2 +- src/bin/dtmt.rs | 2 +- src/binary.rs | 16 ++--- src/bundle/file.rs | 26 ++++++-- src/bundle/mod.rs | 65 +++++++++++++++++--- src/murmur/dictionary.rs | 6 ++ src/murmur/mod.rs | 10 +++- src/oodle/mod.rs | 8 ++- 11 files changed, 224 insertions(+), 27 deletions(-) create mode 100644 src/bin/cmd/bundle/inject.rs diff --git a/.gitignore b/.gitignore index b1addd3..8f1bf6a 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ .envrc liboo2corelinux64.so oo2core_8_win64.dll +dictionary.csv diff --git a/src/bin/cmd/bundle/inject.rs b/src/bin/cmd/bundle/inject.rs new file mode 100644 index 0000000..a96250c --- /dev/null +++ b/src/bin/cmd/bundle/inject.rs @@ -0,0 +1,112 @@ +use std::{path::PathBuf, sync::Arc}; + +use clap::{value_parser, Arg, ArgMatches, Command}; +use color_eyre::{ + eyre::{self, Context, Result}, + Help, +}; +use dtmt::Bundle; +use tokio::{fs::File, io::AsyncReadExt, sync::RwLock}; + +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"), + ) + .arg( + Arg::new("output") + .help( + "The path to write the changed bundle to. \ + If omitted, the input bundle will be overwritten.", + ) + .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( + Arg::new("file") + .help("Path to the file to inject.") + .required(true) + .value_parser(value_parser!(PathBuf)), + ) +} + +#[tracing::instrument(skip_all)] +pub(crate) async fn run(ctx: Arc>, matches: &ArgMatches) -> Result<()> { + let bundle_path = matches + .get_one::("bundle") + .expect("required parameter not found"); + + let file_path = matches + .get_one::("file") + .expect("required parameter not found"); + + tracing::trace!(bundle_path = %bundle_path.display(), file_path = %file_path.display()); + + let mut bundle = Bundle::open(ctx.clone(), bundle_path) + .await + .wrap_err("Failed to open bundle file")?; + + 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()))?; + + if let Some(variant) = bundle + .files_mut() + .filter(|file| file.matches_name(_name)) + // TODO: Handle file variants + .filter_map(|file| file.variants_mut().next()) + .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(|| { + format!( + "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); + } + + let out_path = matches.get_one::("output").unwrap_or(bundle_path); + let mut out_file = File::create(out_path) + .await + .wrap_err_with(|| format!("failed to open output file {}", out_path.display()))?; + bundle + .write(ctx.clone(), &mut out_file) + .await + .wrap_err("failed to write changed bundle to output")?; + + Ok(()) + } else { + eyre::bail!("Currently, only the '--replace' operation is supported."); + } +} diff --git a/src/bin/cmd/bundle/mod.rs b/src/bin/cmd/bundle/mod.rs index a7ce95b..21ecdca 100644 --- a/src/bin/cmd/bundle/mod.rs +++ b/src/bin/cmd/bundle/mod.rs @@ -8,6 +8,7 @@ use dtmt::Oodle; mod decompress; mod extract; +mod inject; mod list; #[cfg(target_os = "windows")] @@ -33,6 +34,7 @@ pub(crate) fn command_definition() -> Command { ) .subcommand(decompress::command_definition()) .subcommand(extract::command_definition()) + .subcommand(inject::command_definition()) .subcommand(list::command_definition()) } @@ -47,6 +49,7 @@ pub(crate) async fn run(ctx: Arc>, matches: &ArgMatches) - match matches.subcommand() { Some(("decompress", sub_matches)) => decompress::run(ctx, sub_matches).await, Some(("extract", sub_matches)) => extract::run(ctx, sub_matches).await, + Some(("inject", sub_matches)) => inject::run(ctx, sub_matches).await, Some(("list", sub_matches)) => list::run(ctx, sub_matches).await, _ => unreachable!( "clap is configured to require a subcommand, and they're all handled above" diff --git a/src/bin/cmd/util.rs b/src/bin/cmd/util.rs index f5846d9..99c2cbd 100644 --- a/src/bin/cmd/util.rs +++ b/src/bin/cmd/util.rs @@ -18,7 +18,7 @@ where } Err(err) => { if err.kind() != io::ErrorKind::NotADirectory { - tracing::error!(%err, "Failed to read path"); + tracing::error!("Failed to read path: {:?}", err); } let paths = vec![PathBuf::from(path.as_ref())]; tracing::debug!(is_dir = false, resolved_paths = ?paths); diff --git a/src/bin/dtmt.rs b/src/bin/dtmt.rs index ed54934..98ecd5c 100644 --- a/src/bin/dtmt.rs +++ b/src/bin/dtmt.rs @@ -88,7 +88,7 @@ async fn main() -> Result<()> { if is_default { return; } - tracing::error!("{}", err); + tracing::error!("{:#}", err); return; } diff --git a/src/binary.rs b/src/binary.rs index 12854b7..2a50211 100644 --- a/src/binary.rs +++ b/src/binary.rs @@ -35,11 +35,11 @@ macro_rules! make_read { macro_rules! make_write { ($func:ident, $op:ident, $type:ty) => { - pub(crate) async fn $func(r: &mut W, val: $type) -> Result<()> + pub(crate) async fn $func(w: &mut W, val: $type) -> Result<()> where W: AsyncWrite + AsyncSeek + std::marker::Unpin, { - let res = r + let res = w .$op(val) .await .wrap_err(concat!("failed to write ", stringify!($type))); @@ -48,7 +48,7 @@ macro_rules! make_write { return res; } - let pos = r.stream_position().await; + let pos = w.stream_position().await; if pos.is_ok() { res.with_section(|| { format!("{pos:#X} ({pos})", pos = pos.unwrap()).header("Position: ") @@ -61,7 +61,7 @@ macro_rules! make_write { } macro_rules! make_skip { - ($func:ident, $read:ident, $op:ident, $type:ty) => { + ($func:ident, $read:ident, $type:ty) => { pub(crate) async fn $func(r: &mut R, cmp: $type) -> Result<()> where R: AsyncRead + AsyncSeek + std::marker::Unpin, @@ -92,8 +92,8 @@ make_write!(write_u8, write_u8, u8); make_write!(write_u32, write_u32_le, u32); make_write!(write_u64, write_u64_le, u64); -make_skip!(skip_u8, read_u8, read_u8, u8); -make_skip!(skip_u32, read_u32, read_u32_le, u32); +make_skip!(skip_u8, read_u8, u8); +make_skip!(skip_u32, read_u32, u32); pub(crate) async fn skip_padding(stream: &mut S) -> Result<()> where @@ -112,7 +112,7 @@ where Ok(()) } -pub(crate) async fn read_up_to(r: &mut R, buf: &mut Vec) -> Result +pub(crate) async fn _read_up_to(r: &mut R, buf: &mut Vec) -> Result where R: AsyncRead + AsyncSeek + std::marker::Unpin, { @@ -142,6 +142,8 @@ where let pos = w.stream_position().await?; let size = 16 - (pos % 16) as usize; + tracing::trace!(padding_size = size, "Writing padding"); + if size > 0 && size < 16 { let buf = vec![0; size]; w.write_all(&buf).await?; diff --git a/src/bundle/file.rs b/src/bundle/file.rs index 74a264e..e88069c 100644 --- a/src/bundle/file.rs +++ b/src/bundle/file.rs @@ -1,4 +1,3 @@ -use std::ops::Deref; use std::sync::Arc; use color_eyre::{Help, Result, SectionExt}; @@ -158,8 +157,8 @@ impl BundleFileType { } } - pub fn hash(&self) -> u64 { - *Murmur64::from(*self).deref() + pub fn hash(&self) -> Murmur64 { + Murmur64::from(*self) } } @@ -171,7 +170,7 @@ impl From for BundleFileType { impl From for BundleFileType { fn from(hash: Murmur64) -> BundleFileType { - match hash.deref() { + match *hash { 0x931e336d7646cc26 => BundleFileType::Animation, 0xdcfb9e18fff13984 => BundleFileType::AnimationCurves, 0x3eed05ba83af5090 => BundleFileType::Apb, @@ -361,6 +360,11 @@ impl BundleFileVariant { pub fn data(&self) -> &Vec { &self.data } + + pub fn set_data(&mut self, data: Vec) { + self.header.size = data.len(); + self.data = data; + } } pub struct BundleFile { @@ -418,7 +422,7 @@ impl BundleFile { where W: AsyncWrite + AsyncSeek + std::marker::Unpin, { - write_u64(w, self.file_type.hash()).await?; + write_u64(w, *self.file_type.hash()).await?; write_u64(w, *self.hash).await?; let header_count = self.variants.len(); @@ -459,6 +463,14 @@ impl BundleFile { } } + pub fn matches_name(&self, name: S) -> bool + where + S: AsRef, + { + let name = name.as_ref(); + self.name == name || self.name(false) == name || self.name(true) == name + } + pub fn hash(&self) -> Murmur64 { self.hash } @@ -471,6 +483,10 @@ impl BundleFile { &self.variants } + pub fn variants_mut(&mut self) -> impl Iterator { + self.variants.iter_mut() + } + pub fn raw(&self) -> Result> { let files = self .variants diff --git a/src/bundle/mod.rs b/src/bundle/mod.rs index 0d429f8..0a78e77 100644 --- a/src/bundle/mod.rs +++ b/src/bundle/mod.rs @@ -56,8 +56,8 @@ impl EntryHeader { where R: AsyncRead + AsyncSeek + std::marker::Unpin, { - let extension_hash = r.read_u64().await?; - let name_hash = r.read_u64().await?; + let extension_hash = read_u64(r).await?; + let name_hash = read_u64(r).await?; let flags = read_u32(r).await?; // NOTE: Known values so far: @@ -67,7 +67,7 @@ impl EntryHeader { if flags != 0x0 { tracing::debug!( flags, - "Unexpected meta flags for file {:08X}.{:08X}", + "Unexpected meta flags for file {:016X}.{:016X}", name_hash, extension_hash ); @@ -113,17 +113,20 @@ impl Bundle { // `AsyncRead` and the bundle name separately. let path = path.as_ref(); let bundle_name = if let Some(name) = path.file_name() { - let hash = Murmur64::try_from(name.to_string_lossy().as_ref()) - .wrap_err_with(|| format!("failed to turn string into hash: {:?}", name))?; - ctx.read().await.lookup_hash(hash, HashGroup::Filename) + match Murmur64::try_from(name.to_string_lossy().as_ref()) { + Ok(hash) => ctx.read().await.lookup_hash(hash, HashGroup::Filename), + Err(err) => { + tracing::debug!("failed to turn bundle name into hash: {}", err); + name.to_string_lossy().to_string() + } + } } else { eyre::bail!("Invalid path to bundle file: {}", path.display()); }; let f = fs::File::open(path) .await - .wrap_err("Failed to open bundle file") - .with_section(|| path.display().to_string().header("Path"))?; + .wrap_err_with(|| format!("failed to open bundle file {}", path.display()))?; let mut r = BufReader::new(f); @@ -171,6 +174,7 @@ impl Bundle { r.seek(SeekFrom::Current(4)).await?; let mut decompressed = Vec::with_capacity(unpacked_size); + let mut unpacked_size_tracked = unpacked_size; for (chunk_index, chunk_size) in chunk_sizes.into_iter().enumerate() { let span = tracing::debug_span!("Decompressing chunk", chunk_index, chunk_size); @@ -200,6 +204,14 @@ impl Bundle { OodleLZ_CheckCRC::No, )?; + if unpacked_size_tracked < CHUNK_SIZE { + raw_buffer.resize(unpacked_size_tracked, 0); + } else { + unpacked_size_tracked -= CHUNK_SIZE; + } + + tracing::trace!(raw_size = raw_buffer.len()); + decompressed.append(&mut raw_buffer); Ok(()) } @@ -254,6 +266,8 @@ impl Bundle { let buf = Vec::new(); let mut c = Cursor::new(buf); + tracing::trace!(num_files = self.files.len()); + async { for file in self.files.iter() { file.write(ctx.clone(), &mut c).await?; @@ -267,19 +281,48 @@ impl Bundle { c.into_inner() }; + // Ceiling division (or division toward infinity) to calculate + // the number of chunks required to fit the unpacked data. + let num_chunks = (unpacked_data.len() + CHUNK_SIZE - 1) / CHUNK_SIZE; + tracing::trace!(num_chunks); + write_u32(w, num_chunks as u32).await?; + + let chunk_sizes_start = w.stream_position().await?; + tracing::trace!(chunk_sizes_start); + w.seek(SeekFrom::Current(num_chunks as i64 * 4)).await?; + + write_padding(w).await?; + + tracing::trace!(unpacked_size = unpacked_data.len()); + write_u32(w, unpacked_data.len() as u32).await?; + // NOTE: Unknown u32 that's always been 0 so far + write_u32(w, 0).await?; + let chunks = unpacked_data.chunks(CHUNK_SIZE); let ctx = ctx.read().await; let oodle_lib = ctx.oodle.as_ref().unwrap(); + let mut chunk_sizes = Vec::with_capacity(num_chunks); for chunk in chunks { let compressed = oodle_lib.compress(chunk)?; + tracing::trace!( + raw_chunk_size = chunk.len(), + compressed_chunk_size = compressed.len() + ); + chunk_sizes.push(compressed.len()); write_u32(w, compressed.len() as u32).await?; write_padding(w).await?; w.write_all(&compressed).await?; } - todo!("compress data and count chunks"); + w.seek(SeekFrom::Start(chunk_sizes_start)).await?; + + for size in chunk_sizes { + write_u32(w, size as u32).await?; + } + + Ok(()) } pub fn name(&self) -> &String { @@ -289,6 +332,10 @@ impl Bundle { pub fn files(&self) -> &Vec { &self.files } + + pub fn files_mut(&mut self) -> impl Iterator { + self.files.iter_mut() + } } /// Returns a decompressed version of the bundle data. diff --git a/src/murmur/dictionary.rs b/src/murmur/dictionary.rs index 2580f32..3f4a280 100644 --- a/src/murmur/dictionary.rs +++ b/src/murmur/dictionary.rs @@ -15,6 +15,12 @@ pub enum HashGroup { Other, } +impl HashGroup { + pub fn all() -> [Self; 3] { + [Self::Filename, Self::Filetype, Self::Other] + } +} + impl std::fmt::Display for HashGroup { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { diff --git a/src/murmur/mod.rs b/src/murmur/mod.rs index 45c9dfa..bc9c54b 100644 --- a/src/murmur/mod.rs +++ b/src/murmur/mod.rs @@ -19,6 +19,14 @@ pub use murmurhash64::hash; pub use murmurhash64::hash32; pub use murmurhash64::hash_inverse as inverse; +fn _swap_bytes_u32(value: u32) -> u32 { + u32::from_le_bytes(value.to_be_bytes()) +} + +fn _swap_bytes_u64(value: u64) -> u64 { + u64::from_le_bytes(value.to_be_bytes()) +} + #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct Murmur64(u64); @@ -74,7 +82,7 @@ impl<'de> Visitor<'de> for Murmur64 { E: serde::de::Error, { let bytes = value.to_le_bytes(); - self.visit_u64(u64::from_le_bytes(bytes)) + Ok(Self::from(u64::from_le_bytes(bytes))) } fn visit_u64(self, value: u64) -> Result diff --git a/src/oodle/mod.rs b/src/oodle/mod.rs index 63343b1..3f7303c 100644 --- a/src/oodle/mod.rs +++ b/src/oodle/mod.rs @@ -79,8 +79,6 @@ impl Oodle { ) }; - tracing::debug!(uncompressed_size = ret, "Decompressed chunk"); - if ret == 0 { eyre::bail!("Failed to decompress chunk."); } @@ -93,7 +91,9 @@ impl Oodle { where I: AsRef<[u8]>, { - let raw = data.as_ref(); + let mut raw = Vec::from(data.as_ref()); + raw.resize(CHUNK_SIZE, 0); + // TODO: Query oodle for buffer size let mut out = vec![0u8; CHUNK_SIZE]; @@ -123,6 +123,8 @@ impl Oodle { eyre::bail!("Failed to compress chunk."); } + out.resize(ret as usize, 0); + Ok(out) } }