bitsquid_dt/dtmt
bitsquid_dt/dtmt
2
Fork 0
Code Issues 34 Pull requests 11 Projects 1 Releases 3 Packages 1 Wiki Activity
Var Dump: 
dumpVar: only available in dev mode
Mailing List

feat: Implement bundle writing and file injecting

Browse source
This commit is contained in:
Lucas Schwiderski 2022-11-16 09:36:46 +01:00
parent 6bb5aef407
commit d500b01709
Signed by: lucas
GPG key ID: AA12679AAA6DF4D8
11 changed files with 224 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.gitignore vendored
View file

@ -3,3 +3,4 @@
.envrc
liboo2corelinux64.so
oo2core_8_win64.dll
dictionary.csv

112
src/bin/cmd/bundle/inject.rs Normal file
View file

@ -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<RwLock<dtmt::Context>>, matches: &ArgMatches) -> Result<()> {
let bundle_path = matches
.get_one::<PathBuf>("bundle")
.expect("required parameter not found");
let file_path = matches
.get_one::<PathBuf>("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::<String>("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::<PathBuf>("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.");
}
}

3
src/bin/cmd/bundle/mod.rs
View file

@ -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<RwLock<dtmt::Context>>, 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"

2
src/bin/cmd/util.rs
View file

@ -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);

2
src/bin/dtmt.rs
View file

@ -88,7 +88,7 @@ async fn main() -> Result<()> {
if is_default {
return;
}
tracing::error!("{}", err);
tracing::error!("{:#}", err);
return;
}

16
src/binary.rs
View file

@ -35,11 +35,11 @@ macro_rules! make_read {
macro_rules! make_write {
($func:ident, $op:ident, $type:ty) => {
pub(crate) async fn $func<W>(r: &mut W, val: $type) -> Result<()>
pub(crate) async fn $func<W>(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>(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<S>(stream: &mut S) -> Result<()>
where
@ -112,7 +112,7 @@ where
Ok(())
}
pub(crate) async fn read_up_to<R>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize>
pub(crate) async fn _read_up_to<R>(r: &mut R, buf: &mut Vec<u8>) -> Result<usize>
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?;

26
src/bundle/file.rs
View file

@ -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<u64> for BundleFileType {
impl From<Murmur64> 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<u8> {
&self.data
}
pub fn set_data(&mut self, data: Vec<u8>) {
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<S>(&self, name: S) -> bool
where
S: AsRef<str>,
{
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<Item = &mut BundleFileVariant> {
self.variants.iter_mut()
}
pub fn raw(&self) -> Result<Vec<UserFile>> {
let files = self
.variants

65
src/bundle/mod.rs
View file

@ -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<BundleFile> {
&self.files
}
pub fn files_mut(&mut self) -> impl Iterator<Item = &mut BundleFile> {
self.files.iter_mut()
}
}
/// Returns a decompressed version of the bundle data.

6
src/murmur/dictionary.rs
View file

@ -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 {

10
src/murmur/mod.rs
View file

@ -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<E>(self, value: u64) -> Result<Self::Value, E>

8
src/oodle/mod.rs
View file

@ -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)
}
}