From 1a282c0912c2a5d0d6b13ab281670cf3b2f93f10 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: Wed, 4 Sep 2024 19:13:03 +0100
Subject: [PATCH] feat: qmc v2 map cipher

---
 um_crypto/qmc/src/lib.rs        |  6 +++-
 um_crypto/qmc/src/v2_map/key.rs | 57 +++++++++++++++++++++++++++++++++
 um_crypto/qmc/src/v2_map/mod.rs | 40 +++++++++++++++++++++++
 3 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 um_crypto/qmc/src/v2_map/key.rs
 create mode 100644 um_crypto/qmc/src/v2_map/mod.rs

diff --git a/um_crypto/qmc/src/lib.rs b/um_crypto/qmc/src/lib.rs
index 65864b6..6110d33 100644
--- a/um_crypto/qmc/src/lib.rs
+++ b/um_crypto/qmc/src/lib.rs
@@ -1,7 +1,11 @@
 use thiserror::Error;
 
 pub mod v1;
+pub mod v2_map;
 pub mod v2_rc4;
 
 #[derive(Error, Debug)]
-pub enum QmcCryptoError {}
+pub enum QmcCryptoError {
+    #[error("QMC V2/Map Cipher: Key is empty")]
+    QMCV2MapKeyEmpty,
+}
diff --git a/um_crypto/qmc/src/v2_map/key.rs b/um_crypto/qmc/src/v2_map/key.rs
new file mode 100644
index 0000000..58b02c9
--- /dev/null
+++ b/um_crypto/qmc/src/v2_map/key.rs
@@ -0,0 +1,57 @@
+use crate::v1::cipher::V1_KEY_SIZE;
+use crate::QmcCryptoError;
+
+const INDEX_OFFSET: usize = 71214;
+
+pub fn key_compress<T: AsRef<[u8]>>(long_key: T) -> anyhow::Result<[u8; V1_KEY_SIZE]> {
+    let long_key = long_key.as_ref();
+    if long_key.is_empty() {
+        Err(QmcCryptoError::QMCV2MapKeyEmpty)?;
+    }
+
+    let n = long_key.len();
+    let mut result = [0u8; V1_KEY_SIZE];
+
+    let key_stream = (0..V1_KEY_SIZE).map(|i| {
+        let i = (i * i + INDEX_OFFSET) % n;
+        let key = long_key[i];
+        let shift = ((i + 4) % 8) as u32;
+        key.wrapping_shl(shift) | key.wrapping_shr(shift)
+    });
+
+    for (key, value) in result.iter_mut().zip(key_stream) {
+        *key = value;
+    }
+
+    Ok(result)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_compress() {
+        let expected = [
+            0x79, 0xf4, 0x00, 0x75, 0x9e, 0x36, 0x00, 0x14, 0x8a, 0x63, 0x00, 0xb4, 0xbe, 0x77,
+            0x00, 0x17, 0xba, 0x00, 0x37, 0x00, 0x00, 0x00, 0xbf, 0x80, 0x41, 0xbf, 0x83, 0xdd,
+            0xbc, 0x5c, 0x02, 0x43, 0x14, 0x82, 0x49, 0x02, 0x00, 0x55, 0xbe, 0x6d, 0xbf, 0x49,
+            0x80, 0x8e, 0x43, 0x00, 0xfa, 0x41, 0x67, 0xa8, 0x17, 0xf4, 0xae, 0x16, 0x15, 0x00,
+            0xc1, 0x37, 0x82, 0xdd, 0x36, 0x21, 0x38, 0x55, 0x00, 0x79, 0x41, 0x9e, 0x42, 0xc1,
+            0x36, 0xfa, 0xcf, 0x35, 0x00, 0x00, 0x41, 0xdd, 0x43, 0x42, 0x17, 0x4d, 0x8e, 0x8a,
+            0xdd, 0x00, 0xbe, 0xf5, 0x38, 0xb4, 0xbf, 0x00, 0x7a, 0xcc, 0x4d, 0x02, 0x00, 0xcf,
+            0xc1, 0xc1, 0x02, 0xa8, 0x00, 0x16, 0xc1, 0xbf, 0xc2, 0x42, 0x00, 0x49, 0x00, 0xc1,
+            0xc2, 0xf5, 0x00, 0x17, 0x41, 0xdc, 0x83, 0xc2, 0x00, 0x9e, 0x41, 0xc1, 0x71, 0x36,
+            0x00, 0x80,
+        ];
+        let key = (b'a'..=b'z')
+            .chain(b'A'..=b'Z')
+            .chain(b'0'..=b'9')
+            .cycle()
+            .take(325)
+            .collect::<Vec<u8>>();
+
+        let actual = key_compress(&key).expect("should compress ok");
+        assert_eq!(expected, actual);
+    }
+}
diff --git a/um_crypto/qmc/src/v2_map/mod.rs b/um_crypto/qmc/src/v2_map/mod.rs
new file mode 100644
index 0000000..66bbb89
--- /dev/null
+++ b/um_crypto/qmc/src/v2_map/mod.rs
@@ -0,0 +1,40 @@
+mod key;
+
+use crate::v1::cipher::{qmc1_transform, V1_KEY_SIZE};
+use crate::v2_map::key::key_compress;
+use anyhow::Result;
+
+pub struct QMC2Map {
+    key: [u8; V1_KEY_SIZE],
+}
+
+impl QMC2Map {
+    pub fn new<T: AsRef<[u8]>>(key: T) -> Result<Self> {
+        let key = key_compress(key)?;
+        Ok(Self { key })
+    }
+
+    pub fn decrypt<T: AsMut<[u8]>>(&self, data: &mut T, offset: usize) {
+        for (i, datum) in data.as_mut().iter_mut().enumerate() {
+            *datum = qmc1_transform(&self.key, *datum, offset + i);
+        }
+    }
+}
+
+#[test]
+fn test_decrypt() {
+    let key = (b'a'..=b'z')
+        .chain(b'A'..=b'Z')
+        .chain(b'0'..=b'9')
+        .cycle()
+        .take(325)
+        .collect::<Vec<u8>>();
+
+    let cipher = QMC2Map::new(key).expect("should not fail");
+    let mut actual = [
+        0x00u8, 0x9e, 0x41, 0xc1, 0x71, 0x36, 0x00, 0x80, 0xf4, 0x00, 0x75, 0x9e, 0x36, 0x00, 0x14,
+        0x8a,
+    ];
+    cipher.decrypt(&mut actual, 32760);
+    assert_eq!(actual, [0u8; 0x10]);
+}