From c440ca67077316fb5a582c435ec9d5c5c8cabdc0 Mon Sep 17 00:00:00 2001
From: spuds <71292624+bananaturtlesandwich@users.noreply.github.com>
Date: Fri, 13 Jan 2023 12:12:39 +0000
Subject: [PATCH] add encryption support

---
 Cargo.toml          |  1 +
 examples/list.rs    |  1 +
 examples/version.rs |  1 +
 src/error.rs        |  9 +++++++--
 src/pak.rs          | 13 +++++++++++++
 5 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a775799..a25d05a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ edition = "2021"
 [dependencies]
 byteorder = "*"
 strum = { version = "*", features = ["derive"] }
+aes = "*"
 hashbrown = "*"
 thiserror = "*"
 
diff --git a/examples/list.rs b/examples/list.rs
index dfbb30b..8d9dc45 100644
--- a/examples/list.rs
+++ b/examples/list.rs
@@ -2,6 +2,7 @@ fn main() -> Result<(), un_pak::Error> {
     let pak = un_pak::Pak::new(
         std::io::BufReader::new(std::io::Cursor::new(include_bytes!("rando_p.pak"))),
         un_pak::Version::CompressionEncryption,
+        None,
     )?;
     for file in pak.entries.keys() {
         println!("{file}")
diff --git a/examples/version.rs b/examples/version.rs
index b69f457..10635fc 100644
--- a/examples/version.rs
+++ b/examples/version.rs
@@ -5,6 +5,7 @@ fn main() -> Result<(), un_pak::Error> {
         match un_pak::Pak::new(
             std::io::BufReader::new(std::fs::OpenOptions::new().read(true).open(&path)?),
             ver,
+            None,
         ) {
             Ok(un_pak::Pak { version, .. }) | Err(un_pak::Error::Version { version, .. }) => {
                 println!("{}", version);
diff --git a/src/error.rs b/src/error.rs
index 3648c6a..db198a6 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -1,9 +1,13 @@
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
-    #[error("io error: {0}")]
-    Io(#[from] std::io::Error),
+    // dependency errors
     #[error("enum conversion: {0}")]
     Strum(#[from] strum::ParseError),
+    #[error("key hash is an incorrect length")]
+    Aes,
+    // std errors
+    #[error("io error: {0}")]
+    Io(#[from] std::io::Error),
     #[error("utf8 conversion: {0}")]
     Utf8(#[from] std::string::FromUtf8Error),
     #[error("utf16 conversion: {0}")]
@@ -11,6 +15,7 @@ pub enum Error {
     #[error("bufwriter dereference: {0}")]
     IntoInner(#[from] std::io::IntoInnerError<std::io::BufWriter<Vec<u8>>>),
     #[error("found magic of {0:#x} instead of {:#x}", super::MAGIC)]
+    // crate errors
     Magic(u32),
     #[error("used version {used} but pak is version {version}")]
     Version {
diff --git a/src/pak.rs b/src/pak.rs
index 68e26cd..253c128 100644
--- a/src/pak.rs
+++ b/src/pak.rs
@@ -1,6 +1,7 @@
 use std::io;
 
 use super::ReadExt;
+use aes::cipher::{BlockDecrypt, KeyInit};
 use byteorder::{ReadBytesExt, LE};
 
 use super::Version;
@@ -17,10 +18,22 @@ impl Pak {
     pub fn new<R: io::Read + io::Seek>(
         mut reader: R,
         version: super::Version,
+        key: Option<&[u8]>,
     ) -> Result<Self, super::Error> {
         reader.seek(io::SeekFrom::End(-version.size()))?;
         let footer = super::Footer::new(&mut reader, version)?;
         reader.seek(io::SeekFrom::Start(footer.offset))?;
+        let mut index = reader.read_len(footer.size as usize)?;
+        if let Some(key) = key {
+            match aes::Aes256Dec::new_from_slice(key) {
+                Ok(key) => {
+                    for chunk in index.chunks_mut(16) {
+                        key.decrypt_block(aes::Block::from_mut_slice(chunk))
+                    }
+                }
+                Err(_) => return Err(super::Error::Aes),
+            }
+        }
         let mount_point = reader.read_string()?;
         let len = reader.read_u32::<LE>()? as usize;
         let mut entries = hashbrown::HashMap::with_capacity(len);