feat: allow qmc2 to read ekey from external file

This commit is contained in:
鲁树人 2024-09-06 01:17:45 +01:00
parent e6fcb07ed2
commit c652066723

View File

@ -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<String>,
@ -28,6 +30,35 @@ pub struct ArgsQMCv2 {
info_only: bool,
}
fn read_ekey(ekey: &str) -> Result<Box<[u8]>> {
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<i32> {
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)?;