ck3_history_extractor/parser/
save_file.rs

1use derive_more::{Display, From};
2use jomini::{
3    self, binary::TokenReader as BinaryTokenReader, text::TokenReader as TextTokenReader,
4};
5use std::{
6    error,
7    fmt::Debug,
8    fs::File,
9    io::{self, Cursor, Read},
10    path::Path,
11    string::FromUtf8Error,
12};
13use zip::{read::ZipArchive, result::ZipError};
14
15use super::types::Tape;
16
17/// The header of an archive within a save file.
18const ARCHIVE_HEADER: &[u8; 4] = b"PK\x03\x04";
19
20const BINARY_HEADER: &[u8; 4] = b"U1\x01\x00";
21
22/// An error that can occur when opening a save file.
23/// Generally things that are the fault of the user, however unintentional those may be
24#[derive(Debug, From, Display)]
25pub enum SaveFileError {
26    /// Something went wrong with stdlib IO.
27    IoError(io::Error),
28    /// We found a problem
29    #[display("{}", _0)]
30    ParseError(&'static str),
31    /// Something went wrong with decompressing the save file.
32    DecompressionError(ZipError),
33    /// Decoding bytes failed
34    DecodingError(FromUtf8Error),
35}
36
37impl error::Error for SaveFileError {
38    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
39        match self {
40            Self::DecompressionError(err) => Some(err),
41            Self::IoError(err) => Some(err),
42            Self::DecodingError(err) => Some(err),
43            Self::ParseError(_) => None,
44        }
45    }
46}
47
48/// A struct that represents a ck3 save file.
49/// It is just a wrapper around the contents of the save file.
50/// This is so that we can abstract away the compression, encoding and just
51/// return an abstract [Tape] that can be used to read from the save file.
52pub struct SaveFile {
53    /// The contents of the save file, shared between all sections
54    contents: Vec<u8>,
55    binary: bool,
56}
57
58impl<'a> SaveFile {
59    /// Open a save file.
60    /// Internally uses [File::open] to open the file and then [SaveFile::read] to read the contents.
61    pub fn open<P: AsRef<Path>>(filename: P) -> Result<SaveFile, SaveFileError> {
62        let mut file = File::open(filename)?;
63        let metadata = file.metadata()?;
64        SaveFile::read(&mut file, Some(metadata.len() as usize))
65    }
66
67    /// Create a new SaveFile instance.
68    ///
69    /// # Compression
70    ///
71    /// The save file can be compressed using the zip format.
72    /// Function will automatically detect if the save file is compressed and decompress it.
73    ///
74    /// # Returns
75    ///
76    /// A new SaveFile instance.
77    /// It is an iterator that returns sections from the save file.
78    pub fn read<F: Read>(
79        file: &mut F,
80        contents_size: Option<usize>,
81    ) -> Result<SaveFile, SaveFileError> {
82        let mut contents = if let Some(size) = contents_size {
83            Vec::with_capacity(size)
84        } else {
85            Vec::new()
86        };
87        let read_size = file.read_to_end(&mut contents)?;
88        if read_size < ARCHIVE_HEADER.len() {
89            return Err(SaveFileError::ParseError("Save file is too small"));
90        }
91        let mut compressed = false;
92        let mut binary = false;
93        // find if ARCHIVE_HEADER is in the file
94        for i in 0..read_size - ARCHIVE_HEADER.len() {
95            if contents[i..i + ARCHIVE_HEADER.len()] == *ARCHIVE_HEADER {
96                compressed = true;
97                break;
98            } else if contents[i..i + BINARY_HEADER.len()] == *BINARY_HEADER {
99                binary = true;
100            }
101        }
102        if compressed {
103            let mut archive = ZipArchive::new(Cursor::new(contents))?;
104            let mut gamestate = archive.by_index(0)?;
105            if gamestate.is_dir() {
106                return Err(SaveFileError::ParseError("Save file is a directory"));
107            }
108            if gamestate.name() != "gamestate" {
109                return Err(SaveFileError::ParseError("Unexpected file name"));
110            }
111            let gamestate_size = gamestate.size() as usize;
112            let mut contents = Vec::with_capacity(gamestate_size);
113            if gamestate.read_to_end(&mut contents)? != gamestate_size {
114                return Err(SaveFileError::ParseError("Failed to read the entire file"));
115            }
116            return Ok(SaveFile { contents, binary });
117        } else {
118            return Ok(SaveFile { contents, binary });
119        }
120    }
121
122    /// Get the tape from the save file.
123    pub fn tape(&'a self) -> Tape<'a> {
124        if self.binary {
125            Tape::Binary(BinaryTokenReader::new(&self.contents))
126        } else {
127            Tape::Text(TextTokenReader::new(&self.contents))
128        }
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use std::io::Write;
135
136    use io::{Seek, SeekFrom};
137    use zip::write::{SimpleFileOptions, ZipWriter};
138
139    use super::*;
140
141    fn create_zipped_test_file(contents: &'static str) -> Cursor<Vec<u8>> {
142        let file = Vec::new();
143        let cur = Cursor::new(file);
144        let mut zip = ZipWriter::new(cur);
145        let options = SimpleFileOptions::default();
146        zip.start_file("gamestate", options).unwrap();
147        if zip.write(contents.as_bytes()).unwrap() != contents.len() {
148            panic!("Failed to write the entire file");
149        }
150        let mut cur = zip.finish().unwrap();
151        cur.seek(SeekFrom::Start(0)).unwrap();
152        return cur;
153    }
154
155    #[test]
156    fn test_open() {
157        let mut file = Cursor::new(b"test");
158        let save = SaveFile::read(&mut file, None).unwrap();
159        assert_eq!(save.contents, b"test");
160    }
161
162    #[test]
163    fn test_compressed_open() {
164        let mut file = create_zipped_test_file("test");
165        let save = SaveFile::read(&mut file, None).unwrap();
166        assert_eq!(save.contents, b"test");
167    }
168
169    #[test]
170    fn test_tape() {
171        let mut file = Cursor::new(b"test=a");
172        let save = SaveFile::read(&mut file, None).unwrap();
173        let tape = save.tape();
174        if let Tape::Binary(_) = tape {
175            panic!("Expected text tape, got binary tape");
176        }
177    }
178}
ck3_history_extractor/parser/save_file.rs

ck3_history_extractor/parser/
save_file.rs
