msg_tool\scripts\artemis\archive/
pfs.rs1use super::detect_script_type;
3use crate::ext::io::*;
4use crate::scripts::base::*;
5use crate::types::*;
6use crate::utils::struct_pack::*;
7use anyhow::Result;
8use digest_io::IoWrapper;
9use msg_tool_macro::*;
10use sha1::Digest;
11use std::collections::HashMap;
12use std::io::{Read, Seek, SeekFrom, Write};
13use std::sync::{Arc, Mutex};
14
15#[derive(Debug)]
16pub struct ArtemisArcBuilder {}
18
19impl ArtemisArcBuilder {
20 pub fn new() -> Self {
22 ArtemisArcBuilder {}
23 }
24}
25
26impl ScriptBuilder for ArtemisArcBuilder {
27 fn default_encoding(&self) -> Encoding {
28 Encoding::Utf8
29 }
30
31 fn default_archive_encoding(&self) -> Option<Encoding> {
32 Some(Encoding::Utf8)
33 }
34
35 fn build_script(
36 &self,
37 buf: Vec<u8>,
38 filename: &str,
39 _encoding: Encoding,
40 archive_encoding: Encoding,
41 config: &ExtraConfig,
42 _archive: Option<&Box<dyn Script>>,
43 ) -> Result<Box<dyn Script + Send + Sync>> {
44 Ok(Box::new(ArtemisArc::new(
45 MemReader::new(buf),
46 archive_encoding,
47 config,
48 filename,
49 )?))
50 }
51
52 fn build_script_from_file(
53 &self,
54 filename: &str,
55 _encoding: Encoding,
56 archive_encoding: Encoding,
57 config: &ExtraConfig,
58 _archive: Option<&Box<dyn Script>>,
59 ) -> Result<Box<dyn Script + Send + Sync>> {
60 let f = std::fs::File::open(filename)?;
61 let f = std::io::BufReader::new(f);
62 Ok(Box::new(ArtemisArc::new(
63 f,
64 archive_encoding,
65 config,
66 filename,
67 )?))
68 }
69
70 fn build_script_from_reader<'a>(
71 &self,
72 reader: Box<dyn ReadSeek + Send + Sync + 'a>,
73 filename: &str,
74 _encoding: Encoding,
75 archive_encoding: Encoding,
76 config: &ExtraConfig,
77 _archive: Option<&Box<dyn Script>>,
78 ) -> Result<Box<dyn Script + Send + Sync + 'a>> {
79 Ok(Box::new(ArtemisArc::new(
80 reader,
81 archive_encoding,
82 config,
83 filename,
84 )?))
85 }
86
87 fn extensions(&self) -> &'static [&'static str] {
88 gen_artemis_arc_ext!()
89 }
90
91 fn is_archive(&self) -> bool {
92 true
93 }
94
95 fn script_type(&self) -> &'static ScriptType {
96 &ScriptType::ArtemisArc
97 }
98
99 fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
100 if buf_len >= 3 && (buf.starts_with(b"pf6") || buf.starts_with(b"pf8")) {
101 return Some(10);
102 }
103 None
104 }
105
106 fn create_archive(
107 &self,
108 filename: &str,
109 files: &[&str],
110 encoding: Encoding,
111 config: &ExtraConfig,
112 ) -> Result<Box<dyn Archive>> {
113 let f = std::fs::File::options()
114 .write(true)
115 .read(true)
116 .create(true)
117 .truncate(true)
118 .open(filename)?;
119 Ok(Box::new(ArtemisArcWriter::new(f, files, encoding, config)?))
120 }
121}
122
123#[derive(Debug, Clone, StructPack, StructUnpack)]
124struct PfsEntryHeader {
125 #[pstring(u32)]
126 name: String,
127 _unk: u32,
128 offset: u32,
129 size: u32,
130}
131
132#[derive(Debug)]
133pub struct ArtemisArc<'a, T: Read + Seek + std::fmt::Debug + Send + Sync + 'a> {
135 reader: Arc<Mutex<T>>,
136 entries: Vec<PfsEntryHeader>,
137 xor_key: Option<[u8; 20]>,
138 output_ext: Option<String>,
139 _mark: std::marker::PhantomData<&'a ()>,
140}
141
142impl<'b, T: Read + Seek + std::fmt::Debug + Send + Sync + 'b> ArtemisArc<'b, T> {
143 pub fn new(
150 mut reader: T,
151 archive_encoding: Encoding,
152 _config: &ExtraConfig,
153 filename: &str,
154 ) -> Result<Self> {
155 let mut magic = [0; 2];
156 reader.read_exact(&mut magic)?;
157 if &magic != b"pf" {
158 return Err(anyhow::anyhow!(
159 "Invalid Artemis archive magic: {:?}",
160 magic
161 ));
162 }
163 let version = reader.read_u8()?;
164 if version != b'6' && version != b'8' {
165 return Err(anyhow::anyhow!(
166 "Unsupported Artemis archive version: {}",
167 version
168 ));
169 }
170 let index_size = reader.read_u32()?;
171 let file_count = reader.read_u32()?;
172 let mut entries = Vec::with_capacity(file_count as usize);
173 for _ in 0..file_count {
174 let header = reader.read_struct(false, archive_encoding, &None)?;
175 entries.push(header);
176 }
177 let xor_key = if version == b'8' {
178 reader.seek(SeekFrom::Start(7))?;
179 let mut sha = IoWrapper(sha1::Sha1::default());
180 let ra = &mut reader;
181 let mut r = ra.take(index_size as u64);
182 std::io::copy(&mut r, &mut sha)?;
183 sha.flush()?;
184 let result = sha.0.finalize();
185 let mut xor_key = [0u8; 20];
186 xor_key.copy_from_slice(&result);
187 Some(xor_key)
188 } else {
189 None
190 };
191 let output_ext = std::path::Path::new(filename)
192 .extension()
193 .filter(|s| *s != "pfs")
194 .map(|s| s.to_string_lossy().to_string());
195 Ok(ArtemisArc {
196 reader: Arc::new(Mutex::new(reader)),
197 entries,
198 xor_key,
199 output_ext,
200 _mark: std::marker::PhantomData,
201 })
202 }
203}
204
205impl<'b, T: Read + Seek + std::fmt::Debug + Send + Sync + 'b> Script for ArtemisArc<'b, T> {
206 fn default_output_script_type(&self) -> OutputScriptType {
207 OutputScriptType::Json
208 }
209
210 fn default_format_type(&self) -> FormatOptions {
211 FormatOptions::None
212 }
213
214 fn is_archive(&self) -> bool {
215 true
216 }
217
218 fn iter_archive_filename<'a>(
219 &'a self,
220 ) -> Result<Box<dyn Iterator<Item = Result<String>> + 'a>> {
221 Ok(Box::new(
222 self.entries.iter().map(|header| Ok(header.name.clone())),
223 ))
224 }
225
226 fn iter_archive_offset<'a>(&'a self) -> Result<Box<dyn Iterator<Item = Result<u64>> + 'a>> {
227 Ok(Box::new(
228 self.entries.iter().map(|header| Ok(header.offset as u64)),
229 ))
230 }
231
232 fn open_file<'a>(&'a self, index: usize) -> Result<Box<dyn ArchiveContent + Send + Sync + 'a>> {
233 if index >= self.entries.len() {
234 return Err(anyhow::anyhow!(
235 "Index out of bounds: {} (max: {})",
236 index,
237 self.entries.len()
238 ));
239 }
240 let header = &self.entries[index];
241 let mut entry = Entry {
242 header: header.clone(),
243 reader: self.reader.clone(),
244 pos: 0,
245 script_type: None,
246 xor_key: self.xor_key.clone(),
247 };
248 let mut header = [0; 0x20];
249 let readed = entry.read(&mut header)?;
250 entry.pos = 0;
251 entry.script_type = detect_script_type(&header, readed, &entry.header.name);
252 Ok(Box::new(entry))
253 }
254
255 fn archive_output_ext<'a>(&'a self) -> Option<&'a str> {
256 self.output_ext.as_ref().map(|s| s.as_str())
257 }
258}
259
260#[derive(Debug)]
261struct Entry<T: Read + Seek + std::fmt::Debug> {
262 header: PfsEntryHeader,
263 reader: Arc<Mutex<T>>,
264 pos: u64,
265 script_type: Option<ScriptType>,
266 xor_key: Option<[u8; 20]>,
267}
268
269impl<T: Read + Seek + std::fmt::Debug + Send + Sync> ArchiveContent for Entry<T> {
270 fn name(&self) -> &str {
271 &self.header.name
272 }
273
274 fn size(&self) -> Option<u64> {
275 Some(self.header.size as u64)
276 }
277
278 fn script_type(&self) -> Option<&ScriptType> {
279 self.script_type.as_ref()
280 }
281
282 fn to_data<'a>(&'a mut self) -> Result<Box<dyn ReadSeek + Send + Sync + 'a>> {
283 Ok(Box::new(self))
284 }
285}
286
287impl<T: Read + Seek + std::fmt::Debug> Read for Entry<T> {
288 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
289 let mut reader = self.reader.lock().map_err(|e| {
290 std::io::Error::new(
291 std::io::ErrorKind::Other,
292 format!("Failed to lock mutex: {}", e),
293 )
294 })?;
295 reader.seek(SeekFrom::Start(self.header.offset as u64 + self.pos))?;
296 let bytes_read = buf.len().min(self.header.size as usize - self.pos as usize);
297 if bytes_read == 0 {
298 return Ok(0);
299 }
300 let bytes_read = reader.read(&mut buf[..bytes_read])?;
301 if let Some(xor_key) = &self.xor_key {
302 for i in 0..bytes_read {
303 let l = (self.pos + i as u64) % 20;
304 buf[i] ^= xor_key[l as usize];
305 }
306 }
307 self.pos += bytes_read as u64;
308 Ok(bytes_read)
309 }
310}
311
312impl<T: Read + Seek + std::fmt::Debug> Seek for Entry<T> {
313 fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
314 let new_pos = match pos {
315 SeekFrom::Start(offset) => offset,
316 SeekFrom::End(offset) => {
317 if offset < 0 {
318 if (-offset) as u64 > self.header.size as u64 {
319 return Err(std::io::Error::new(
320 std::io::ErrorKind::InvalidInput,
321 "Seek from end exceeds file length",
322 ));
323 }
324 self.header.size as u64 - (-offset) as u64
325 } else {
326 self.header.size as u64 + offset as u64
327 }
328 }
329 SeekFrom::Current(offset) => {
330 if offset < 0 {
331 if (-offset) as u64 > self.pos {
332 return Err(std::io::Error::new(
333 std::io::ErrorKind::InvalidInput,
334 "Seek from current exceeds current position",
335 ));
336 }
337 self.pos.saturating_sub((-offset) as u64)
338 } else {
339 self.pos + offset as u64
340 }
341 }
342 };
343 self.pos = new_pos;
344 Ok(self.pos)
345 }
346
347 fn stream_position(&mut self) -> std::io::Result<u64> {
348 Ok(self.pos)
349 }
350}
351
352pub struct ArtemisArcWriter<T: Write + Seek + Read> {
354 writer: T,
355 headers: HashMap<String, PfsEntryHeader>,
356 encoding: Encoding,
357 disable_xor: bool,
358 index_size: u32,
359}
360
361impl<T: Write + Seek + Read> ArtemisArcWriter<T> {
362 pub fn new(
369 mut writer: T,
370 files: &[&str],
371 encoding: Encoding,
372 config: &ExtraConfig,
373 ) -> Result<Self> {
374 writer.write_all(if config.artemis_arc_disable_xor {
375 b"pf6"
376 } else {
377 b"pf8"
378 })?;
379 writer.write_u32(0)?; writer.write_u32(files.len() as u32)?;
381 let mut headers = HashMap::new();
382 for file in files {
383 let header = PfsEntryHeader {
384 name: file.to_string(),
385 _unk: 0,
386 offset: 0,
387 size: 0,
388 };
389 header.pack(&mut writer, false, encoding, &None)?;
390 headers.insert(file.to_string(), header);
391 }
392 let size = writer.stream_position()?;
393 let index_size = size as u32 - 7;
394 writer.write_u32_at(3, index_size)?;
395 Ok(ArtemisArcWriter {
396 writer,
397 headers,
398 encoding,
399 disable_xor: config.artemis_arc_disable_xor,
400 index_size,
401 })
402 }
403}
404
405impl<T: Write + Seek + Read> Archive for ArtemisArcWriter<T> {
406 fn new_file<'a>(
407 &'a mut self,
408 name: &str,
409 _size: Option<u64>,
410 ) -> Result<Box<dyn WriteSeek + 'a>> {
411 let entry = self
412 .headers
413 .get_mut(name)
414 .ok_or_else(|| anyhow::anyhow!("File '{}' not found in archive", name))?;
415 if entry.offset != 0 || entry.size != 0 {
416 return Err(anyhow::anyhow!("File '{}' already exists in archive", name));
417 }
418 self.writer.seek(SeekFrom::End(0))?;
419 entry.offset = self.writer.stream_position()? as u32;
420 let file = ArtemisArcFile {
421 header: entry,
422 writer: &mut self.writer,
423 pos: 0,
424 };
425 Ok(Box::new(file))
426 }
427
428 fn write_header(&mut self) -> Result<()> {
429 self.writer.seek(SeekFrom::Start(11))?;
430 let mut files = self.headers.values().collect::<Vec<_>>();
431 files.sort_by_key(|d| d.offset);
432 for file in files.iter() {
433 file.pack(&mut self.writer, false, self.encoding, &None)?;
434 }
435 if !self.disable_xor {
436 self.writer.seek(SeekFrom::Start(7))?;
437 let mut sha = IoWrapper(sha1::Sha1::default());
438 let w = &mut self.writer;
439 let mut header = w.take(self.index_size as u64);
440 std::io::copy(&mut header, &mut sha)?;
441 sha.flush()?;
442 let result = sha.0.finalize();
443 let mut xor_key = [0u8; 20];
444 xor_key.copy_from_slice(&result);
445 let mut buf = [0u8; 1024];
446 for file in files.iter() {
447 self.writer.seek(SeekFrom::Start(file.offset as u64))?;
448 let mut pos = 0u32;
449 while pos < file.size {
450 let bytes_to_read = (file.size - pos).min(1024) as usize;
451 let bytes_read = self.writer.read(&mut buf[..bytes_to_read])?;
452 if bytes_read == 0 {
453 return Err(anyhow::anyhow!(
454 "Unexpected end of file while reading '{}'",
455 file.name
456 ));
457 }
458 for i in 0..bytes_read {
459 let l = (pos as u64 + i as u64) % 20;
460 buf[i] ^= xor_key[l as usize];
461 }
462 self.writer.seek_relative(-(bytes_read as i64))?;
463 self.writer.write_all(&buf[..bytes_read])?;
464 pos += bytes_read as u32;
465 }
466 }
467 }
468 Ok(())
469 }
470}
471
472pub struct ArtemisArcFile<'a, T: Write + Seek> {
474 header: &'a mut PfsEntryHeader,
475 writer: &'a mut T,
476 pos: u64,
477}
478
479impl<'a, T: Write + Seek> Write for ArtemisArcFile<'a, T> {
480 fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
481 self.writer
482 .seek(SeekFrom::Start(self.header.offset as u64 + self.pos))?;
483 let bytes_written = self.writer.write(buf)?;
484 self.pos += bytes_written as u64;
485 self.header.size = self.header.size.max(self.pos as u32);
486 Ok(bytes_written)
487 }
488
489 fn flush(&mut self) -> std::io::Result<()> {
490 self.writer.flush()
491 }
492}
493
494impl<'a, T: Write + Seek> Seek for ArtemisArcFile<'a, T> {
495 fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
496 let new_pos = match pos {
497 SeekFrom::Start(offset) => offset,
498 SeekFrom::End(offset) => {
499 if offset < 0 {
500 if (-offset) as u64 > self.header.size as u64 {
501 return Err(std::io::Error::new(
502 std::io::ErrorKind::InvalidInput,
503 "Seek from end exceeds file length",
504 ));
505 }
506 self.header.size as u64 - (-offset) as u64
507 } else {
508 self.header.size as u64 + offset as u64
509 }
510 }
511 SeekFrom::Current(offset) => {
512 if offset < 0 {
513 if (-offset) as u64 > self.pos {
514 return Err(std::io::Error::new(
515 std::io::ErrorKind::InvalidInput,
516 "Seek from current exceeds current position",
517 ));
518 }
519 self.pos.saturating_sub((-offset) as u64)
520 } else {
521 self.pos + offset as u64
522 }
523 }
524 };
525 self.pos = new_pos;
526 Ok(self.pos)
527 }
528
529 fn stream_position(&mut self) -> std::io::Result<u64> {
530 Ok(self.pos)
531 }
532}