From a9c7ba9fd4ecfe81297134e94ef90beb5677cac0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=B2=81=E6=A0=91=E4=BA=BA?= <lu.shuren@um-react.app>
Date: Fri, 6 Sep 2024 00:52:13 +0100
Subject: [PATCH] refactor: unify qmc2 decrypt api

---
 um_crypto/qmc/src/footer/mod.rs    |  2 ++
 um_crypto/qmc/src/lib.rs           | 34 ++++++++++++++++++++++++++++++
 um_crypto/qmc/src/v2_map/mod.rs    |  5 ++++-
 um_crypto/qmc/src/v2_rc4/cipher.rs | 27 +++++++++++++-----------
 4 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/um_crypto/qmc/src/footer/mod.rs b/um_crypto/qmc/src/footer/mod.rs
index faecfe0..83e5ec1 100644
--- a/um_crypto/qmc/src/footer/mod.rs
+++ b/um_crypto/qmc/src/footer/mod.rs
@@ -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")]
diff --git a/um_crypto/qmc/src/lib.rs b/um_crypto/qmc/src/lib.rs
index 355151b..090b952 100644
--- a/um_crypto/qmc/src/lib.rs
+++ b/um_crypto/qmc/src/lib.rs
@@ -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> {
diff --git a/um_crypto/qmc/src/v2_map/mod.rs b/um_crypto/qmc/src/v2_map/mod.rs
index 66bbb89..c39ae2f 100644
--- a/um_crypto/qmc/src/v2_map/mod.rs
+++ b/um_crypto/qmc/src/v2_map/mod.rs
@@ -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);
         }
diff --git a/um_crypto/qmc/src/v2_rc4/cipher.rs b/um_crypto/qmc/src/v2_rc4/cipher.rs
index 6069e82..22ac55b 100644
--- a/um_crypto/qmc/src/v2_rc4/cipher.rs
+++ b/um_crypto/qmc/src/v2_rc4/cipher.rs
@@ -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]);
     }
 }