From 3af2f1e131dd0398e718691e60070cd67983b281 Mon Sep 17 00:00:00 2001
From: Sosthene00 <674694@protonmail.ch>
Date: Fri, 29 Mar 2024 20:08:02 +0100
Subject: [PATCH] Allow encryption of 32 bytes array

---
 crates/sp_client/src/aesgcm.rs | 37 ++++++++++++++++++++++++++++++++++
 crates/sp_client/src/user.rs   |  2 +-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/crates/sp_client/src/aesgcm.rs b/crates/sp_client/src/aesgcm.rs
index b3d55c7..ff4a8b0 100644
--- a/crates/sp_client/src/aesgcm.rs
+++ b/crates/sp_client/src/aesgcm.rs
@@ -29,6 +29,8 @@ use rand::{thread_rng, RngCore};
 
 const HALFKEYSIZE: usize = SECRET_KEY_SIZE / 2;
 
+const THIRTYTWO: usize = 32;
+
 pub struct HalfKey([u8; HALFKEYSIZE]);
 
 impl TryFrom<Vec<u8>> for HalfKey {
@@ -56,6 +58,7 @@ impl HalfKey {
 
 pub enum Purpose {
     Login,
+    ThirtyTwoBytes,
 }
 
 pub type CipherText = Vec<u8>;
@@ -118,6 +121,10 @@ impl Aes256Decryption {
             Purpose::Login => {
                 let half_key = self.decrypt_login()?;
                 Ok(half_key.to_inner())
+            },
+            Purpose::ThirtyTwoBytes => {
+                let thirty_two_buf = self.decrypt_thirty_two()?;
+                Ok(thirty_two_buf.to_vec())
             }
         }
     }
@@ -134,6 +141,19 @@ impl Aes256Decryption {
         key_half.copy_from_slice(&plain);
         Ok(HalfKey(key_half))
     }
+
+    fn decrypt_thirty_two(&self) -> Result<[u8; THIRTYTWO]> {
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let plain = cipher
+            .decrypt(&self.nonce.into(), &*self.cipher_text)
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        if plain.len() != THIRTYTWO {
+            return Err(Error::msg("Plain text of invalid length, should be 32"));
+        }
+        let mut thirty_two = [0u8; THIRTYTWO];
+        thirty_two.copy_from_slice(&plain);
+        Ok(thirty_two)
+    }
 }
 
 pub struct Aes256Encryption {
@@ -205,6 +225,7 @@ impl Aes256Encryption {
     pub fn encrypt_with_aes_key(&self) -> Result<CipherText> {
         match self.purpose {
             Purpose::Login => self.encrypt_login(),
+            Purpose::ThirtyTwoBytes => self.encrypt_thirty_two()
         }
     }
 
@@ -219,6 +240,22 @@ impl Aes256Encryption {
         res.extend_from_slice(&cipher_text);
         Ok(res)
     }
+
+    fn encrypt_thirty_two(&self) -> Result<CipherText> {
+        if self.plaintext.len() != 32 {
+            return Err(Error::msg("Invalid length, should be 32"));
+        }
+        let mut thirty_two = [0u8;32];
+        thirty_two.copy_from_slice(&self.plaintext);
+        let cipher = Aes256Gcm::new(&self.aes_key.into());
+        let cipher_text = cipher
+            .encrypt(&self.nonce.into(), thirty_two.as_slice())
+            .map_err(|e| Error::msg(format!("{}", e)))?;
+        let mut res = Vec::with_capacity(self.nonce.len() + cipher_text.len());
+        res.extend_from_slice(&self.nonce);
+        res.extend_from_slice(&cipher_text);
+        Ok(res)
+    }
 }
 
 #[cfg(test)]
diff --git a/crates/sp_client/src/user.rs b/crates/sp_client/src/user.rs
index 6dba477..d2f15db 100644
--- a/crates/sp_client/src/user.rs
+++ b/crates/sp_client/src/user.rs
@@ -191,7 +191,7 @@ impl User {
         let hash3 = sha256::Hash::from_engine(engine);
 
         let scan_key_encryption = Aes256Encryption::import_key(
-            Purpose::Login,
+            Purpose::ThirtyTwoBytes,
             recover_scan_key.secret_bytes().to_vec(),
             hash3.to_byte_array(),
             Aes256Gcm::generate_nonce(&mut rng).into(),