msg_tool\scripts\artemis\archive/
pfs.rs

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