rust_unreal_unpak/pak/
pak_reader.rs

1use std::io::{Cursor, Read, Seek, SeekFrom};
2use thiserror::Error;
3use tokio_byteorder::{LittleEndian, BigEndian, AsyncReadBytesExt};
4use std::collections::HashMap;
5
6use super::pak_file::PakInfo;
7use super::{PakVersions, PakVersionSizes};
8use crate::cursor_ext::{CursorExt, DecompressType};
9use crate::pak::pak_file::PakCompressionBlock;
10
11// static PAK file magic
12static PAK_MAGIC: u32 = 0x5A6F12E1;
13static BUFFER_SIZE: u32 = 32 * 1024;
14
15pub type PakData = (String, HashMap<String, PakEntry>);
16
17#[derive(Error, Debug)]
18pub enum PakReaderError {
19    #[error("Error occurred while reading: {0}")]
20    ReadError(#[from] std::io::Error),
21    #[error("Error reading header")]
22    ReadHeaderError,
23    #[error("Mismatching magic header")]
24    MagicMismatch,
25    #[error("Unknown version")]
26    UnknownVersion,
27    #[error("Encryption not supported")]
28    EncryptionNotSupported
29}
30
31#[derive(Debug)]
32pub struct PakEntry {
33    start: u64,
34    offset: u64,
35    size: u64,
36    flags: u8,
37    timestamp: u64,
38    hash: Vec<u8>,
39    uncompressed_size: u64,
40    compression_index: u32,
41    compression_block_size: u32,
42    compression_blocks: Vec<PakCompressionBlock>,
43    header_size: u64
44}
45
46#[derive(Debug)]
47pub(crate) struct PakReader {
48    reader: Cursor<Vec<u8>>,
49}
50
51impl PakReader {
52    pub fn new(buffer: Vec<u8>) -> Self {
53        Self {
54            reader: Cursor::new(buffer)
55        }
56    }
57
58    pub async fn get_pak_info(&mut self) -> Result<PakInfo, PakReaderError> {
59        let sizes = PakVersionSizes::get_sizes();
60        for size in sizes {
61            let result = self.read_pak_info(size).await;
62            if let Ok(info) = result {
63                return Ok(info)
64            }
65        }
66
67        Err(PakReaderError::UnknownVersion)
68    }
69
70    pub async fn get_pak_entries(&mut self, info: &PakInfo) -> Result<PakData, PakReaderError> {
71        let index = self.read_pak_index(info).await?;
72        let mut reader = Cursor::new(index);
73
74        let mut entries = HashMap::new();
75        let mut mount_point = reader.read_fstring().await?;
76        let entry_count = reader.read_i32::<LittleEndian>().await?;
77
78        mount_point = mount_point.replace("../../../", "");
79
80        for _ in 0..entry_count {
81            let file_name = reader.read_fstring().await?;
82            let entry = self.read_pak_entry(&mut reader, info).await?;
83
84            entries.insert(file_name, entry);
85        }
86
87        Ok((mount_point, entries))
88    }
89
90    pub async fn get_pak_entry_data(&mut self, entry: &PakEntry) -> Result<Vec<u8>, PakReaderError> {
91        self.reader.set_position(entry.offset + entry.header_size);
92
93        if entry.compression_index == 0 {
94            Ok(self.reader.read_buffer(entry.size as usize).await?)
95        } else {
96            let mut index = 0;
97            let mut offset = 0;
98            let mut decompressed = vec![0u8; entry.uncompressed_size as usize];
99
100            for block in &entry.compression_blocks {
101                let uncompressed_block_size = (entry.uncompressed_size - entry.compression_block_size as u64 * index).min(entry.compression_block_size as u64);
102
103                let compressed_size = block.get_size() as usize;
104                let compressed_buffer = self.reader.read_buffer(uncompressed_block_size as usize).await?;
105                let mut compression_reader = Cursor::new(compressed_buffer);
106                let compression_method = match entry.compression_index {
107                    1 => DecompressType::Zlib,
108                    2 => DecompressType::GZip,
109                    _ => panic!("invalid/unsupported compression index for compressed block")
110                };
111
112                let (bytes_read, decompressed_bytes) = compression_reader.read_decompress(compressed_size, compression_method).await?;
113                decompressed.splice(offset..bytes_read, decompressed_bytes.iter().cloned());
114
115                offset += bytes_read;
116                index += 1;
117            }
118
119            Ok(decompressed)
120        }
121
122    }
123
124    async fn read_pak_entry(&mut self, reader: &mut Cursor<Vec<u8>>, info: &PakInfo) -> Result<PakEntry, PakReaderError> {
125        let start = reader.position();
126        let offset = reader.read_u64::<LittleEndian>().await?;
127        let size = reader.read_u64::<LittleEndian>().await?;
128        let uncompressed_size = reader.read_u64::<LittleEndian>().await?;
129        let mut compression_index: u32 = 0;
130        let mut compression_block_size = 0;
131        let mut flags = 0;
132        let mut timestamp = 0;
133        let mut compression_blocks = vec![];
134
135        if info.version >= PakVersions::FNameBasedCompressionMethod as i32 {
136            if info.sub_version == 1 {
137                compression_index = reader.read_u8().await? as u32;
138            } else {
139                compression_index = reader.read_u32::<LittleEndian>().await?;
140            }
141        } else {
142            let compression_flags = reader.read_u32::<LittleEndian>().await?;
143            if compression_flags == 0 { // No commpression
144                compression_index = 0;
145            } else if compression_flags & 0x01 != 0 { // Zlib compression
146                compression_index = 1;
147            } else if compression_flags & 0x02 != 0 { // GZip Compression
148                compression_index = 2;
149            } else if compression_flags & 0x04 != 0 { // Custom Compression
150                compression_index = 3;
151            }
152        }
153
154        if info.version < PakVersions::NoTimestamps as i32 {
155            timestamp = reader.read_u64::<LittleEndian>().await?;
156        }
157
158        // read hash
159        let hash = reader.read_buffer(20).await?;
160
161        if info.version >= PakVersions::CompressionEncryption as i32 {
162            if compression_index != 0 {
163                let size = reader.read_i32::<LittleEndian>().await?;
164                if size > 0 {
165                    for _ in 0..size {
166                        let compression_start = reader.read_i64::<LittleEndian>().await?;
167                        let compression_end = reader.read_i64::<LittleEndian>().await?;
168                        compression_blocks.push(PakCompressionBlock { compression_start, compression_end });
169                    }
170                }
171            }
172
173            flags = reader.read_u8().await?;
174            compression_block_size = reader.read_u32::<LittleEndian>().await?;
175        }
176
177        let header_size = reader.position() - start;
178
179        Ok(PakEntry {
180            start,
181            offset,
182            size,
183            flags,
184            timestamp,
185            hash,
186            uncompressed_size,
187            compression_index,
188            compression_block_size,
189            compression_blocks,
190            header_size
191        })
192    }
193
194    async fn read_pak_index(&mut self, info: &PakInfo) -> Result<Vec<u8>, PakReaderError> {
195        let position = self.reader.position();
196        self.reader.seek_index(info.index_offset).await;
197        let buffer = self.reader.read_buffer(info.index_size as usize).await?;
198        self.reader.set_position(position);
199        // TODO: decrypt memory
200        Ok(buffer)
201    }
202
203    async fn read_pak_info(&mut self, version_size: PakVersionSizes) -> Result<PakInfo, PakReaderError> {
204        // start reading from the end
205        let version_size = version_size as usize;
206        self.reader.seek(SeekFrom::End(-(version_size as i64)))?;
207
208        // initialize buffer, and reader
209        let mut header = self.reader.read_buffer(version_size).await?;
210        let mut reader = &mut Cursor::new(header);
211
212        // reset the main cursor back to the start
213        self.reader.seek(SeekFrom::Start(0));
214
215        // read the encryption guid
216        let mut encryption_index_guid =  reader.read_buffer(16).await?;
217
218        // is the pak file encrypted?
219        let mut is_encrypted = reader.read_u8().await? > 0;
220        let magic = reader.read_u32::<LittleEndian>().await?;
221
222        // if the magic doesn't match, theres no point continuing
223        if magic != PAK_MAGIC {
224            return Err(PakReaderError::MagicMismatch);
225        }
226
227        let mut compression = vec![];
228        let mut index_frozen = 0u8;
229        let version = reader.read_i32::<LittleEndian>().await?;
230        let index_offset = reader.read_i64::<LittleEndian>().await?;
231        let index_size = reader.read_i64::<LittleEndian>().await?;
232        let index_hash = reader.read_buffer( 20).await?;
233        let sub_version = if version_size == PakVersionSizes::SizeV8A as usize && version == 8 {
234            1
235        } else {
236            0
237        };
238
239        if version < PakVersions::IndexEncryption as i32 {
240            is_encrypted = false;
241        }
242
243        if version < PakVersions::EncryptionKeyGuid as i32 {
244            encryption_index_guid = Vec::default();
245        }
246
247        if version >= PakVersions::FrozenIndex as i32 {
248            index_frozen = reader.read_u8().await?;
249        }
250
251        if version < PakVersions::FNameBasedCompressionMethod as i32 {
252            compression = vec!["Zlib".into(), "Gzip".into(), "Oodle".into()];
253        } else {
254            let mut start = 0;
255            let mut buffer = vec![];
256            let remaining = version_size - reader.position() as usize;
257
258            for idx in 0..remaining {
259                let char = reader.read_u8().await?;
260                if char == 0 {
261                    if buffer.len() > 0 {
262                        compression.push(buffer.iter().collect());
263                        buffer.clear();
264                    }
265
266                    start = idx + 1;
267                } else {
268                    buffer.push(char as char);
269                }
270            }
271        }
272
273        if is_encrypted {
274            return Err(PakReaderError::EncryptionNotSupported)
275        }
276
277        Ok(PakInfo {
278            encryption_index_guid,
279            is_encrypted,
280            magic,
281            version,
282            index_offset,
283            index_size,
284            index_hash,
285            index_frozen,
286            sub_version,
287            compression_methods: compression,
288        })
289    }
290}
291
292impl PakCompressionBlock {
293    pub fn get_size(&self) -> i64 {
294        self.compression_end - self.compression_start
295    }
296}
rust_unreal_unpak/pak/pak_reader.rs

rust_unreal_unpak/pak/
pak_reader.rs
