refactor: unify qmc2 decrypt api

This commit is contained in:
鲁树人 2024-09-06 00:52:13 +01:00
parent c4249226a2
commit a9c7ba9fd4
4 changed files with 55 additions and 13 deletions

View File

@ -12,6 +12,8 @@ use crate::footer::{
use anyhow::Result;
use thiserror::Error;
pub const INITIAL_DETECTION_LEN: usize = 1024;
#[derive(Error, Debug)]
pub enum FooterParseError {
#[error("Footer: Buffer too small, require at least {0} bytes")]

View File

@ -1,5 +1,9 @@
use crate::v2_map::QMC2Map;
use crate::v2_rc4::cipher::QMC2RC4;
use anyhow::Result;
use thiserror::Error;
pub mod ekey;
pub mod footer;
pub mod v1;
pub mod v2_map;
@ -11,6 +15,36 @@ pub enum QmcCryptoError {
QMCV2MapKeyEmpty,
}
pub enum QMCv2Cipher {
MapL(QMC2Map),
RC4(QMC2RC4),
}
impl QMCv2Cipher {
pub fn new<T>(key: T) -> Result<Self>
where
T: AsRef<[u8]>,
{
let key = key.as_ref();
let cipher = match key.len() {
0 => Err(QmcCryptoError::QMCV2MapKeyEmpty)?,
..=300 => QMCv2Cipher::MapL(QMC2Map::new(key)?),
_ => QMCv2Cipher::RC4(QMC2RC4::new(key)),
};
Ok(cipher)
}
pub fn decrypt<T>(&self, data: &mut T, offset: usize)
where
T: AsMut<[u8]> + ?Sized,
{
match self {
QMCv2Cipher::MapL(cipher) => cipher.decrypt(data, offset),
QMCv2Cipher::RC4(cipher) => cipher.decrypt(data, offset),
}
}
}
#[cfg(test)]
mod test {
pub fn generate_key(len: usize) -> Vec<u8> {

View File

@ -14,7 +14,10 @@ impl QMC2Map {
Ok(Self { key })
}
pub fn decrypt<T: AsMut<[u8]>>(&self, data: &mut T, offset: usize) {
pub fn decrypt<T>(&self, data: &mut T, offset: usize)
where
T: AsMut<[u8]> + ?Sized,
{
for (i, datum) in data.as_mut().iter_mut().enumerate() {
*datum = qmc1_transform(&self.key, *datum, offset + i);
}

View File

@ -27,16 +27,16 @@ impl QMC2RC4 {
}
}
fn transform_first_segment(&mut self, offset: usize, dst: &mut [u8]) {
fn process_first_segment(&self, data: &mut [u8], offset: usize) {
let n = self.key.len();
for (value, offset) in dst.iter_mut().zip(offset..) {
for (datum, offset) in data.iter_mut().zip(offset..) {
let idx = get_segment_key(offset as u64, self.key[offset % n], self.hash) as usize;
*value ^= self.key[idx % n];
*datum ^= self.key[idx % n];
}
}
fn transform_other_segment(&mut self, offset: usize, data: &mut [u8]) {
fn process_other_segment(&self, data: &mut [u8], offset: usize) {
let n = self.key.len();
let id = offset / OTHER_SEGMENT_SIZE;
@ -52,14 +52,17 @@ impl QMC2RC4 {
}
}
pub fn transform(&mut self, start_offset: usize, data: &mut [u8]) {
let mut offset = start_offset;
let mut buffer = data;
pub fn decrypt<T>(&self, data: &mut T, offset: usize)
where
T: AsMut<[u8]> + ?Sized,
{
let mut offset = offset;
let mut buffer = data.as_mut();
if offset < FIRST_SEGMENT_SIZE {
let n = min(FIRST_SEGMENT_SIZE - offset, buffer.len());
let (block, rest) = buffer.split_at_mut(n);
buffer = rest;
self.transform_first_segment(offset, block);
self.process_first_segment(block, offset);
offset += n;
}
@ -69,7 +72,7 @@ impl QMC2RC4 {
let n = min(OTHER_SEGMENT_SIZE - excess, buffer.len());
let (block, rest) = buffer.split_at_mut(n);
buffer = rest;
self.transform_other_segment(offset, block);
self.process_other_segment(block, offset);
offset += n;
}
};
@ -78,7 +81,7 @@ impl QMC2RC4 {
let n = min(OTHER_SEGMENT_SIZE, buffer.len());
let (block, rest) = buffer.split_at_mut(n);
buffer = rest;
self.transform_other_segment(offset, block);
self.process_other_segment(block, offset);
offset += n;
}
}
@ -118,8 +121,8 @@ mod tests {
.take(512)
.collect::<Vec<u8>>();
let mut cipher = QMC2RC4::new(&key);
cipher.transform(0, &mut data);
let cipher = QMC2RC4::new(&key);
cipher.decrypt(&mut data, 0);
assert_eq!(data, [0u8; 256]);
}
}