Improve AES key handling and accept both hex string and base64 formats

repak/src/error.rs

@@ -6,12 +6,9 @@ pub enum Error {
     #[error("enum conversion: {0}")]
     Strum(#[from] strum::ParseError),
 
-    #[error("key hash is an incorrect length")]
+    #[error("expect 256 bit AES key as base64 or hex string")]
     Aes,
 
-    #[error("malformed base64")]
-    Base64,
-
     // std errors
     #[error("io error: {0}")]
     Io(#[from] std::io::Error),
@@ -68,8 +65,8 @@ pub enum Error {
     #[error("error with OsString")]
     OsString(std::ffi::OsString),
 
-    #[error("{0}")]
+    #[error("version unsupported or is encrypted (possibly missing --aes-key?)")]
-    Other(String),
+    UnsuportedOrEncrypted,
 }
 
 impl std::fmt::Debug for Error {

repak/src/pak.rs

@@ -91,7 +91,7 @@ impl PakReader {
                 _ => continue,
             }
         }
-        Err(super::Error::Other("version unsupported".to_owned()))
+        Err(super::Error::UnsuportedOrEncrypted)
     }
 
     pub fn version(&self) -> super::Version {

repak/tests/test.rs

@@ -88,7 +88,7 @@ fn test_read(version: repak::Version, _file_name: &str, bytes: &[u8]) {
     let key = general_purpose::STANDARD
         .decode(AES_KEY)
         .as_ref()
-        .map_err(|_| repak::Error::Base64)
+        .map_err(|_| repak::Error::Aes)
         .and_then(|bytes| aes::Aes256::new_from_slice(bytes).map_err(|_| repak::Error::Aes))
         .unwrap();
 
@@ -154,7 +154,7 @@ fn test_write(_version: repak::Version, _file_name: &str, bytes: &[u8]) {
     let key = general_purpose::STANDARD
         .decode(AES_KEY)
         .as_ref()
-        .map_err(|_| repak::Error::Base64)
+        .map_err(|_| repak::Error::Aes)
         .and_then(|bytes| aes::Aes256::new_from_slice(bytes).map_err(|_| repak::Error::Aes))
         .unwrap();
 
@@ -186,7 +186,7 @@ fn test_rewrite_index(_version: repak::Version, _file_name: &str, bytes: &[u8])
     let key = general_purpose::STANDARD
         .decode(AES_KEY)
         .as_ref()
-        .map_err(|_| repak::Error::Base64)
+        .map_err(|_| repak::Error::Aes)
         .and_then(|bytes| aes::Aes256::new_from_slice(bytes).map_err(|_| repak::Error::Aes))
         .unwrap();
 

repak_cli/Cargo.toml

@@ -14,6 +14,7 @@ path = "src/main.rs"
 aes = { workspace = true }
 base64 = { workspace = true }
 clap = { version = "4.1.4", features = ["derive"] }
+hex = "0.4.3"
 indicatif = { version = "0.17.3", features = ["rayon"] }
 path-clean = "0.1.0"
 path-slash = "0.2.1"

repak_cli/src/main.rs

@@ -14,10 +14,6 @@ struct ActionInfo {
     /// Input .pak path
     #[arg(index = 1)]
     input: String,
-
-    /// Base64 encoded AES encryption key if the pak is encrypted
-    #[arg(short, long)]
-    aes_key: Option<String>,
 }
 
 #[derive(Parser, Debug)]
@@ -25,10 +21,6 @@ struct ActionList {
     /// Input .pak path
     #[arg(index = 1)]
     input: String,
-
-    /// Base64 encoded AES encryption key if the pak is encrypted
-    #[arg(short, long)]
-    aes_key: Option<String>,
 }
 
 #[derive(Parser, Debug)]
@@ -45,10 +37,6 @@ struct ActionUnpack {
     #[arg(short, long, default_value = "../../../")]
     strip_prefix: String,
 
-    /// Base64 encoded AES encryption key if the pak is encrypted
-    #[arg(short, long)]
-    aes_key: Option<String>,
-
     /// Verbose
     #[arg(short, long, default_value = "false")]
     verbose: bool,
@@ -106,10 +94,6 @@ struct ActionGet {
     /// Prefix to strip from entry path
     #[arg(short, long, default_value = "../../../")]
     strip_prefix: String,
-
-    /// Base64 encoded AES encryption key if the pak is encrypted
-    #[arg(short, long)]
-    aes_key: Option<String>,
 }
 
 #[derive(Subcommand, Debug)]
@@ -129,37 +113,50 @@ enum Action {
 #[derive(Parser, Debug)]
 #[command(author, version)]
 struct Args {
+    /// 256 bit AES encryption key as base64 or hex string if the pak is encrypted
+    #[arg(short, long)]
+    aes_key: Option<AesKey>,
+
     #[command(subcommand)]
     action: Action,
 }
 
-fn main() -> Result<(), repak::Error> {
+#[derive(Debug, Clone)]
-    let args = Args::parse();
+struct AesKey(aes::Aes256);
-
+impl std::str::FromStr for AesKey {
-    match args.action {
+    type Err = repak::Error;
-        Action::Info(args) => info(args),
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
-        Action::List(args) => list(args),
+        use aes::cipher::KeyInit;
-        Action::Unpack(args) => unpack(args),
+        use base64::{engine::general_purpose, Engine as _};
-        Action::Pack(args) => pack(args),
+        let try_parse = |bytes: Vec<_>| aes::Aes256::new_from_slice(&bytes).ok().map(AesKey);
-        Action::Get(args) => get(args),
+        hex::decode(s.strip_prefix("0x").unwrap_or(s))
+            .ok()
+            .and_then(try_parse)
+            .or_else(|| {
+                general_purpose::STANDARD_NO_PAD
+                    .decode(s.trim_end_matches('='))
+                    .ok()
+                    .and_then(try_parse)
+            })
+            .ok_or(repak::Error::Aes)
     }
 }
 
-fn aes_key(key: &str) -> Result<aes::Aes256, repak::Error> {
+fn main() -> Result<(), repak::Error> {
-    use aes::cipher::KeyInit;
+    let args = Args::parse();
-    use base64::{engine::general_purpose, Engine as _};
+    let aes_key = args.aes_key.map(|k| k.0);
-    general_purpose::STANDARD
+
-        .decode(key)
+    match args.action {
-        .as_ref()
+        Action::Info(action) => info(aes_key, action),
-        .map_err(|_| repak::Error::Base64)
+        Action::List(action) => list(aes_key, action),
-        .and_then(|bytes| aes::Aes256::new_from_slice(bytes).map_err(|_| repak::Error::Aes))
+        Action::Unpack(action) => unpack(aes_key, action),
+        Action::Pack(action) => pack(action),
+        Action::Get(action) => get(aes_key, action),
+    }
 }
 
-fn info(args: ActionInfo) -> Result<(), repak::Error> {
+fn info(aes_key: Option<aes::Aes256>, action: ActionInfo) -> Result<(), repak::Error> {
-    let pak = repak::PakReader::new_any(
+    let pak = repak::PakReader::new_any(BufReader::new(File::open(action.input)?), aes_key)?;
-        BufReader::new(File::open(&args.input)?),
-        args.aes_key.map(|k| aes_key(k.as_str())).transpose()?,
-    )?;
     println!("mount point: {}", pak.mount_point());
     println!("version: {}", pak.version());
     println!("version major: {}", pak.version().version_major());
@@ -167,11 +164,8 @@ fn info(args: ActionInfo) -> Result<(), repak::Error> {
     Ok(())
 }
 
-fn list(args: ActionInfo) -> Result<(), repak::Error> {
+fn list(aes_key: Option<aes::Aes256>, action: ActionInfo) -> Result<(), repak::Error> {
-    let pak = repak::PakReader::new_any(
+    let pak = repak::PakReader::new_any(BufReader::new(File::open(action.input)?), aes_key)?;
-        BufReader::new(File::open(&args.input)?),
-        args.aes_key.map(|k| aes_key(k.as_str())).transpose()?,
-    )?;
     for f in pak.files() {
         println!("{f}");
     }
@@ -180,29 +174,26 @@ fn list(args: ActionInfo) -> Result<(), repak::Error> {
 
 const STYLE: &str = "[{elapsed_precise}] [{wide_bar}] {pos}/{len} ({eta})";
 
-fn unpack(args: ActionUnpack) -> Result<(), repak::Error> {
+fn unpack(aes_key: Option<aes::Aes256>, action: ActionUnpack) -> Result<(), repak::Error> {
-    let pak = repak::PakReader::new_any(
+    let pak = repak::PakReader::new_any(BufReader::new(File::open(&action.input)?), aes_key)?;
-        BufReader::new(File::open(&args.input)?),
+    let output = action
-        args.aes_key.map(|k| aes_key(k.as_str())).transpose()?,
-    )?;
-    let output = args
         .output
         .map(PathBuf::from)
-        .unwrap_or_else(|| Path::new(&args.input).with_extension(""));
+        .unwrap_or_else(|| Path::new(&action.input).with_extension(""));
     match fs::create_dir(&output) {
         Ok(_) => Ok(()),
         Err(ref e) if e.kind() == std::io::ErrorKind::AlreadyExists => Ok(()),
         Err(e) => Err(e),
     }?;
-    if !args.force && output.read_dir()?.next().is_some() {
+    if !action.force && output.read_dir()?.next().is_some() {
         return Err(repak::Error::OutputNotEmpty(
             output.to_string_lossy().to_string(),
         ));
     }
     let mount_point = PathBuf::from(pak.mount_point());
-    let prefix = Path::new(&args.strip_prefix);
+    let prefix = Path::new(&action.strip_prefix);
 
-    let includes = args
+    let includes = action
         .include
         .iter()
         .map(|i| prefix.join(Path::new(i)))
@@ -255,9 +246,9 @@ fn unpack(args: ActionUnpack) -> Result<(), repak::Error> {
         .progress_with_style(indicatif::ProgressStyle::with_template(STYLE).unwrap());
     let progress = iter.progress.clone();
     iter.try_for_each_init(
-        || (progress.clone(), File::open(&args.input)),
+        || (progress.clone(), File::open(&action.input)),
         |(progress, file), entry| -> Result<(), repak::Error> {
-            if args.verbose {
+            if action.verbose {
                 progress.println(format!("unpacking {}", entry.entry_path));
             }
             fs::create_dir_all(&entry.out_dir)?;
@@ -335,12 +326,9 @@ fn pack(args: ActionPack) -> Result<(), repak::Error> {
     Ok(())
 }
 
-fn get(args: ActionGet) -> Result<(), repak::Error> {
+fn get(aes_key: Option<aes::Aes256>, args: ActionGet) -> Result<(), repak::Error> {
     let mut reader = BufReader::new(File::open(&args.input)?);
-    let pak = repak::PakReader::new_any(
+    let pak = repak::PakReader::new_any(&mut reader, aes_key)?;
-        &mut reader,
-        args.aes_key.map(|k| aes_key(k.as_str())).transpose()?,
-    )?;
     let mount_point = PathBuf::from(pak.mount_point());
     let prefix = Path::new(&args.strip_prefix);