diff src/th06/pbg3.rs @ 637:afa012bb8021

Hello Rust!
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Wed, 03 Jul 2019 16:27:12 +0200
parents
children 7bde50132735
line wrap: on
line diff
new file mode 100644
--- /dev/null
+++ b/src/th06/pbg3.rs
@@ -0,0 +1,190 @@
+//! PBG3 archive files handling.
+//!
+//! This module provides classes for handling the PBG3 file format.
+//! The PBG3 format is the archive format used by Touhou 6: EoSD.
+//!
+//! PBG3 files are merely a bitstream composed of a header, a file
+//! table, and LZSS-compressed files.
+
+use crate::util::bitstream::BitStream;
+use crate::util::lzss;
+use std::io;
+use std::collections::hash_map::{self, HashMap};
+
+/// Helper struct to handle strings and integers in PBG3 bitstreams.
+pub struct PBG3BitStream<R: io::Read + io::Seek> {
+    bitstream: BitStream<R>,
+}
+
+impl<R: io::Read + io::Seek> PBG3BitStream<R> {
+    /// Create a bitstream capable of reading u32 and strings.
+    pub fn new(bitstream: BitStream<R>) -> PBG3BitStream<R> {
+        PBG3BitStream {
+            bitstream,
+        }
+    }
+
+    /// Seek inside the bitstream, ditching any unused data read.
+    pub fn seek(&mut self, seek_from: io::SeekFrom) -> io::Result<u64> {
+        self.bitstream.seek(seek_from)
+    }
+
+    /// Return the current position in the stream.
+    pub fn tell(&mut self) -> io::Result<u64> {
+        self.bitstream.seek(io::SeekFrom::Current(0))
+    }
+
+    /// Read a given amount of bits.
+    pub fn read(&mut self, nb_bits: usize) -> io::Result<usize> {
+        self.bitstream.read(nb_bits)
+    }
+
+    /// Read a given amount of bytes.
+    pub fn read_bytes(&mut self, nb_bytes: usize) -> io::Result<Vec<u8>> {
+        self.bitstream.read_bytes(nb_bytes)
+    }
+
+    /// Read an integer from the bitstream.
+    ///
+    /// Integers have variable sizes. They begin with a two-bit value indicating
+    /// the number of (non-aligned) bytes to read.
+    pub fn read_u32(&mut self) -> io::Result<u32> {
+        let size = self.read(2)?;
+        Ok(self.read((size + 1) * 8)? as u32)
+    }
+
+    /// Read a string from the bitstream.
+    ///
+    /// Strings are stored as NULL-terminated sequences of bytes.
+    /// The only catch is that they are not byte-aligned.
+    pub fn read_string(&mut self, mut max_size: usize) -> io::Result<Vec<u8>> {
+        let mut buf = Vec::new();
+        while max_size > 0 {
+            let byte = self.read(8)? as u8;
+            if byte == 0 {
+                break;
+            }
+            buf.push(byte);
+            max_size -= 1;
+        }
+        Ok(buf)
+    }
+}
+
+type Entry = (u32, u32, u32, u32, u32);
+
+/// Handle PBG3 archive files.
+///
+/// PBG3 is a file archive format used in Touhou 6: EoSD.
+/// This class provides a representation of such files, as well as functions to
+/// read and extract files from a PBG3 archive.
+pub struct PBG3<R: io::Read + io::Seek> {
+    /// List of PBG3Entry objects describing files present in the archive.
+    entries: HashMap<String, Entry>,
+
+    /// PBG3BitStream struct.
+    bitstream: PBG3BitStream<R>,
+}
+
+impl<R: io::Read + io::Seek> PBG3<R> {
+    /// Create a PBG3 archive.
+    fn new(entries: HashMap<String, Entry>, bitstream: PBG3BitStream<R>) -> PBG3<R> {
+        PBG3 {
+            entries,
+            bitstream,
+        }
+    }
+
+    /// Open a PBG3 archive.
+    pub fn from_file(mut file: R) -> io::Result<PBG3<R>> {
+        let mut magic = [0; 4];
+        file.read(&mut magic)?;
+        if &magic != b"PBG3" {
+            return Err(io::Error::new(io::ErrorKind::Other, "Wrong magic!"));
+        }
+
+        let bitstream = BitStream::new(file);
+        let mut bitstream = PBG3BitStream::new(bitstream);
+        let mut entries = HashMap::new();
+
+        let nb_entries = bitstream.read_u32()?;
+        let offset = bitstream.read_u32()?;
+        bitstream.seek(io::SeekFrom::Start(offset as u64))?;
+
+        for _ in 0..nb_entries {
+            let unknown_1 = bitstream.read_u32()?;
+            let unknown_2 = bitstream.read_u32()?;
+            let checksum = bitstream.read_u32()?; // Checksum of *compressed data*
+            let offset = bitstream.read_u32()?;
+            let size = bitstream.read_u32()?;
+            let name = bitstream.read_string(255)?;
+            // XXX: no unwrap!
+            let name = String::from_utf8(name).unwrap();
+            entries.insert(name, (unknown_1, unknown_2, checksum, offset, size));
+        }
+
+        Ok(PBG3::new(entries, bitstream))
+    }
+
+    /// List all file entries in this PBG3 archive.
+    pub fn list_files(&self) -> hash_map::Keys<String, Entry> {
+        self.entries.keys()
+    }
+
+    /// Read a single file from this PBG3 archive.
+    pub fn get_file(&mut self, filename: String, check: bool) -> io::Result<Vec<u8>> {
+        // XXX: no unwrap!
+        let (_unknown_1, _unknown_2, checksum, offset, size) = self.entries.get(&filename).unwrap();
+        self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?;
+        let data = lzss::decompress(&mut self.bitstream.bitstream, *size as usize, 0x2000, 13, 4, 3)?;
+        if check {
+            // Verify the checksum.
+            let compressed_size = self.bitstream.tell()? as u32 - *offset;
+            self.bitstream.seek(io::SeekFrom::Start(*offset as u64))?;
+            let mut value: u32 = 0;
+            for c in self.bitstream.read_bytes(compressed_size as usize)? {
+                value += c as u32;
+                value &= 0xffffffff;
+            }
+            if value != *checksum {
+                return Err(io::Error::new(io::ErrorKind::Other, "Corrupted data!"));
+            }
+        }
+        Ok(data)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::util::SeekableSlice;
+    use std::fs::File;
+
+    #[test]
+    fn bitstream() {
+        let data = SeekableSlice::new(b"Hello world!\0");
+        let bitstream = BitStream::new(data);
+        let mut pbg3 = PBG3BitStream::new(bitstream);
+        assert_eq!(pbg3.read_string(42).unwrap(), b"Hello world!");
+    }
+
+    #[test]
+    fn file_present() {
+        let file = File::open("/home/linkmauve/games/pc/東方/TH06 ~ The Embodiment of Scarlet Devil/MD.DAT").unwrap();
+        let file = io::BufReader::new(file);
+        let pbg3 = PBG3::from_file(file).unwrap();
+        let files = pbg3.list_files().cloned().collect::<Vec<String>>();
+        assert!(files.contains(&String::from("th06_01.pos")));
+    }
+
+    #[test]
+    fn check_all_files() {
+        let file = File::open("/home/linkmauve/games/pc/東方/TH06 ~ The Embodiment of Scarlet Devil/MD.DAT").unwrap();
+        let file = io::BufReader::new(file);
+        let mut pbg3 = PBG3::from_file(file).unwrap();
+        let files = pbg3.list_files().cloned().collect::<Vec<String>>();
+        for filename in files {
+            pbg3.get_file(filename, true).unwrap();
+        }
+    }
+}