diff --git a/src/ext/io.rs b/src/ext/io.rs index 57880b1..0b13c7b 100644 --- a/src/ext/io.rs +++ b/src/ext/io.rs @@ -1719,6 +1719,7 @@ impl CPeek for MemWriter { } /// A region of a stream that can be read/write and seeked within a specified range. +#[derive(Debug)] pub struct StreamRegion { stream: T, start_pos: u64, @@ -2001,7 +2002,7 @@ impl Result, O: Fn(u64) -> R } /// A thread-safe wrapper around a Mutex-protected writer/reader. -#[derive(Clone)] +#[derive(Debug)] pub struct MutexWrapper { inner: Arc>, pos: u64, @@ -2078,3 +2079,90 @@ impl Write for EmptyWriter { Ok(()) } } + +#[derive(Debug)] +/// A readable stream that starts with a given prefix before the actual data. +pub struct PrefixStream { + prefix: Vec, + pos: usize, + inner: T, +} + +impl PrefixStream { + /// Creates a new `PrefixStream` with the given prefix and inner stream. + pub fn new(prefix: Vec, inner: T) -> Self { + PrefixStream { + prefix, + pos: 0, + inner, + } + } +} + +impl Read for PrefixStream { + fn read(&mut self, buf: &mut [u8]) -> Result { + if self.pos < self.prefix.len() { + let bytes_to_read = std::cmp::min(buf.len(), self.prefix.len() - self.pos); + buf[..bytes_to_read].copy_from_slice(&self.prefix[self.pos..self.pos + bytes_to_read]); + self.pos += bytes_to_read; + Ok(bytes_to_read) + } else { + self.inner.read(buf) + } + } +} + +impl Seek for PrefixStream { + fn seek(&mut self, pos: SeekFrom) -> Result { + let prefix_len = self.prefix.len() as u64; + let new_pos = match pos { + SeekFrom::Start(offset) => offset, + SeekFrom::End(offset) => { + let inner_len = self.inner.stream_length()?; + if offset < 0 { + if (-offset) as u64 > inner_len + prefix_len { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Seek position is before the start of the stream", + )); + } + inner_len + prefix_len - (-offset) as u64 + } else { + inner_len + prefix_len + offset as u64 + } + } + SeekFrom::Current(offset) => { + let current_pos = self.stream_position()?; + if offset < 0 { + if (-offset) as u64 > current_pos + prefix_len { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Seek position is before the start of the stream", + )); + } + prefix_len + current_pos - (-offset) as u64 + } else { + prefix_len + current_pos + offset as u64 + } + } + }; + if new_pos < prefix_len { + self.pos = new_pos as usize; + self.inner.rewind()?; + } else { + self.pos = self.prefix.len(); + self.inner.seek(SeekFrom::Start(new_pos - prefix_len))?; + } + Ok(new_pos) + } + + fn stream_position(&mut self) -> Result { + Ok(self.pos as u64 + self.inner.stream_position()?) + } + + fn rewind(&mut self) -> Result<()> { + self.pos = 0; + self.inner.rewind()?; + Ok(()) + } +} diff --git a/src/scripts/kirikiri/archive/xp3.rs b/src/scripts/kirikiri/archive/xp3.rs index 5d1c918..7537a89 100644 --- a/src/scripts/kirikiri/archive/xp3.rs +++ b/src/scripts/kirikiri/archive/xp3.rs @@ -3,7 +3,8 @@ use crate::scripts::base::*; use crate::types::*; use anyhow::Result; use flate2::read::ZlibDecoder; -use std::io::{Read, Seek, Take}; +use overf::wrapping; +use std::io::{Read, Seek, SeekFrom, Take}; use std::sync::{Arc, Mutex}; use xp3::XP3Reader; use xp3::index::file::{IndexSegmentFlag, XP3FileIndex}; @@ -138,12 +139,32 @@ impl Script for Xp3Archive { .ok_or(anyhow::anyhow!("Index out of bounds: {}", index))? .1 .clone(); - let entry = Entry::new(self.reader.clone(), index); + let mut entry = Entry::new(self.reader.clone(), index); + let mut header = [0u8; 16]; + let header_len = entry.read(&mut header)?; + entry.rewind()?; + if header_len >= 5 + && header[0] == 0xFE + && header[1] == 0xFE + && header[3] == 0xFF + && header[4] == 0xFE + { + let crypt = header[2]; + if crypt == 2 { + let index = entry.index.clone(); + return Ok(Box::new(SimpleCryptZlib::new(entry, index)?)); + } + if matches!(crypt, 0 | 1) { + let index = entry.index.clone(); + return Ok(Box::new(SimpleCrypt::new(entry, index, crypt)?)); + } + } Ok(Box::new(entry)) } } -struct Entry { +#[derive(Debug)] +struct Entry { reader: Arc>, index: XP3FileIndex, cache: Option>>>, @@ -151,7 +172,7 @@ struct Entry { entries_pos: Vec, } -impl Entry { +impl Entry { fn new(reader: Arc>, index: XP3FileIndex) -> Self { let mut pos = 0; let entries_pos = index @@ -173,13 +194,17 @@ impl Entry { } } -impl ArchiveContent for Entry { +impl ArchiveContent for Entry { fn name(&self) -> &str { &self.index.info().name() } + + fn to_data<'a>(&'a mut self) -> Result> { + Ok(Box::new(self)) + } } -impl Read for Entry { +impl Read for Entry { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { if self.pos >= self.index.info().file_size() { self.cache.take(); @@ -231,3 +256,208 @@ impl Read for Entry { } } } + +impl Seek for Entry { + fn seek(&mut self, pos: SeekFrom) -> std::io::Result { + let new_pos = match pos { + SeekFrom::Start(p) => p, + SeekFrom::End(offset) => { + if offset < 0 { + if (-offset) as u64 > self.index.info().file_size() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Seek from end exceeds file length", + )); + } + self.index.info().file_size() - (-offset) as u64 + } else { + self.index.info().file_size() + offset as u64 + } + } + SeekFrom::Current(offset) => { + if offset < 0 { + if (-offset) as u64 > self.pos { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "Seek from current exceeds file start", + )); + } + self.pos - (-offset) as u64 + } else { + self.pos + offset as u64 + } + } + }; + if let Some(cache) = self.cache.as_mut() { + let old_seg_index = match self.entries_pos.binary_search(&self.pos) { + Ok(i) => i, + Err(i) => { + if i == 0 { + 0 + } else { + i - 1 + } + } + }; + let new_seg_index = match self.entries_pos.binary_search(&new_pos) { + Ok(i) => i, + Err(i) => { + if i == 0 { + 0 + } else { + i - 1 + } + } + }; + if old_seg_index != new_seg_index { + self.cache.take(); + } else { + if new_pos >= self.pos { + let skip_pos = new_pos - self.pos; + let mut e = EmptyWriter::new(); + std::io::copy(&mut cache.take(skip_pos), &mut e)?; // skip + } else { + self.cache.take(); + } + } + } + self.pos = new_pos; + Ok(self.pos) + } + + fn rewind(&mut self) -> std::io::Result<()> { + self.pos = 0; + self.cache.take(); + Ok(()) + } + + fn stream_position(&mut self) -> std::io::Result { + Ok(self.pos) + } +} + +struct SimpleCryptZlib { + inner: PrefixStream>>>, + index: XP3FileIndex, +} + +impl SimpleCryptZlib { + fn new(mut entry: Entry, index: XP3FileIndex) -> Result { + entry.seek(SeekFrom::Start(0x15))?; + let entry = StreamRegion::new(entry, 0x15, index.info().file_size())?; + let inner = PrefixStream::new(vec![0xFF, 0xFE], ZlibDecoder::new(entry)); + Ok(Self { inner, index }) + } +} + +impl ArchiveContent for SimpleCryptZlib { + fn name(&self) -> &str { + &self.index.info().name() + } +} + +impl Read for SimpleCryptZlib { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.inner.read(buf) + } +} + +#[derive(Debug)] +struct SimpleCryptInner { + inner: StreamRegion>, + crypt: u8, +} + +impl SimpleCryptInner { + fn new(mut entry: Entry, crypt: u8) -> Result { + entry.seek(SeekFrom::Start(5))?; + let size = entry.index.info().file_size(); + let entry = StreamRegion::new(entry, 5, size)?; + Ok(Self { + inner: entry, + crypt, + }) + } +} + +impl Read for SimpleCryptInner { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + let readed = self.inner.read(buf)?; + match self.crypt { + 0 => { + for b in &mut buf[..readed] { + let ch = *b as u16; + if ch >= 20 { + *b = wrapping! {ch ^ (((ch & 0xfe) << 8) ^ 1)} as u8; + } + } + } + 1 => { + for b in &mut buf[..readed] { + let mut ch = *b as u32; + ch = wrapping! {((ch & 0xaaaaaaaa) >> 1) | ((ch & 0x55555555) << 1)}; + *b = ch as u8; + } + } + _ => {} + } + Ok(readed) + } +} + +impl Seek for SimpleCryptInner { + fn seek(&mut self, pos: 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() + } +} + +#[derive(Debug)] +struct SimpleCrypt { + inner: PrefixStream>, + index: XP3FileIndex, +} + +impl SimpleCrypt { + fn new(entry: Entry, index: XP3FileIndex, crypt: u8) -> Result { + let inner = PrefixStream::new(vec![0xFF, 0xFE], SimpleCryptInner::new(entry, crypt)?); + Ok(Self { inner, index }) + } +} + +impl ArchiveContent for SimpleCrypt { + fn name(&self) -> &str { + &self.index.info().name() + } + + fn to_data<'a>(&'a mut self) -> Result> { + Ok(Box::new(self)) + } +} + +impl Read for SimpleCrypt { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.inner.read(buf) + } +} + +impl Seek for SimpleCrypt { + fn seek(&mut self, pos: 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() + } +}