From ed653969a7d75543f6984ac60b5934a2b263d2f9 Mon Sep 17 00:00:00 2001 From: Truman Kilen Date: Sat, 21 Jan 2023 22:40:53 -0600 Subject: [PATCH] Implement v8b PakWriter --- Cargo.toml | 1 + examples/subcommands/unpack.rs | 2 +- src/entry.rs | 92 ++++++++++- src/ext.rs | 26 ++- src/footer.rs | 29 +++- src/pak.rs | 292 ++++++++++++++++++++++++++------- tests/test.rs | 2 +- 7 files changed, 369 insertions(+), 75 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index d156923..2ccd325 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ aes = "0.8" flate2 = "1.0" hashbrown = "0.13" thiserror = "1.0" +sha1 = "0.10.5" [dev-dependencies] base64 = "0.21.0" diff --git a/examples/subcommands/unpack.rs b/examples/subcommands/unpack.rs index 715f6e9..8bf268e 100644 --- a/examples/subcommands/unpack.rs +++ b/examples/subcommands/unpack.rs @@ -8,7 +8,7 @@ pub fn unpack(path: String, key: Option) -> Result<(), unpak::Error> { let mut pak = super::load_pak(path.clone(), key)?; for file in pak.files() { std::fs::create_dir_all(folder.join(&file).parent().expect("will be a file"))?; - match pak.read(&file, &mut std::fs::File::create(folder.join(&file))?) { + match pak.read_file(&file, &mut std::fs::File::create(folder.join(&file))?) { Ok(_) => println!("{file}"), Err(e) => eprintln!("{e}"), } diff --git a/src/entry.rs b/src/entry.rs index 966c193..8ba1070 100644 --- a/src/entry.rs +++ b/src/entry.rs @@ -1,7 +1,13 @@ -use super::{ext::ReadExt, Compression, Version, VersionMajor}; -use byteorder::{ReadBytesExt, LE}; +use super::{ext::ReadExt, ext::WriteExt, Compression, Version, VersionMajor}; +use byteorder::{ReadBytesExt, WriteBytesExt, LE}; use std::io; +#[derive(Debug)] +pub enum EntryLocation { + Data, + Index, +} + #[derive(Debug)] pub struct Block { pub start: u64, @@ -9,12 +15,18 @@ pub struct Block { } impl Block { - pub fn new(reader: &mut R) -> Result { + pub fn read(reader: &mut R) -> Result { Ok(Self { start: reader.read_u64::()?, end: reader.read_u64::()?, }) } + + pub fn write(&self, writer: &mut W) -> Result<(), super::Error> { + writer.write_u64::(self.start)?; + writer.write_u64::(self.end)?; + Ok(()) + } } fn align(offset: u64) -> u64 { @@ -66,7 +78,10 @@ impl Entry { size } - pub fn new(reader: &mut R, version: super::Version) -> Result { + pub fn read( + reader: &mut R, + version: super::Version, + ) -> Result { // since i need the compression flags, i have to store these as variables which is mildly annoying let offset = reader.read_u64::()?; let compressed = reader.read_u64::()?; @@ -92,7 +107,7 @@ impl Entry { blocks: match version.version_major() >= VersionMajor::CompressionEncryption && compression != Compression::None { - true => Some(reader.read_array(Block::new)?), + true => Some(reader.read_array(Block::read)?), false => None, }, encrypted: version.version_major() >= VersionMajor::CompressionEncryption @@ -104,8 +119,50 @@ impl Entry { }, }) } + pub fn write( + &self, + writer: &mut W, + version: super::Version, + location: EntryLocation, + ) -> Result<(), super::Error> { + writer.write_u64::(match location { + EntryLocation::Data => 0, + EntryLocation::Index => self.offset, + })?; + writer.write_u64::(self.compressed)?; + writer.write_u64::(self.uncompressed)?; + let compression: u8 = match self.compression { + Compression::None => 0, + Compression::Zlib => 1, + Compression::Gzip => todo!(), + Compression::Oodle => todo!(), + }; + match version { + Version::V8A => writer.write_u8(compression)?, + _ => writer.write_u32::(compression.into())?, + } - pub fn new_encoded( + if version.version_major() == VersionMajor::Initial { + writer.write_u64::(self.timestamp.unwrap_or_default())?; + } + if let Some(hash) = self.hash { + writer.write_all(&hash)?; + } else { + panic!("hash missing"); + } + if version.version_major() >= VersionMajor::CompressionEncryption { + if let Some(blocks) = &self.blocks { + for block in blocks { + block.write(writer)?; + } + } + writer.write_bool(self.encrypted)?; + writer.write_u32::(self.block_uncompressed.unwrap_or_default())?; + } + Ok(()) + } + + pub fn read_encoded( reader: &mut R, version: super::Version, ) -> Result { @@ -194,7 +251,7 @@ impl Entry { }) } - pub fn read( + pub fn read_file( &self, reader: &mut R, version: Version, @@ -202,7 +259,7 @@ impl Entry { buf: &mut W, ) -> Result<(), super::Error> { reader.seek(io::SeekFrom::Start(self.offset))?; - Entry::new(reader, version)?; + Entry::read(reader, version)?; let data_offset = reader.stream_position()?; let mut data = reader.read_len(match self.encrypted { true => align(self.compressed), @@ -258,3 +315,22 @@ impl Entry { Ok(()) } } + +mod test { + #[test] + fn test_entry() { + let data = vec![ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x54, 0x02, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x54, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0xDD, 0x94, 0xFD, 0xC3, 0x5F, 0xF5, 0x91, 0xA9, 0x9A, 0x5E, 0x14, 0xDC, 0x9B, + 0xD3, 0x58, 0x89, 0x78, 0xA6, 0x1C, 0x00, 0x00, 0x00, 0x00, 0x00, + ]; + let mut out = vec![]; + let entry = super::Entry::read(&mut std::io::Cursor::new(data.clone()), super::Version::V5) + .unwrap(); + entry + .write(&mut out, super::Version::V5, super::EntryLocation::Data) + .unwrap(); + assert_eq!(&data, &out); + } +} diff --git a/src/ext.rs b/src/ext.rs index 90e8835..5dd196c 100644 --- a/src/ext.rs +++ b/src/ext.rs @@ -1,4 +1,4 @@ -use byteorder::{ReadBytesExt, LE}; +use byteorder::{ReadBytesExt, WriteBytesExt, LE}; pub trait ReadExt { fn read_bool(&mut self) -> Result; @@ -11,6 +11,11 @@ pub trait ReadExt { fn read_len(&mut self, len: usize) -> Result, super::Error>; } +pub trait WriteExt { + fn write_bool(&mut self, value: bool) -> Result<(), super::Error>; + fn write_string(&mut self, value: &str) -> Result<(), super::Error>; +} + impl ReadExt for R { fn read_bool(&mut self) -> Result { match self.read_u8()? { @@ -37,7 +42,7 @@ impl ReadExt for R { Ok(buf) } - fn read_string(&mut self) -> Result { + fn read_string(&mut self) -> Result { let mut buf = match self.read_i32::()? { size if size.is_negative() => { let mut buf = Vec::with_capacity(-size as usize); @@ -59,3 +64,20 @@ impl ReadExt for R { Ok(buf) } } + +impl WriteExt for W { + fn write_bool(&mut self, value: bool) -> Result<(), super::Error> { + self.write_u8(match value { + true => 1, + false => 0, + })?; + Ok(()) + } + fn write_string(&mut self, value: &str) -> Result<(), super::Error> { + let bytes = value.as_bytes(); + self.write_u32::(bytes.len() as u32 + 1)?; + self.write_all(bytes)?; + self.write_u8(0)?; + Ok(()) + } +} diff --git a/src/footer.rs b/src/footer.rs index e30d3c1..2539760 100644 --- a/src/footer.rs +++ b/src/footer.rs @@ -1,5 +1,7 @@ +use crate::ext::WriteExt; + use super::{ext::ReadExt, Compression, Version, VersionMajor}; -use byteorder::{ReadBytesExt, LE}; +use byteorder::{ReadBytesExt, WriteBytesExt, LE}; use std::str::FromStr; #[derive(Debug)] @@ -17,7 +19,7 @@ pub struct Footer { } impl Footer { - pub fn new(reader: &mut R, version: Version) -> Result { + pub fn read(reader: &mut R, version: Version) -> Result { let footer = Self { encryption_uuid: match version.version_major() >= VersionMajor::EncryptionKeyGuid { true => Some(reader.read_u128::()?), @@ -66,4 +68,27 @@ impl Footer { } Ok(footer) } + + pub fn write(&self, writer: &mut W) -> Result<(), super::Error> { + if self.version_major >= VersionMajor::EncryptionKeyGuid { + writer.write_u128::(0)?; + } + if self.version_major >= VersionMajor::IndexEncryption { + writer.write_bool(self.encrypted)?; + } + writer.write_u32::(self.magic)?; + writer.write_u32::(self.version_major as u32)?; + writer.write_u64::(self.index_offset)?; + writer.write_u64::(self.index_size)?; + writer.write_all(&self.hash)?; + let algo_size = match self.version { + ver if ver < Version::V8A => 0, + ver if ver < Version::V8B => 4, + _ => 5, + }; + for _ in 0..algo_size { + writer.write_all(&[0; 32])?; + } + Ok(()) + } } diff --git a/src/pak.rs b/src/pak.rs index a3756cb..0893ca9 100644 --- a/src/pak.rs +++ b/src/pak.rs @@ -1,20 +1,38 @@ +use super::ext::{ReadExt, WriteExt}; use super::{Version, VersionMajor}; -use hashbrown::HashMap; -use std::io::{self, Seek}; +use byteorder::{ReadBytesExt, WriteBytesExt, LE}; +use std::collections::BTreeMap; +use std::io::{self, Read, Seek, Write}; #[derive(Debug)] -pub struct PakReader { +pub struct PakReader { pak: Pak, reader: R, + key: Option, +} +#[derive(Debug)] +pub struct PakWriter { + pak: Pak, + writer: W, + key: Option, } #[derive(Debug)] pub struct Pak { version: Version, mount_point: String, - key: Option, index: Index, } +impl Pak { + fn new(version: Version, mount_point: String) -> Self { + Pak { + version, + mount_point, + index: Index::new(version), + } + } +} + #[derive(Debug)] pub enum Index { V1(IndexV1), @@ -22,26 +40,41 @@ pub enum Index { } impl Index { - fn entries(&self) -> &HashMap { + fn new(version: Version) -> Self { + if version < Version::V10 { + Self::V1(IndexV1::default()) + } else { + Self::V2(IndexV2::default()) + } + } + + fn entries(&self) -> &BTreeMap { match self { Index::V1(index) => &index.entries, Index::V2(index) => &index.entries_by_path, } } + + fn add_entry(&mut self, path: &str, entry: super::entry::Entry) { + match self { + Index::V1(index) => index.entries.insert(path.to_string(), entry), + Index::V2(_index) => todo!(), + }; + } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct IndexV1 { - entries: HashMap, + entries: BTreeMap, } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct IndexV2 { path_hash_seed: u64, path_hash_index: Option>, - full_directory_index: Option>>, + full_directory_index: Option>>, encoded_entries: Vec, - entries_by_path: HashMap, + entries_by_path: BTreeMap, } fn decrypt(key: &Option, bytes: &mut [u8]) -> Result<(), super::Error> { @@ -56,18 +89,12 @@ fn decrypt(key: &Option, bytes: &mut [u8]) -> Result<(), super:: } } -impl PakReader { - pub fn into_reader(self) -> R { - self.reader - } -} - -impl PakReader { +impl PakReader { pub fn new_any(mut reader: R, key: Option) -> Result { for ver in Version::iter() { - match PakReader::new(&mut reader, ver, key.clone()) { + match Pak::read(&mut reader, ver, key.clone()) { Ok(pak) => { - return Ok(PakReader { pak, reader }); + return Ok(PakReader { pak, reader, key }); } _ => continue, } @@ -75,16 +102,119 @@ impl PakReader { Err(super::Error::Other("version unsupported")) } + pub fn into_reader(self) -> R { + self.reader + } + + pub fn version(&self) -> super::Version { + self.pak.version + } + + pub fn mount_point(&self) -> &str { + &self.pak.mount_point + } + + pub fn get(&mut self, path: &str) -> Result, super::Error> { + let mut data = Vec::new(); + self.read_file(path, &mut data)?; + Ok(data) + } + + pub fn read_file( + &mut self, + path: &str, + writer: &mut W, + ) -> Result<(), super::Error> { + match self.pak.index.entries().get(path) { + Some(entry) => entry.read_file( + &mut self.reader, + self.pak.version, + self.key.as_ref(), + writer, + ), + None => Err(super::Error::Other("no file found at given path")), + } + } + + pub fn files(&self) -> std::vec::IntoIter { + self.pak + .index + .entries() + .keys() + .cloned() + .collect::>() + .into_iter() + } +} + +impl PakWriter { pub fn new( + writer: W, + key: Option, + version: Version, + mount_point: String, + ) -> Self { + PakWriter { + pak: Pak::new(version, mount_point), + writer, + key, + } + } + + pub fn into_writer(self) -> W { + self.writer + } + + pub fn write_file(&mut self, path: &str, reader: &mut R) -> Result<(), super::Error> { + let mut data = vec![]; + reader.read_to_end(&mut data)?; + + use sha1::{Digest, Sha1}; + let mut hasher = Sha1::new(); + hasher.update(&data); + + let offset = self.writer.stream_position()?; + let len = data.len() as u64; + + let entry = super::entry::Entry { + offset, + compressed: len, + uncompressed: len, + compression: super::Compression::None, + timestamp: None, + hash: Some(hasher.finalize().into()), + blocks: None, + encrypted: false, + block_uncompressed: None, + }; + + entry.write( + &mut self.writer, + self.pak.version, + super::entry::EntryLocation::Data, + )?; + + self.pak.index.add_entry(path, entry); + + self.writer.write_all(&data)?; + Ok(()) + } + + pub fn write_index(mut self) -> Result { + self.pak.write(&mut self.writer, self.key)?; + Ok(self.writer) + } +} + +impl Pak { + fn read( mut reader: R, version: super::Version, key: Option, - ) -> Result { - use super::ext::ReadExt; - use byteorder::{ReadBytesExt, LE}; + ) -> Result { // read footer to get index, encryption & compression info reader.seek(io::SeekFrom::End(-version.size()))?; - let footer = super::footer::Footer::new(&mut reader, version)?; + let footer = super::footer::Footer::read(&mut reader, version)?; // read index to get all the entry info reader.seek(io::SeekFrom::Start(footer.index_offset))?; let mut index = reader.read_len(footer.index_size as usize)?; @@ -134,11 +264,11 @@ impl PakReader { let mut fdi = io::Cursor::new(full_directory_index); let dir_count = fdi.read_u32::()? as usize; - let mut directories = HashMap::with_capacity(dir_count); + let mut directories = BTreeMap::new(); for _ in 0..dir_count { let dir_name = fdi.read_string()?; let file_count = fdi.read_u32::()? as usize; - let mut files = HashMap::with_capacity(file_count); + let mut files = BTreeMap::new(); for _ in 0..file_count { let file_name = fdi.read_string()?; files.insert(file_name, fdi.read_u32::()?); @@ -152,14 +282,14 @@ impl PakReader { let size = index.read_u32::()? as usize; let encoded_entries = index.read_len(size)?; - let mut entries_by_path = HashMap::new(); + let mut entries_by_path = BTreeMap::new(); if let Some(fdi) = &full_directory_index { let mut encoded_entries = io::Cursor::new(&encoded_entries); for (dir_name, dir) in fdi { for (file_name, encoded_offset) in dir { encoded_entries.seek(io::SeekFrom::Start(*encoded_offset as u64))?; let entry = - super::entry::Entry::new_encoded(&mut encoded_entries, version)?; + super::entry::Entry::read_encoded(&mut encoded_entries, version)?; // entry next to file contains full metadata //reader.seek(io::SeekFrom::Start(entry.offset))?; @@ -186,11 +316,11 @@ impl PakReader { entries_by_path, }) } else { - let mut entries = HashMap::with_capacity(len); + let mut entries = BTreeMap::new(); for _ in 0..len { entries.insert( index.read_string()?, - super::entry::Entry::new(&mut index, version)?, + super::entry::Entry::read(&mut index, version)?, ); } Index::V1(IndexV1 { entries }) @@ -199,44 +329,84 @@ impl PakReader { Ok(Pak { version, mount_point, - key, index, }) } + fn write( + &self, + writer: &mut W, + _key: Option, + ) -> Result<(), super::Error> { + let index_offset = writer.stream_position()?; - pub fn version(&self) -> super::Version { - self.pak.version - } + let mut index_cur = std::io::Cursor::new(vec![]); + index_cur.write_string(&self.mount_point)?; - pub fn mount_point(&self) -> &str { - &self.pak.mount_point - } - - pub fn get(&mut self, path: &str) -> Result, super::Error> { - let mut data = Vec::new(); - self.read(path, &mut data)?; - Ok(data) - } - - pub fn read(&mut self, path: &str, writer: &mut W) -> Result<(), super::Error> { - match self.pak.index.entries().get(path) { - Some(entry) => entry.read( - &mut self.reader, - self.pak.version, - self.pak.key.as_ref(), - writer, - ), - None => Err(super::Error::Other("no file found at given path")), + match &self.index { + Index::V1(index) => { + index_cur.write_u32::(index.entries.len() as u32)?; + for (path, entry) in &index.entries { + index_cur.write_string(path)?; + entry.write( + &mut index_cur, + self.version, + super::entry::EntryLocation::Index, + )?; + } + } + Index::V2(_index) => todo!(), } - } - pub fn files(&self) -> std::vec::IntoIter { - self.pak - .index - .entries() - .keys() - .cloned() - .collect::>() - .into_iter() + let index_data = index_cur.into_inner(); + + use sha1::{Digest, Sha1}; + let mut hasher = Sha1::new(); + hasher.update(&index_data); + + let footer = super::footer::Footer { + encryption_uuid: None, + encrypted: false, + magic: super::MAGIC, + version: self.version, + version_major: self.version.version_major(), + index_offset, + index_size: index_data.len() as u64, + hash: hasher.finalize().into(), + frozen: false, + compression: vec![], + }; + + writer.write_all(&index_data)?; + + footer.write(writer)?; + + Ok(()) + } +} + +mod test { + #[test] + fn test_rewrite_pak() { + use std::io::Cursor; + let bytes = include_bytes!("../tests/packs/pack_v8b.pak"); + + let mut reader = super::PakReader::new_any(Cursor::new(bytes), None).unwrap(); + let writer = Cursor::new(vec![]); + let mut pak_writer = super::PakWriter::new( + writer, + None, + super::Version::V8B, + reader.mount_point().to_owned(), + ); + + for path in reader.files() { + let data = reader.get(&path).unwrap(); + pak_writer + .write_file(&path, &mut std::io::Cursor::new(data)) + .unwrap(); + } + + let out_bytes = pak_writer.write_index().unwrap().into_inner(); + assert_eq!(bytes.to_vec(), out_bytes); } } diff --git a/tests/test.rs b/tests/test.rs index 0e9bdd2..cf623da 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -134,7 +134,7 @@ macro_rules! encryptindex { for file in files { let mut buf = vec![]; let mut writer = std::io::Cursor::new(&mut buf); - pak.read(&file, &mut writer).unwrap(); + pak.read_file(&file, &mut writer).unwrap(); match file.as_str() { "test.txt" => assert_eq!(buf, include_bytes!("pack/root/test.txt"), "test.txt incorrect contents"), "test.png" => assert_eq!(buf, include_bytes!("pack/root/test.png"), "test.png incorrect contents"),