emote_psb/
writer.rs

1/*
2 * Created on Fri Dec 25 2020
3 *
4 * Copyright (c) storycraft. Licensed under the MIT Licence.
5 */
6
7use std::io::{self, BufReader, Cursor, Read, Seek, SeekFrom, Write};
8
9use adler::Adler32;
10use byteorder::{LittleEndian, WriteBytesExt};
11use flate2::{Compression, bufread::ZlibEncoder};
12
13use crate::{PSB_MDF_SIGNATURE, PSB_SIGNATURE, PsbError, PsbRefs, VirtualPsb, header::MdfHeader, offsets::{PsbOffsets, PsbResourcesOffset, PsbStringOffset}, types::{PsbValue, binary_tree::PsbBinaryTree, collection::PsbUintArray}};
14
15pub struct PsbWriter<T> {
16
17    pub psb: VirtualPsb,
18
19    stream: T
20
21}
22
23impl<T: Write> PsbWriter<T> {
24
25    pub fn write_names(names: &Vec<String>, stream: &mut T) -> Result<u64, PsbError> {
26        let mut buffer_list = Vec::<Vec<u8>>::new();
27
28        for name in names.iter() {
29            buffer_list.push(name.as_bytes().into());
30        }
31
32        PsbBinaryTree::from(buffer_list).write_bytes(stream)
33    }
34
35}
36
37impl<T: Write + Seek> PsbWriter<T> {
38
39    pub fn new(
40        psb: VirtualPsb,
41        stream: T
42    ) -> Self {
43        Self {
44            psb,
45            stream
46        }
47    }
48
49    /// Write file and finish stream
50    pub fn finish(mut self) -> Result<u64, PsbError> {
51        let file_start = self.stream.seek(SeekFrom::Current(0)).unwrap();
52
53        let (header, resources, extra, root) = self.psb.unwrap();
54
55        self.stream.write_u32::<LittleEndian>(PSB_SIGNATURE)?;
56        header.write_bytes(&mut self.stream)?;
57
58        let offsets_end_pos = self.stream.seek(SeekFrom::Current(0)).unwrap() - file_start;
59        self.stream.write_u32::<LittleEndian>(0)?;
60
61        // Offsets
62        let offset_start_pos = self.stream.seek(SeekFrom::Current(0)).unwrap() - file_start;
63        let mut offsets = PsbOffsets::default();
64
65        // Offsets prefill
66        offsets.write_bytes(header.version, &mut self.stream)?;
67
68        let offsets_end = self.stream.seek(SeekFrom::Current(0)).unwrap() - file_start;
69
70        let refs = {
71            let mut names = Vec::new();
72            let mut strings = Vec::new();
73
74            root.collect_names(&mut names);
75            root.collect_strings(&mut strings);
76
77            names.sort();
78            strings.sort();
79
80            PsbRefs::new(names, strings)
81        };
82
83        // Names
84        {
85            offsets.name_offset = (self.stream.seek(SeekFrom::Current(0)).unwrap() - file_start) as u32;
86            Self::write_names(refs.names(), &mut self.stream)?;
87        }
88
89        // Root Entry
90        {
91            offsets.entry_point = (self.stream.seek(SeekFrom::Current(0)).unwrap() - file_start) as u32;
92            PsbValue::Object(root).write_bytes_refs(&mut self.stream, &refs)?;
93        }
94
95        // Strings
96        {
97            let (_, strings) = Self::write_strings(refs.strings(), &mut self.stream)?;
98
99            offsets.strings = strings;
100        }
101
102        // Resources
103        {
104            let (_, res_offsets) = Self::write_resources(&resources, &mut self.stream)?;
105            offsets.resources = res_offsets;
106        }
107
108        // Extra resources support from 4
109        if header.version > 3 {
110            let (_, extra_offsets) = Self::write_resources(&extra, &mut self.stream)?;
111            offsets.extra = Some(extra_offsets);
112        }
113
114        // Rewrite entries
115        let file_end = self.stream.seek(SeekFrom::Current(0)).unwrap();
116
117        self.stream.seek(SeekFrom::Start(offsets_end_pos))?;
118        self.stream.write_u32::<LittleEndian>(offsets_end as u32)?;
119
120        if header.version > 2 {
121            let mut adler = Adler32::new();
122
123            adler.write_slice(&(offset_start_pos as u32).to_le_bytes());
124            adler.write_slice(&offsets.name_offset.to_le_bytes());
125            adler.write_slice(&offsets.strings.offset_pos.to_le_bytes());
126            adler.write_slice(&offsets.strings.data_pos.to_le_bytes());
127            adler.write_slice(&offsets.resources.offset_pos.to_le_bytes());
128            adler.write_slice(&offsets.resources.lengths_pos.to_le_bytes());
129            adler.write_slice(&offsets.resources.data_pos.to_le_bytes());
130            adler.write_slice(&offsets.entry_point.to_le_bytes());
131            
132            offsets.checksum = Some(adler.checksum());
133        }
134
135        self.stream.seek(SeekFrom::Start(offset_start_pos))?;
136        offsets.write_bytes(header.version, &mut self.stream)?;
137
138        self.stream.seek(SeekFrom::Start(file_end))?;
139
140        Ok(file_end - file_start)
141    }
142
143    /// Write resources. Returns written size, PsbResourcesOffset tuple
144    pub fn write_resources(resources: &Vec<Vec<u8>>, stream: &mut T) -> Result<(u64, PsbResourcesOffset), PsbError> {
145        let mut offset_list = Vec::<u64>::new();
146        let mut length_list = Vec::<u64>::new();
147
148        let mut total_len = 0_u64;
149        for res in resources.iter() {
150            let len = res.len() as u64;
151
152            offset_list.push(total_len);
153            length_list.push(len);
154
155            total_len += len;
156        }
157
158        let offset_pos = (stream.seek(SeekFrom::Current(0)).unwrap()) as u32;
159        let offsets_written = PsbValue::IntArray(PsbUintArray::from(offset_list)).write_bytes(stream)?;
160
161        let lengths_pos = (stream.seek(SeekFrom::Current(0)).unwrap()) as u32;
162        let lengths_written = PsbValue::IntArray(PsbUintArray::from(length_list)).write_bytes(stream)?;
163
164        let data_pos = (stream.seek(SeekFrom::Current(0)).unwrap()) as u32;
165        let mut data_written = 0_u64;
166        for res in resources.iter() {
167            data_written += res.len() as u64;
168            stream.write_all(res)?;
169        }
170
171        Ok((offsets_written + lengths_written + data_written, PsbResourcesOffset {
172            offset_pos,
173            lengths_pos,
174            data_pos
175        }))
176    }
177
178    /// Write strings. Returns written size, PsbStringOffset tuple
179    pub fn write_strings(strings: &Vec<String>, stream: &mut T) -> Result<(u64, PsbStringOffset), PsbError> {
180        let mut offset_list = Vec::<u64>::new();
181
182        let mut total_len = 0_u64;
183        for string in strings.iter() {
184            let len = string.as_bytes().len() as u64;
185            
186            offset_list.push(total_len);
187
188            total_len += len + 1;
189        }
190
191        let offset_pos = stream.seek(SeekFrom::Current(0)).unwrap() as u32;
192        let offset_written = PsbValue::IntArray(PsbUintArray::from(offset_list)).write_bytes(stream)?;
193
194        let data_pos = stream.seek(SeekFrom::Current(0)).unwrap() as u32;
195        for string in strings.iter() {
196            stream.write_all(string.as_bytes())?;
197            stream.write_u8(0)?;
198        }
199
200        Ok((offset_written + total_len as u64, PsbStringOffset {
201            offset_pos,
202            data_pos
203        }))
204    }
205
206}
207
208pub struct MdfWriter<R, W> {
209
210    read: R,
211    stream: W
212
213}
214
215impl<R: Read, W: Write + Seek> MdfWriter<R, W> {
216
217    pub fn new(read: R, stream: W) -> Self {
218        Self {
219            read,
220            stream
221        }
222    }
223
224    /// Write mdf file.
225    /// Returns written size
226    pub fn finish(mut self) -> Result<u64, PsbError> {
227        let mut reader = BufReader::new(self.read);
228
229        let mut encoder = ZlibEncoder::new(&mut reader, Compression::best());
230
231        // Write signature first
232        self.stream.write_u32::<LittleEndian>(PSB_MDF_SIGNATURE)?;
233
234        let header_pos = self.stream.seek(SeekFrom::Current(0)).unwrap();
235        // Prefill header
236        MdfHeader { size: 0 }.write_bytes(&mut self.stream)?;
237        
238        io::copy(&mut encoder, &mut self.stream)?;
239        let total_out = encoder.total_out();
240
241        let end_pos = self.stream.seek(SeekFrom::Current(0)).unwrap();
242
243        // Fill header
244        self.stream.seek(SeekFrom::Start(header_pos)).unwrap();
245        MdfHeader { size: total_out as u32 }.write_bytes(&mut self.stream)?;
246
247        self.stream.seek(SeekFrom::Start(end_pos)).unwrap();
248        Ok(total_out + 8)
249    }
250
251}
252
253pub struct PsbMdfWriter<T> {
254
255    buffer: Cursor<Vec<u8>>,
256    psb: VirtualPsb,
257    stream: T
258
259}
260
261impl<T: Write + Seek> PsbMdfWriter<T> {
262
263    pub fn new(psb: VirtualPsb, stream: T) -> Self {
264        Self {
265            buffer: Default::default(),
266            psb,
267            stream
268        }
269    }
270
271    /// Write mdf file.
272    /// Returns written size
273    pub fn finish(mut self) -> Result<u64, PsbError> {
274        let psb_writer = PsbWriter::new(self.psb, &mut self.buffer);
275        psb_writer.finish()?;
276
277        let mdf_writer = MdfWriter::new(self.buffer, self.stream);
278
279        mdf_writer.finish()
280    }
281}
emote_psb/writer.rs

emote_psb/
writer.rs
