diff --git a/um_cli/src/cmd/qmc2.rs b/um_cli/src/cmd/qmc2.rs index 05942b2..56cc067 100644 --- a/um_cli/src/cmd/qmc2.rs +++ b/um_cli/src/cmd/qmc2.rs @@ -1,6 +1,7 @@ use crate::Cli; use anyhow::{bail, Result}; use clap::Args; +use std::fs; use std::fs::File; use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write}; use std::path::PathBuf; @@ -20,6 +21,7 @@ pub struct ArgsQMCv2 { /// Override EKey for this file. /// Prefix with "decrypted:" to use base64 encoded raw key. + /// Prefix with "@" to read ekey from external file #[arg(short = 'K', long = "ekey")] ekey: Option, @@ -28,6 +30,35 @@ pub struct ArgsQMCv2 { info_only: bool, } +fn read_ekey(ekey: &str) -> Result> { + let mut external_file = false; + let mut decrypt_ekey = true; + + let mut ekey = ekey; + loop { + if let Some(stripped) = ekey.strip_prefix("@") { + ekey = stripped; + external_file = true; + } else if let Some(stripped) = ekey.strip_prefix("decrypted:") { + ekey = stripped; + decrypt_ekey = false; + } else { + break; + } + } + + let ekey = match external_file { + true => fs::read_to_string(ekey)?, + false => ekey.into(), + }; + let ekey = ekey.trim(); + let ekey = match decrypt_ekey { + true => umc_qmc::ekey::decrypt(ekey)?, + false => base64::decode(ekey)?.into_boxed_slice(), + }; + Ok(ekey) +} + impl ArgsQMCv2 { pub fn run(&self, cli: &Cli) -> Result { let mut file_input = File::open(&self.input)?; @@ -60,10 +91,7 @@ impl ArgsQMCv2 { let key = match ekey { None => bail!("--ekey is required when embedded ekey is not present."), - Some(ekey) => match ekey.strip_suffix("decrypted:") { - Some(decrypted) => base64::decode(decrypted)?.into_boxed_slice(), - None => umc_qmc::ekey::decrypt(ekey)?, - }, + Some(ekey) => read_ekey(ekey.as_str())?, }; let cipher = QMCv2Cipher::new(key)?;