msg_tool\scripts\artemis\archive/
pfs.rs

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