Mercurial > touhou
annotate formats/src/th06/pbg3.rs @ 770:f6c287745a67
Rust: Add a libtouhou Python wrapper using pyo3
| author | Emmanuel Gil Peyrot <linkmauve@linkmauve.fr> |
|---|---|
| date | Tue, 30 Aug 2022 18:23:55 +0200 |
| parents | cae5f15ca5ed |
| children |
| rev | line source |
|---|---|
| 637 | 1 //! PBG3 archive files handling. |
| 2 //! | |
| 3 //! This module provides classes for handling the PBG3 file format. | |
| 4 //! The PBG3 format is the archive format used by Touhou 6: EoSD. | |
| 5 //! | |
| 6 //! PBG3 files are merely a bitstream composed of a header, a file | |
| 7 //! table, and LZSS-compressed files. | |
| 8 | |
|
757
21b186be2590
Split the Rust version into multiple crates.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
746
diff
changeset
|
9 use touhou_utils::bitstream::BitStream; |
|
21b186be2590
Split the Rust version into multiple crates.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
746
diff
changeset
|
10 use touhou_utils::lzss; |
|
746
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
11 use std::fs::File; |
| 637 | 12 use std::io; |
| 13 use std::collections::hash_map::{self, HashMap}; | |
|
746
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
14 use std::path::Path; |
| 637 | 15 |
| 16 /// Helper struct to handle strings and integers in PBG3 bitstreams. | |
| 17 pub struct PBG3BitStream<R: io::Read + io::Seek> { | |
| 18 bitstream: BitStream<R>, | |
| 19 } | |
| 20 | |
| 21 impl<R: io::Read + io::Seek> PBG3BitStream<R> { | |
| 22 /// Create a bitstream capable of reading u32 and strings. | |
| 23 pub fn new(bitstream: BitStream<R>) -> PBG3BitStream<R> { | |
| 24 PBG3BitStream { | |
| 25 bitstream, | |
| 26 } | |
| 27 } | |
| 28 | |
| 29 /// Seek inside the bitstream, ditching any unused data read. | |
| 30 pub fn seek(&mut self, seek_from: io::SeekFrom) -> io::Result<u64> { | |
| 31 self.bitstream.seek(seek_from) | |
| 32 } | |
| 33 | |
| 34 /// Return the current position in the stream. | |
| 35 pub fn tell(&mut self) -> io::Result<u64> { | |
| 36 self.bitstream.seek(io::SeekFrom::Current(0)) | |
| 37 } | |
| 38 | |
| 39 /// Read a given amount of bits. | |
| 40 pub fn read(&mut self, nb_bits: usize) -> io::Result<usize> { | |
| 41 self.bitstream.read(nb_bits) | |
| 42 } | |
| 43 | |
| 44 /// Read a given amount of bytes. | |
| 45 pub fn read_bytes(&mut self, nb_bytes: usize) -> io::Result<Vec<u8>> { | |
| 46 self.bitstream.read_bytes(nb_bytes) | |
| 47 } | |
| 48 | |
| 49 /// Read an integer from the bitstream. | |
| 50 /// | |
| 51 /// Integers have variable sizes. They begin with a two-bit value indicating | |
| 52 /// the number of (non-aligned) bytes to read. | |
| 53 pub fn read_u32(&mut self) -> io::Result<u32> { | |
| 54 let size = self.read(2)?; | |
| 55 Ok(self.read((size + 1) * 8)? as u32) | |
| 56 } | |
| 57 | |
| 58 /// Read a string from the bitstream. | |
| 59 /// | |
| 60 /// Strings are stored as NULL-terminated sequences of bytes. | |
| 61 /// The only catch is that they are not byte-aligned. | |
| 62 pub fn read_string(&mut self, mut max_size: usize) -> io::Result<Vec<u8>> { | |
| 63 let mut buf = Vec::new(); | |
| 64 while max_size > 0 { | |
| 65 let byte = self.read(8)? as u8; | |
| 66 if byte == 0 { | |
| 67 break; | |
| 68 } | |
| 69 buf.push(byte); | |
| 70 max_size -= 1; | |
| 71 } | |
| 72 Ok(buf) | |
| 73 } | |
| 74 } | |
| 75 | |
| 76 type Entry = (u32, u32, u32, u32, u32); | |
| 77 | |
| 78 /// Handle PBG3 archive files. | |
| 79 /// | |
| 80 /// PBG3 is a file archive format used in Touhou 6: EoSD. | |
| 81 /// This class provides a representation of such files, as well as functions to | |
| 82 /// read and extract files from a PBG3 archive. | |
| 83 pub struct PBG3<R: io::Read + io::Seek> { | |
| 84 /// List of PBG3Entry objects describing files present in the archive. | |
| 85 entries: HashMap<String, Entry>, | |
| 86 | |
| 87 /// PBG3BitStream struct. | |
| 88 bitstream: PBG3BitStream<R>, | |
| 89 } | |
| 90 | |
| 91 impl<R: io::Read + io::Seek> PBG3<R> { | |
| 92 /// Create a PBG3 archive. | |
| 93 fn new(entries: HashMap<String, Entry>, bitstream: PBG3BitStream<R>) -> PBG3<R> { | |
| 94 PBG3 { | |
| 95 entries, | |
| 96 bitstream, | |
| 97 } | |
| 98 } | |
| 99 | |
| 100 /// Open a PBG3 archive. | |
| 101 pub fn from_file(mut file: R) -> io::Result<PBG3<R>> { | |
| 102 let mut magic = [0; 4]; | |
|
767
ccb04468c5fa
formats: Use Read::read_exact() and Write::write_all()
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
761
diff
changeset
|
103 file.read_exact(&mut magic)?; |
| 637 | 104 if &magic != b"PBG3" { |
| 105 return Err(io::Error::new(io::ErrorKind::Other, "Wrong magic!")); | |
| 106 } | |
| 107 | |
| 108 let bitstream = BitStream::new(file); | |
| 109 let mut bitstream = PBG3BitStream::new(bitstream); | |
| 110 let mut entries = HashMap::new(); | |
| 111 | |
| 112 let nb_entries = bitstream.read_u32()?; | |
| 113 let offset = bitstream.read_u32()?; | |
| 114 bitstream.seek(io::SeekFrom::Start(offset as u64))?; | |
| 115 | |
| 116 for _ in 0..nb_entries { | |
| 117 let unknown_1 = bitstream.read_u32()?; | |
| 118 let unknown_2 = bitstream.read_u32()?; | |
| 119 let checksum = bitstream.read_u32()?; // Checksum of *compressed data* | |
| 120 let offset = bitstream.read_u32()?; | |
| 121 let size = bitstream.read_u32()?; | |
| 122 let name = bitstream.read_string(255)?; | |
|
761
f506ad5c9b17
formats: Replace unwrap() with proper io::Error.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
760
diff
changeset
|
123 let name = String::from_utf8(name).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; |
| 637 | 124 entries.insert(name, (unknown_1, unknown_2, checksum, offset, size)); |
| 125 } | |
| 126 | |
| 127 Ok(PBG3::new(entries, bitstream)) | |
| 128 } | |
| 129 | |
| 130 /// List all file entries in this PBG3 archive. | |
| 131 pub fn list_files(&self) -> hash_map::Keys<String, Entry> { | |
| 132 self.entries.keys() | |
| 133 } | |
| 134 | |
| 135 /// Read a single file from this PBG3 archive. | |
|
746
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
136 pub fn get_file(&mut self, filename: &str, check: bool) -> io::Result<Vec<u8>> { |
|
761
f506ad5c9b17
formats: Replace unwrap() with proper io::Error.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
760
diff
changeset
|
137 let (_unknown_1, _unknown_2, checksum, offset, size) = |
|
f506ad5c9b17
formats: Replace unwrap() with proper io::Error.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
760
diff
changeset
|
138 self.entries.get(filename) |
|
f506ad5c9b17
formats: Replace unwrap() with proper io::Error.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
760
diff
changeset
|
139 .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, format!("File not found in PBG3: {}", filename)))?; |
| 637 | 140 self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?; |
| 141 let data = lzss::decompress(&mut self.bitstream.bitstream, *size as usize, 0x2000, 13, 4, 3)?; | |
| 142 if check { | |
| 143 // Verify the checksum. | |
| 144 let compressed_size = self.bitstream.tell()? as u32 - *offset; | |
| 145 self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?; | |
| 146 let mut value: u32 = 0; | |
| 147 for c in self.bitstream.read_bytes(compressed_size as usize)? { | |
|
769
cae5f15ca5ed
formats: Fix possible panic in PBG3 checksum check
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
767
diff
changeset
|
148 value = value.wrapping_add(c as u32); |
| 637 | 149 } |
| 150 if value != *checksum { | |
| 151 return Err(io::Error::new(io::ErrorKind::Other, "Corrupted data!")); | |
| 152 } | |
| 153 } | |
| 154 Ok(data) | |
| 155 } | |
| 156 } | |
| 157 | |
|
746
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
158 /// Open a PBG3 archive from its path. |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
159 pub fn from_path_buffered<P: AsRef<Path>>(path: P) -> io::Result<PBG3<io::BufReader<File>>> { |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
160 let file = File::open(path)?; |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
161 let buf_file = io::BufReader::new(file); |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
162 PBG3::from_file(buf_file) |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
163 } |
|
0ebf6467e4ff
examples: Add a menu example.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
645
diff
changeset
|
164 |
| 637 | 165 #[cfg(test)] |
| 166 mod tests { | |
| 167 use super::*; | |
|
760
eba9a3d0c484
formats: Fix the tests.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
757
diff
changeset
|
168 use std::io::Cursor; |
| 637 | 169 |
| 170 #[test] | |
| 171 fn bitstream() { | |
|
760
eba9a3d0c484
formats: Fix the tests.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
757
diff
changeset
|
172 let data = Cursor::new(b"Hello world!\0"); |
| 637 | 173 let bitstream = BitStream::new(data); |
| 174 let mut pbg3 = PBG3BitStream::new(bitstream); | |
| 175 assert_eq!(pbg3.read_string(42).unwrap(), b"Hello world!"); | |
| 176 } | |
| 177 | |
| 178 #[test] | |
| 179 fn file_present() { | |
|
645
7bde50132735
Don’t hardcode my home directory in tests.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
637
diff
changeset
|
180 let file = File::open("EoSD/MD.DAT").unwrap(); |
| 637 | 181 let file = io::BufReader::new(file); |
| 182 let pbg3 = PBG3::from_file(file).unwrap(); | |
| 183 let files = pbg3.list_files().cloned().collect::<Vec<String>>(); | |
| 184 assert!(files.contains(&String::from("th06_01.pos"))); | |
| 185 } | |
| 186 | |
| 187 #[test] | |
| 188 fn check_all_files() { | |
|
645
7bde50132735
Don’t hardcode my home directory in tests.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
637
diff
changeset
|
189 let file = File::open("EoSD/MD.DAT").unwrap(); |
| 637 | 190 let file = io::BufReader::new(file); |
| 191 let mut pbg3 = PBG3::from_file(file).unwrap(); | |
| 192 let files = pbg3.list_files().cloned().collect::<Vec<String>>(); | |
| 193 for filename in files { | |
|
760
eba9a3d0c484
formats: Fix the tests.
Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
parents:
757
diff
changeset
|
194 pbg3.get_file(&filename, true).unwrap(); |
| 637 | 195 } |
| 196 } | |
| 197 } |