Implement v8b PakWriter

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"

examples/subcommands/unpack.rs

@@ -8,7 +8,7 @@ pub fn unpack(path: String, key: Option<String>) -> 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}"),
         }

src/entry.rs

@@ -1,7 +1,13 @@
-use super::{ext::ReadExt, Compression, Version, VersionMajor};
+use super::{ext::ReadExt, ext::WriteExt, Compression, Version, VersionMajor};
-use byteorder::{ReadBytesExt, LE};
+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<R: io::Read>(reader: &mut R) -> Result<Self, super::Error> {
+    pub fn read<R: io::Read>(reader: &mut R) -> Result<Self, super::Error> {
         Ok(Self {
             start: reader.read_u64::<LE>()?,
             end: reader.read_u64::<LE>()?,
         })
     }
+
+    pub fn write<W: io::Write>(&self, writer: &mut W) -> Result<(), super::Error> {
+        writer.write_u64::<LE>(self.start)?;
+        writer.write_u64::<LE>(self.end)?;
+        Ok(())
+    }
 }
 
 fn align(offset: u64) -> u64 {
@@ -66,7 +78,10 @@ impl Entry {
         size
     }
 
-    pub fn new<R: io::Read>(reader: &mut R, version: super::Version) -> Result<Self, super::Error> {
+    pub fn read<R: io::Read>(
+        reader: &mut R,
+        version: super::Version,
+    ) -> Result<Self, super::Error> {
         // since i need the compression flags, i have to store these as variables which is mildly annoying
         let offset = reader.read_u64::<LE>()?;
         let compressed = reader.read_u64::<LE>()?;
@@ -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<W: io::Write>(
+        &self,
+        writer: &mut W,
+        version: super::Version,
+        location: EntryLocation,
+    ) -> Result<(), super::Error> {
+        writer.write_u64::<LE>(match location {
+            EntryLocation::Data => 0,
+            EntryLocation::Index => self.offset,
+        })?;
+        writer.write_u64::<LE>(self.compressed)?;
+        writer.write_u64::<LE>(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::<LE>(compression.into())?,
+        }
 
-    pub fn new_encoded<R: io::Read>(
+        if version.version_major() == VersionMajor::Initial {
+            writer.write_u64::<LE>(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::<LE>(self.block_uncompressed.unwrap_or_default())?;
+        }
+        Ok(())
+    }
+
+    pub fn read_encoded<R: io::Read>(
         reader: &mut R,
         version: super::Version,
     ) -> Result<Self, super::Error> {
@@ -194,7 +251,7 @@ impl Entry {
         })
     }
 
-    pub fn read<R: io::Read + io::Seek, W: io::Write>(
+    pub fn read_file<R: io::Read + io::Seek, W: io::Write>(
         &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);
+    }
+}

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<bool, super::Error>;
@@ -11,6 +11,11 @@ pub trait ReadExt {
     fn read_len(&mut self, len: usize) -> Result<Vec<u8>, 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<R: std::io::Read> ReadExt for R {
     fn read_bool(&mut self) -> Result<bool, super::Error> {
         match self.read_u8()? {
@@ -37,7 +42,7 @@ impl<R: std::io::Read> ReadExt for R {
         Ok(buf)
     }
 
-    fn read_string(&mut self) -> Result<String, crate::Error> {
+    fn read_string(&mut self) -> Result<String, super::Error> {
         let mut buf = match self.read_i32::<LE>()? {
             size if size.is_negative() => {
                 let mut buf = Vec::with_capacity(-size as usize);
@@ -59,3 +64,20 @@ impl<R: std::io::Read> ReadExt for R {
         Ok(buf)
     }
 }
+
+impl<W: std::io::Write> 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::<LE>(bytes.len() as u32 + 1)?;
+        self.write_all(bytes)?;
+        self.write_u8(0)?;
+        Ok(())
+    }
+}

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<R: std::io::Read>(reader: &mut R, version: Version) -> Result<Self, super::Error> {
+    pub fn read<R: std::io::Read>(reader: &mut R, version: Version) -> Result<Self, super::Error> {
         let footer = Self {
             encryption_uuid: match version.version_major() >= VersionMajor::EncryptionKeyGuid {
                 true => Some(reader.read_u128::<LE>()?),
@@ -66,4 +68,27 @@ impl Footer {
         }
         Ok(footer)
     }
+
+    pub fn write<W: std::io::Write>(&self, writer: &mut W) -> Result<(), super::Error> {
+        if self.version_major >= VersionMajor::EncryptionKeyGuid {
+            writer.write_u128::<LE>(0)?;
+        }
+        if self.version_major >= VersionMajor::IndexEncryption {
+            writer.write_bool(self.encrypted)?;
+        }
+        writer.write_u32::<LE>(self.magic)?;
+        writer.write_u32::<LE>(self.version_major as u32)?;
+        writer.write_u64::<LE>(self.index_offset)?;
+        writer.write_u64::<LE>(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(())
+    }
 }

src/pak.rs

@@ -1,20 +1,38 @@
+use super::ext::{ReadExt, WriteExt};
 use super::{Version, VersionMajor};
-use hashbrown::HashMap;
+use byteorder::{ReadBytesExt, WriteBytesExt, LE};
-use std::io::{self, Seek};
+use std::collections::BTreeMap;
+use std::io::{self, Read, Seek, Write};
 
 #[derive(Debug)]
-pub struct PakReader<R: io::Read + io::Seek> {
+pub struct PakReader<R: Read + Seek> {
     pak: Pak,
     reader: R,
+    key: Option<aes::Aes256Dec>,
+}
+#[derive(Debug)]
+pub struct PakWriter<W: Write + Seek> {
+    pak: Pak,
+    writer: W,
+    key: Option<aes::Aes256Enc>,
 }
 #[derive(Debug)]
 pub struct Pak {
     version: Version,
     mount_point: String,
-    key: Option<aes::Aes256Dec>,
     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<String, super::entry::Entry> {
+    fn new(version: Version) -> Self {
+        if version < Version::V10 {
+            Self::V1(IndexV1::default())
+        } else {
+            Self::V2(IndexV2::default())
+        }
+    }
+
+    fn entries(&self) -> &BTreeMap<String, super::entry::Entry> {
         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<String, super::entry::Entry>,
+    entries: BTreeMap<String, super::entry::Entry>,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct IndexV2 {
     path_hash_seed: u64,
     path_hash_index: Option<Vec<u8>>,
-    full_directory_index: Option<HashMap<String, HashMap<String, u32>>>,
+    full_directory_index: Option<BTreeMap<String, BTreeMap<String, u32>>>,
     encoded_entries: Vec<u8>,
-    entries_by_path: HashMap<String, super::entry::Entry>,
+    entries_by_path: BTreeMap<String, super::entry::Entry>,
 }
 
 fn decrypt(key: &Option<aes::Aes256Dec>, bytes: &mut [u8]) -> Result<(), super::Error> {
@@ -56,18 +89,12 @@ fn decrypt(key: &Option<aes::Aes256Dec>, bytes: &mut [u8]) -> Result<(), super::
     }
 }
 
-impl<R: io::Seek + io::Read> PakReader<R> {
+impl<R: Read + Seek> PakReader<R> {
-    pub fn into_reader(self) -> R {
-        self.reader
-    }
-}
-
-impl<R: io::Read + io::Seek> PakReader<R> {
     pub fn new_any(mut reader: R, key: Option<aes::Aes256Dec>) -> Result<Self, super::Error> {
         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<R: io::Read + io::Seek> PakReader<R> {
         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<Vec<u8>, super::Error> {
+        let mut data = Vec::new();
+        self.read_file(path, &mut data)?;
+        Ok(data)
+    }
+
+    pub fn read_file<W: io::Write>(
+        &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<String> {
+        self.pak
+            .index
+            .entries()
+            .keys()
+            .cloned()
+            .collect::<Vec<String>>()
+            .into_iter()
+    }
+}
+
+impl<W: Write + io::Seek> PakWriter<W> {
     pub fn new(
+        writer: W,
+        key: Option<aes::Aes256Enc>,
+        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<R: Read>(&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<W, super::Error> {
+        self.pak.write(&mut self.writer, self.key)?;
+        Ok(self.writer)
+    }
+}
+
+impl Pak {
+    fn read<R: Read + Seek>(
         mut reader: R,
         version: super::Version,
         key: Option<aes::Aes256Dec>,
-    ) -> Result<Pak, super::Error> {
+    ) -> Result<Self, super::Error> {
-        use super::ext::ReadExt;
-        use byteorder::{ReadBytesExt, LE};
         // 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<R: io::Read + io::Seek> PakReader<R> {
                 let mut fdi = io::Cursor::new(full_directory_index);
 
                 let dir_count = fdi.read_u32::<LE>()? 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::<LE>()? 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::<LE>()?);
@@ -152,14 +282,14 @@ impl<R: io::Read + io::Seek> PakReader<R> {
             let size = index.read_u32::<LE>()? 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<R: io::Read + io::Seek> PakReader<R> {
                 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<R: io::Read + io::Seek> PakReader<R> {
         Ok(Pak {
             version,
             mount_point,
-            key,
             index,
         })
     }
+    fn write<W: Write + Seek>(
+        &self,
+        writer: &mut W,
+        _key: Option<aes::Aes256Enc>,
+    ) -> Result<(), super::Error> {
+        let index_offset = writer.stream_position()?;
 
-    pub fn version(&self) -> super::Version {
+        let mut index_cur = std::io::Cursor::new(vec![]);
-        self.pak.version
+        index_cur.write_string(&self.mount_point)?;
-    }
 
-    pub fn mount_point(&self) -> &str {
+        match &self.index {
-        &self.pak.mount_point
+            Index::V1(index) => {
-    }
+                index_cur.write_u32::<LE>(index.entries.len() as u32)?;
-
+                for (path, entry) in &index.entries {
-    pub fn get(&mut self, path: &str) -> Result<Vec<u8>, super::Error> {
+                    index_cur.write_string(path)?;
-        let mut data = Vec::new();
+                    entry.write(
-        self.read(path, &mut data)?;
+                        &mut index_cur,
-        Ok(data)
+                        self.version,
-    }
+                        super::entry::EntryLocation::Index,
-
+                    )?;
-    pub fn read<W: io::Write>(&mut self, path: &str, writer: &mut W) -> Result<(), super::Error> {
+                }
-        match self.pak.index.entries().get(path) {
+            }
-            Some(entry) => entry.read(
+            Index::V2(_index) => todo!(),
-                &mut self.reader,
-                self.pak.version,
-                self.pak.key.as_ref(),
-                writer,
-            ),
-            None => Err(super::Error::Other("no file found at given path")),
         }
-    }
 
-    pub fn files(&self) -> std::vec::IntoIter<String> {
+        let index_data = index_cur.into_inner();
-        self.pak
+
-            .index
+        use sha1::{Digest, Sha1};
-            .entries()
+        let mut hasher = Sha1::new();
-            .keys()
+        hasher.update(&index_data);
-            .cloned()
+
-            .collect::<Vec<String>>()
+        let footer = super::footer::Footer {
-            .into_iter()
+            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);
     }
 }

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"),