emote_psb/
reader.rs

1/*
2 * Created on Fri Dec 25 2020
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7use std::io::{BufRead, BufReader, Cursor, Read, Seek, SeekFrom};
8
9use byteorder::{ReadBytesExt, LittleEndian};
10use encoding::{Encoding, all::UTF_8};
11use flate2::read::ZlibDecoder;
12
13use crate::{PSB_MDF_SIGNATURE, PSB_SIGNATURE, PsbError, PsbErrorKind, PsbFile, PsbRefs, header::{MdfHeader, PsbHeader}, offsets::PsbOffsets, types::{PsbValue, binary_tree::PsbBinaryTree}};
14
15pub struct PsbReader;
16
17impl PsbReader {
18
19    /// Read as PsbFile
20    pub fn open_psb<T: Read + Seek>(mut stream: T) -> Result<PsbFile<T>, PsbError> {
21        let start = stream.seek(SeekFrom::Current(0)).unwrap();
22
23        let signature = stream.read_u32::<LittleEndian>()?;
24        if signature != PSB_SIGNATURE {
25            return Err(PsbError::new(PsbErrorKind::InvalidFile, None));
26        }
27
28        let (_, header) = PsbHeader::from_bytes(&mut stream)?;
29
30        let _ = stream.read_u32::<LittleEndian>()?;
31
32        // offsets
33        let (_, offsets) = PsbOffsets::from_bytes(header.version, &mut stream)?;
34
35        stream.seek(SeekFrom::Start(start + offsets.name_offset as u64))?;
36        let (_, names) = Self::read_names(&mut stream)?;
37
38        stream.seek(SeekFrom::Start(start + offsets.strings.offset_pos as u64))?;
39        let (_, strings) = Self::read_strings(offsets.strings.data_pos + start as u32, &mut stream)?;
40
41        let refs = PsbRefs::new(names, strings);
42
43        Ok(
44            PsbFile::new(
45                header,
46                refs,
47                offsets,
48                stream
49            )
50        )
51    }
52
53    pub fn read_names<T: Read + Seek>(stream: &mut T) -> Result<(u64, Vec<String>), PsbError> {
54        let mut names = Vec::<String>::new();
55
56        let (read, btree) = PsbBinaryTree::from_bytes(stream)?;
57
58        for raw_string in btree.unwrap() {
59            let name = UTF_8.decode(&raw_string, encoding::DecoderTrap::Replace).unwrap();
60            
61            names.push(name);
62        }
63
64        Ok((read, names))
65    }
66
67    pub fn read_strings<T: Read + Seek>(data_pos: u32, stream: &mut T) -> Result<(u64, Vec<String>), PsbError> {
68        let mut strings = Vec::<String>::new();
69
70        let (offsets_read, string_offsets) = match PsbValue::from_bytes(stream)? {
71    
72            (read, PsbValue::IntArray(array)) => Ok((read, array)),
73
74            _ => Err(PsbError::new(PsbErrorKind::InvalidOffsetTable, None))
75
76        }?;
77
78        let mut reader = BufReader::new(stream.by_ref());
79        let string_offsets = string_offsets.unwrap();
80
81        let mut read = 0_usize;
82        for offset in string_offsets {
83            let mut buffer = Vec::new();
84
85            reader.seek(SeekFrom::Start(data_pos as u64 + offset as u64))?;
86            read += reader.read_until(0x00, &mut buffer)?;
87
88            // Decode excluding nul
89            let string = UTF_8.decode(&buffer[..buffer.len() - 1], encoding::DecoderTrap::Replace).unwrap();
90            strings.push(string);
91        }
92
93        Ok((offsets_read + read as u64, strings))
94    }
95
96}
97
98pub struct MdfReader;
99
100impl MdfReader {
101    
102    pub fn open_mdf<T: Read + Seek>(mut stream: T) -> Result<PsbFile<Cursor<Vec<u8>>>, PsbError> {
103        let signature = stream.read_u32::<LittleEndian>()?;
104        if signature != PSB_MDF_SIGNATURE {
105            return Err(PsbError::new(PsbErrorKind::InvalidFile, None));
106        }
107
108        let (_, mdf_header) = MdfHeader::from_bytes(&mut stream)?;
109
110        let mut compressed_buffer = Vec::new();
111
112        stream.take(mdf_header.size as u64).read_to_end(&mut compressed_buffer)?;
113
114        let mut decoder = ZlibDecoder::new(&compressed_buffer[..]);
115
116        let mut buffer = Vec::new();
117        decoder.read_to_end(&mut buffer)?;
118
119        PsbReader::open_psb(Cursor::new(buffer))
120    }
121
122}
emote_psb/reader.rs

emote_psb/
reader.rs
