diff --git a/src/ext/io.rs b/src/ext/io.rs index 0c1f1e9..96affa3 100644 --- a/src/ext/io.rs +++ b/src/ext/io.rs @@ -1808,6 +1808,25 @@ impl Seek for StreamRegion { } } +impl Write for StreamRegion { + fn write(&mut self, buf: &[u8]) -> Result { + if self.cur_pos + self.start_pos >= self.end_pos { + return Ok(0); // EOF + } + self.stream + .seek(SeekFrom::Start(self.start_pos + self.cur_pos))?; + let bytes_to_write = (self.end_pos - self.start_pos - self.cur_pos) as usize; + let m = buf.len().min(bytes_to_write); + let written = self.stream.write(&buf[..m])?; + self.cur_pos += written as u64; + Ok(written) + } + + fn flush(&mut self) -> Result<()> { + self.stream.flush() + } +} + struct RangeMap { original: (u64, u64), new: (u64, u64), diff --git a/src/scripts/musica/archive/paz.rs b/src/scripts/musica/archive/paz.rs index 507b7e1..367d511 100644 --- a/src/scripts/musica/archive/paz.rs +++ b/src/scripts/musica/archive/paz.rs @@ -670,6 +670,32 @@ impl Archive for PazArcWriter { version: self.schema.version, }; return Ok(Box::new(writer)); + } else if let Some(mov_key) = &self.mov_key { + entry.offset = self.writer.stream_position()?; + let stream = XoredStream::new(&mut self.writer, self.xor_key); + if self.schema.version < 1 { + let stream = TableEncryptedStream::new(stream, mov_key.clone())?; + let writer = MovDataWriter { + inner: Box::new(stream), + entry, + }; + return Ok(Box::new(writer)); + } + let type_key = self + .schema + .get_type_key(&entry, self.is_audio) + .ok_or_else(|| { + anyhow::anyhow!("Data decryption key not found for entry '{}'.", entry.name) + })?; + let writer = MemMovDataKeyWriter { + inner: Box::new(stream), + cache: MemWriter::new(), + type_key: type_key.to_string(), + mov_key: mov_key.clone(), + entry, + encoding: self.encoding, + }; + return Ok(Box::new(writer)); } Err(anyhow::anyhow!("Data encryption key not found.")) } @@ -722,6 +748,59 @@ impl Archive for PazArcWriter { entry, }; return Ok(Box::new(writer)); + } else if let Some(mov_key) = &self.mov_key { + let size = match size { + Some(size) => size, + None => { + return Ok(Box::new(self.new_file(name, None)?)); + } + }; + let entry = self.headers.get_mut(name).ok_or_else(|| { + anyhow::anyhow!("File '{}' not found in PAZ archive headers", name) + })?; + if entry.offset != 0 || entry.size != 0 { + return Err(anyhow::anyhow!( + "File '{}' already exists in PAZ archive", + name + )); + } + entry.offset = self.writer.stream_position()?; + let stream = XoredStream::new(&mut self.writer, self.xor_key); + if self.schema.version < 1 { + let stream = TableEncryptedStream::new(stream, mov_key.clone())?; + let writer = MovDataWriter { + inner: Box::new(stream), + entry, + }; + return Ok(Box::new(writer)); + } + let type_key = self + .schema + .get_type_key(&entry, self.is_audio) + .ok_or_else(|| { + anyhow::anyhow!("Data decryption key not found for entry '{}'.", entry.name) + })?; + let key = format!( + "{} {:08X} {}", + entry.name.to_ascii_lowercase(), + size, + type_key + ); + let key = encode_string(self.encoding, &key, false)?; + let mut rkey = mov_key.clone(); + let key_len = key.len(); + for i in 0..0x100 { + rkey[i] ^= key[i % key_len]; + } + let mut rc4 = Rc4::new(&rkey); + let key_block = rc4.generate_block((size as usize).min(0x10000)); + let region = StreamRegion::new(stream, entry.offset, entry.offset + size)?; + let stream = XoredKeyStream::new(region, key_block, 0); + let writer = DateKeyWriter { + inner: Box::new(stream), + entry, + }; + return Ok(Box::new(writer)); } Err(anyhow::anyhow!("Data encryption key not found.")) } @@ -769,6 +848,54 @@ impl<'a> Drop for DateKeyWriter<'a> { } } +trait MyWriteSeek: Write + Seek {} +impl MyWriteSeek for T {} + +struct MovDataWriter<'a> { + inner: Box, + entry: &'a mut PazEntry, +} + +impl<'a> Write for MovDataWriter<'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> Seek for MovDataWriter<'a> { + fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { + self.inner.seek(pos) + } + + fn rewind(&mut self) -> std::io::Result<()> { + self.inner.rewind() + } + + fn stream_position(&mut self) -> std::io::Result { + self.inner.stream_position() + } +} + +impl<'a> Drop for MovDataWriter<'a> { + fn drop(&mut self) { + if let Ok(pos) = self.inner.stream_position() { + self.entry.unpacked_size = (pos - self.entry.offset) as u32; + self.entry.size = self.entry.unpacked_size; + self.entry.aligned_size = self.entry.size; + } else { + eprintln!( + "Error getting stream position for PAZ file entry '{}'", + self.entry.name + ); + crate::COUNTER.inc_error(); + } + } +} + struct MemDataKeyWriter<'a> { inner: Box, cache: MemWriter, @@ -843,6 +970,92 @@ impl<'a> Drop for MemDataKeyWriter<'a> { } } +struct MemMovDataKeyWriter<'a> { + inner: Box, + cache: MemWriter, + type_key: String, + entry: &'a mut PazEntry, + encoding: Encoding, + mov_key: Vec, +} + +impl<'a> Write for MemMovDataKeyWriter<'a> { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.cache.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.cache.flush() + } +} + +impl<'a> Seek for MemMovDataKeyWriter<'a> { + fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result { + self.cache.seek(pos) + } + + fn rewind(&mut self) -> std::io::Result<()> { + self.cache.rewind() + } + + fn stream_position(&mut self) -> std::io::Result { + self.cache.stream_position() + } +} + +impl<'a> Drop for MemMovDataKeyWriter<'a> { + fn drop(&mut self) { + let data = &self.cache.data; + self.entry.unpacked_size = data.len() as u32; + self.entry.size = self.entry.unpacked_size; + self.entry.aligned_size = self.entry.size; + let key = format!( + "{} {:08X} {}", + self.entry.name.to_ascii_lowercase(), + self.entry.unpacked_size, + self.type_key + ); + 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 rkey = self.mov_key.clone(); + let key_len = key.len(); + for i in 0..0x100 { + rkey[i] ^= key[i % key_len]; + } + let mut rc4 = Rc4::new(&rkey); + let key_block = rc4.generate_block(data.len().min(0x10000)); + let region = match StreamRegion::new( + &mut self.inner, + self.entry.offset, + self.entry.offset + self.entry.size as u64, + ) { + Ok(region) => region, + Err(e) => { + eprintln!( + "Error creating stream region for PAZ file entry '{}': {}", + self.entry.name, e + ); + crate::COUNTER.inc_error(); + return; + } + }; + let mut stream = XoredKeyStream::new(region, key_block, 0); + if let Err(e) = stream.write_all(&data) { + eprintln!("Error writing PAZ file entry '{}': {}", self.entry.name, e); + crate::COUNTER.inc_error(); + } + } +} + #[test] fn test_deserialize_paz() { for (game, schema) in PAZ_SCHEMA.iter() {