diff --git a/src/args.rs b/src/args.rs index 840ac3a..a71dd47 100644 --- a/src/args.rs +++ b/src/args.rs @@ -545,6 +545,10 @@ pub struct Arg { #[arg(long, global = true)] /// Musica xor key for paz archive. pub musica_xor_key: Option, + #[cfg(feature = "musica-arc")] + #[arg(long, global = true)] + /// Compress files in Musica paz archive when packing paz archive. + pub musica_compress: bool, #[command(subcommand)] /// Command pub command: Command, diff --git a/src/ext/io.rs b/src/ext/io.rs index 96affa3..cf7a949 100644 --- a/src/ext/io.rs +++ b/src/ext/io.rs @@ -2301,3 +2301,34 @@ impl Seek for MultipleReadStream { Ok(()) } } + +pub struct TrackStream<'a, T> { + inner: T, + total: &'a mut u64, +} + +impl<'a, T> TrackStream<'a, T> { + pub fn new(inner: T, total: &'a mut u64) -> Self { + TrackStream { inner, total } + } +} + +impl<'a, T: Read> Read for TrackStream<'a, T> { + fn read(&mut self, buf: &mut [u8]) -> Result { + let readed = self.inner.read(buf)?; + *self.total += readed as u64; + Ok(readed) + } +} + +impl<'a, T: Write> Write for TrackStream<'a, T> { + fn write(&mut self, buf: &[u8]) -> Result { + let written = self.inner.write(buf)?; + *self.total += written as u64; + Ok(written) + } + + fn flush(&mut self) -> Result<()> { + self.inner.flush() + } +} diff --git a/src/main.rs b/src/main.rs index 16f7093..5726a48 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2930,6 +2930,8 @@ fn main() { musica_game_title: arg.musica_game_title.clone(), #[cfg(feature = "musica-arc")] musica_xor_key: arg.musica_xor_key, + #[cfg(feature = "musica-arc")] + musica_compress: arg.musica_compress, }); match &arg.command { args::Command::Export { input, output } => { diff --git a/src/scripts/musica/archive/paz.rs b/src/scripts/musica/archive/paz.rs index a64dd7e..348a403 100644 --- a/src/scripts/musica/archive/paz.rs +++ b/src/scripts/musica/archive/paz.rs @@ -9,6 +9,7 @@ use crate::utils::struct_pack::*; use crate::utils::xored_stream::*; use anyhow::Result; use flate2::read::ZlibDecoder; +use flate2::write::ZlibEncoder; use msg_tool_macro::{StructPack, StructUnpack}; use serde::Deserialize; use std::collections::{BTreeMap, HashMap}; @@ -197,6 +198,14 @@ impl PazEntry { pub fn is_compressed(&self) -> bool { (self.flags & 0x1) != 0 } + + pub fn set_is_compressed(&mut self, compressed: bool) { + if compressed { + self.flags |= 0x1; + } else { + self.flags &= !0x1; + } + } } #[derive(Debug)] @@ -377,7 +386,7 @@ impl Script for PazArc { 0, entry.size as u64, )?; - if self.schema.version > 0 { + if self.schema.version > 0 && !entry.is_compressed() { if let Some(type_key) = self.schema.get_type_key(&entry, self.is_audio) { let key = format!( "{} {:08X} {}", @@ -393,10 +402,6 @@ impl Script for PazArc { rc4.skip_bytes(skip as usize); } let stream = Rc4Stream::new(stream, rc4); - if entry.is_compressed() { - let stream = ZlibDecoder::new(stream); - return Ok(Box::new(PazFileEntry::new(entry, stream))); - } return Ok(Box::new(PazFileEntry::new(entry, stream))); } } @@ -559,6 +564,8 @@ pub struct PazArcWriter { schema: Schema, arc_key: ArcKey, xor_key: u8, + compress: bool, + compress_level: u32, } impl PazArcWriter { @@ -653,6 +660,8 @@ impl PazArcWriter { schema: schema.clone(), arc_key: arc_key.clone(), xor_key, + compress: config.musica_compress, + compress_level: config.zlib_compression_level, }) } } @@ -679,7 +688,8 @@ impl Archive for PazArcWriter { let stream = XoredStream::new(&mut self.writer, self.xor_key); let stream = BlowfishEncryptor::new(blowfish, stream); let mut type_key = None; - if self.schema.version > 0 { + entry.set_is_compressed(self.compress); + if self.schema.version > 0 && !self.compress { if let Some(tkey) = self.schema.get_type_key(&entry, self.is_audio) { type_key = Some(tkey.to_string()); } @@ -691,6 +701,9 @@ impl Archive for PazArcWriter { entry, encoding: self.encoding, version: self.schema.version, + compress: self.compress, + compress_level: self.compress_level, + compressed_size: 0, }; return Ok(Box::new(writer)); } else if let Some(mov_key) = &self.mov_key { @@ -748,7 +761,7 @@ impl Archive for PazArcWriter { entry.offset = self.writer.stream_position()?; let stream = XoredStream::new(&mut self.writer, self.xor_key); let stream = BlowfishEncryptor::new(blowfish, stream); - if self.schema.version > 0 { + if self.schema.version > 0 && !self.compress { if let Some(tkey) = self.schema.get_type_key(&entry, self.is_audio) { let key = format!("{} {:08X} {}", entry.name.to_ascii_lowercase(), size, tkey); let key = encode_string(self.encoding, &key, false)?; @@ -766,6 +779,11 @@ impl Archive for PazArcWriter { return Ok(Box::new(writer)); } } + if self.compress { + entry.set_is_compressed(true); + let writer = DataKeyComWriter::new(Box::new(stream), entry, self.compress_level); + return Ok(Box::new(writer)); + } let writer = DateKeyWriter { inner: Box::new(stream), entry, @@ -871,6 +889,46 @@ impl<'a> Drop for DateKeyWriter<'a> { } } +struct DataKeyComWriter<'a> { + inner: ZlibEncoder>, + entry: &'a mut PazEntry, +} + +impl<'a> DataKeyComWriter<'a> { + pub fn new(inner: Box, entry: &'a mut PazEntry, level: u32) -> Self { + DataKeyComWriter { + inner: ZlibEncoder::new(inner, flate2::Compression::new(level)), + entry, + } + } +} + +impl<'a> Write for DataKeyComWriter<'a> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.inner.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +impl<'a> Drop for DataKeyComWriter<'a> { + fn drop(&mut self) { + if let Err(e) = self.inner.try_finish() { + eprintln!( + "Error finishing compression for PAZ file entry '{}': {}", + self.entry.name, e + ); + crate::COUNTER.inc_error(); + return; + } + self.entry.size = self.inner.total_out() as u32; + self.entry.unpacked_size = self.inner.total_in() as u32; + self.entry.aligned_size = (self.entry.size + 7) & !7; + } +} + trait MyWriteSeek: Write + Seek {} impl MyWriteSeek for T {} @@ -926,6 +984,9 @@ struct MemDataKeyWriter<'a> { entry: &'a mut PazEntry, encoding: Encoding, version: u32, + compress: bool, + compress_level: u32, + compressed_size: u64, } impl<'a> Write for MemDataKeyWriter<'a> { @@ -958,37 +1019,49 @@ impl<'a> Drop for MemDataKeyWriter<'a> { self.entry.unpacked_size = data.len() as u32; self.entry.size = self.entry.unpacked_size; self.entry.aligned_size = (self.entry.size + 7) & !7; - let mut stream = if let Some(tkey) = &self.type_key { - let key = format!( - "{} {:08X} {}", - self.entry.name.to_ascii_lowercase(), - self.entry.unpacked_size, - tkey - ); - let key = match encode_string(self.encoding, &key, false) { - Ok(key) => key, - Err(e) => { - eprintln!( - "Error encoding key for PAZ file entry '{}': {}", - self.entry.name, e - ); - crate::COUNTER.inc_error(); - return; + { + let mut stream = if let Some(tkey) = &self.type_key { + let key = format!( + "{} {:08X} {}", + self.entry.name.to_ascii_lowercase(), + self.entry.unpacked_size, + tkey + ); + let key = match encode_string(self.encoding, &key, false) { + Ok(key) => key, + Err(e) => { + eprintln!( + "Error encoding key for PAZ file entry '{}': {}", + self.entry.name, e + ); + crate::COUNTER.inc_error(); + return; + } + }; + let mut rc4 = Rc4::new(&key); + if self.version >= 2 { + let crc = crc32fast::hash(&key); + let skip = ((crc >> 12) as i32) & 0xFF; + rc4.skip_bytes(skip as usize); } + Box::new(Rc4Stream::new(&mut self.inner, rc4)) as Box + } else if self.compress { + let stream = ZlibEncoder::new( + TrackStream::new(&mut self.inner, &mut self.compressed_size), + flate2::Compression::new(self.compress_level), + ); + Box::new(stream) as Box + } else { + Box::new(&mut self.inner) as Box }; - let mut rc4 = Rc4::new(&key); - if self.version >= 2 { - let crc = crc32fast::hash(&key); - let skip = ((crc >> 12) as i32) & 0xFF; - rc4.skip_bytes(skip as usize); + if let Err(e) = stream.write_all(&data) { + eprintln!("Error writing PAZ file entry '{}': {}", self.entry.name, e); + crate::COUNTER.inc_error(); } - Box::new(Rc4Stream::new(&mut self.inner, rc4)) as Box - } else { - Box::new(&mut self.inner) as Box - }; - if let Err(e) = stream.write_all(&data) { - eprintln!("Error writing PAZ file entry '{}': {}", self.entry.name, e); - crate::COUNTER.inc_error(); + } + if self.compress { + self.entry.size = self.compressed_size as u32; + self.entry.aligned_size = (self.entry.size + 7) & !7; } } } diff --git a/src/types.rs b/src/types.rs index a7dd5f8..c89b3ff 100644 --- a/src/types.rs +++ b/src/types.rs @@ -512,6 +512,9 @@ pub struct ExtraConfig { #[cfg(feature = "musica-arc")] /// Musica xor key for paz archive. pub musica_xor_key: Option, + #[cfg(feature = "musica-arc")] + /// Whether to compress files in Musica paz archive when packing paz archive. + pub musica_compress: bool, } #[derive(Clone, Copy, Debug, ValueEnum, PartialEq, Eq, PartialOrd, Ord)]