Refactore user + test

crates/sp_client/src/user.rs

@@ -5,7 +5,6 @@ use aes_gcm::KeyInit;
 use aes_gcm::{aead::Buffer, Aes256Gcm, Key};
 use anyhow::{Error, Result};
 use bytes::Buf;
-use js_sys::JsString;
 use rand::{self, thread_rng, Rng, RngCore};
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
@@ -15,7 +14,6 @@ use sp_backend::bitcoin::hex::{DisplayHex, FromHex};
 use sp_backend::bitcoin::secp256k1::SecretKey;
 use tsify::Tsify;
 use wasm_bindgen::prelude::*;
-use wasm_bindgen::JsValue;
 
 use bytes::Bytes;
 use shamir::SecretData;
@@ -33,8 +31,9 @@ use sp_backend::spclient::{OutputList, SpClient};
 use img_parts::jpeg::Jpeg;
 use img_parts::{ImageEXIF, ImageICC};
 
+use crate::aesgcm::Aes256Decryption;
 use crate::aesgcm::HalfKey;
-use crate::aesgcm::{Aes256Encryption, EncryptionTarget};
+use crate::aesgcm::{Aes256Encryption, Purpose};
 
 //extern crate shamir;
 //use shamir::SecretData;
@@ -42,8 +41,8 @@ use crate::aesgcm::{Aes256Encryption, EncryptionTarget};
 #[derive(Debug, Serialize, Deserialize, Clone, Tsify)]
 #[tsify(into_wasm_abi, from_wasm_abi)]
 pub struct User {
-    image_recover: BackUpImage,
+    recover_data: Vec<u8>,
-    image_revoke: BackUpImage,
+    revoke_data: Vec<u8>,
     sharding: Sharding,
     pre_id: String,
     recovered_spend_key: Option<String>,
@@ -54,19 +53,15 @@ impl User {
         recover_spend_key: SecretKey,
         revoke_spend_key: SecretKey,
         revoke_scan_key: SecretKey,
-        user_password: &JsString,
+        user_password: String,
-        recover_image: &[u8],
-        revoke_image: &[u8],
     ) -> Result<Self> {
         let mut rng = thread_rng();
 
         // image revoke
-        // We just take the 2 revoke keys and put it in the revoke_img file
+        // We just take the 2 revoke keys
-        let mut revoke_data = [0u8; SECRET_KEY_SIZE * 2];
+        let mut revoke_data = Vec::with_capacity(64);
-        revoke_data[..SECRET_KEY_SIZE].copy_from_slice(revoke_scan_key.as_ref());
+        revoke_data.extend_from_slice(revoke_scan_key.as_ref());
-        revoke_data[SECRET_KEY_SIZE..].copy_from_slice(revoke_spend_key.as_ref());
+        revoke_data.extend_from_slice(revoke_spend_key.as_ref());
-
-        let revoke_img_with_data = BackUpImage::new_revoke(revoke_image, &revoke_data)?;
 
         // split recover spend key
         let (part1_key, part2_key) = recover_spend_key.as_ref().split_at(SECRET_KEY_SIZE / 2);
@@ -79,20 +74,16 @@ impl User {
         recover_data.extend_from_slice(&entropy_1);
         recover_data.extend_from_slice(&entropy_2);
 
-        // convert the password in a String, i.e. a Vec<u8>
-        // Be careful of javascript strings: https://rustwasm.github.io/wasm-bindgen/reference/types/str.html#utf-16-vs-utf-8
-        assert!(user_password.is_valid_utf16()); // we can think better than panicking in this case
-        let password: String = user_password.into();
-
         // hash the concatenation
         let mut engine = sha256::HashEngine::default();
-        engine.write_all(&password.as_bytes());
+        engine.write_all(&user_password.as_bytes());
         engine.write_all(&entropy_1);
         let hash1 = sha256::Hash::from_engine(engine);
 
         // take it as a AES key
         let part1_encryption = Aes256Encryption::import_key(
-            EncryptionTarget::Login(part1_key.try_into()?),
+            Purpose::Login,
+            part1_key.to_vec(),
             hash1.to_byte_array(),
             Aes256Gcm::generate_nonce(&mut rng).into(),
         )?;
@@ -100,22 +91,18 @@ impl User {
         // encrypt the part1 of the key
         let cipher_recover_part1 = part1_encryption.encrypt_with_aes_key()?;
 
-        log::debug!("cipher_part1 length: {}", cipher_recover_part1.len());
-
         recover_data.extend_from_slice(&cipher_recover_part1);
 
-        //image recover
-        let recover_img_with_data = BackUpImage::new_recover(recover_image, &recover_data)?;
-
         // encrypt the part 2 of the key
         let mut engine = sha256::HashEngine::default();
-        engine.write_all(&password.as_bytes());
+        engine.write_all(&user_password.as_bytes());
         engine.write_all(&entropy_2);
         let hash2 = sha256::Hash::from_engine(engine);
 
         // take it as a AES key
         let part2_encryption = Aes256Encryption::import_key(
-            EncryptionTarget::Login(part2_key.try_into()?),
+            Purpose::Login,
+            part2_key.to_vec(),
             hash2.to_byte_array(),
             Aes256Gcm::generate_nonce(&mut rng).into(),
         )?;
@@ -128,7 +115,7 @@ impl User {
 
         //Pre ID
         let mut engine = sha256::HashEngine::default();
-        engine.write_all(&password.as_bytes());
+        engine.write_all(&user_password.as_bytes());
         engine.write_all(&cipher_recover_part1);
         let pre_id = sha256::Hash::from_engine(engine);
 
@@ -139,46 +126,34 @@ impl User {
         //Receive List Items (PCD)
 
         Ok(User {
-            image_recover: recover_img_with_data,
+            recover_data,
-            image_revoke: revoke_img_with_data,
+            revoke_data,
             sharding,
             pre_id: pre_id.to_string(),
             recovered_spend_key: None,
         })
     }
 
-    pub fn login(&self, user_password: JsString, image_recover: &[u8]) -> Result<SecretKey> {
+    pub fn login(user_password: String, recover_data: &[u8], sharding: Sharding) -> Result<SecretKey> {
         let mut retrieved_key = [0u8; 32];
         let mut entropy1 = [0u8; 32];
         let mut entropy2 = [0u8; 32];
-        let mut nonce1 = [0u8; 12];
-        let mut nonce2 = [0u8; 12];
         let mut part1_ciphertext = Vec::with_capacity(32); // just a guess
 
-        assert!(user_password.is_valid_utf16());
+        let mut reader = recover_data.reader();
-        let password: String = user_password.into();
-
-        let exif_image_bytes =
-            Bytes::from(read_exif(image_recover).map_err(|e| Error::msg(format!("{}", e)))?);
-        let mut reader = exif_image_bytes.reader();
         reader.read_exact(&mut entropy1)?;
         reader.read_exact(&mut entropy2)?;
-        reader.read_exact(&mut nonce1)?;
-        reader.read_exact(&mut nonce2)?;
         reader.read_to_end(&mut part1_ciphertext)?;
 
-        retrieved_key[..16].copy_from_slice(&Self::recover_part1(
+        retrieved_key[..16].copy_from_slice(&Self::recover_part1(&user_password, &entropy1, part1_ciphertext)?);
-            &password,
-            &entropy1,
-            &nonce1,
-            &part1_ciphertext,
-        )?);
 
         //@todo: get shardings from member managers!
-        let shardings = self.sharding.shares_vec.clone(); // temporary
+        let shardings = sharding.shares_vec.clone(); // temporary
 
         retrieved_key[16..].copy_from_slice(&Self::recover_part2(
-            &password, &entropy2, &nonce2, shardings,
+            &user_password,
+            &entropy2,
+            shardings,
         )?);
 
         let key = SecretKey::from_slice(&retrieved_key)?;
@@ -189,26 +164,21 @@ impl User {
     fn recover_part1(
         password: &str,
         entropy: &[u8],
-        nonce: &[u8],
+        part1_ciphertext: Vec<u8>,
-        part1_ciphertext: &[u8],
     ) -> Result<Vec<u8>> {
         let mut engine = sha256::HashEngine::default();
         engine.write_all(&password.as_bytes());
         engine.write_all(&entropy);
         let hash = sha256::Hash::from_engine(engine);
 
-        let key: &Key<Aes256Gcm> = hash.as_byte_array().try_into()?;
+        let aes_dec = Aes256Decryption::new(Purpose::Login, part1_ciphertext, hash.to_byte_array().to_vec(), None)?;
-        let cipher = Aes256Gcm::new(&key);
 
-        cipher
+        aes_dec.decrypt_with_key()
-            .decrypt(nonce.into(), part1_ciphertext)
-            .map_err(|e| anyhow::Error::msg(format!("{}", e)))
     }
 
     fn recover_part2(
         password: &str,
         entropy: &[u8],
-        nonce: &[u8],
         shares_vec: Vec<Vec<u8>>,
     ) -> Result<Vec<u8>> {
         let mut engine = sha256::HashEngine::default();
@@ -224,12 +194,9 @@ impl User {
                 .ok_or_else(|| anyhow::Error::msg("Failed to retrieve the sharded secret"))?,
         )?;
 
-        let key: &Key<Aes256Gcm> = hash.as_byte_array().try_into()?;
+        let aes_dec = Aes256Decryption::new(Purpose::Login, part2_key_enc, hash.to_byte_array().to_vec(), None)?;
-        let cipher = Aes256Gcm::new(&key);
 
-        cipher
+        aes_dec.decrypt_with_key()
-            .decrypt(nonce.into(), &*part2_key_enc)
-            .map_err(|e| anyhow::Error::msg(format!("{}", e)))
     }
 
     //not used
@@ -246,10 +213,6 @@ impl User {
     //     sha_256(&password_hash)
     // }
 
-    pub fn get_exif_image(&self, image: &[u8]) -> Vec<u8> {
-        return read_exif(image).expect("Error reading the exif");
-    }
-
     // Test sharing JS side
     pub fn get_shares(&self) -> Vec<String> {
         self.sharding.shares_format_str.clone()
@@ -285,7 +248,7 @@ impl BackUpImage {
 
     pub fn new_revoke(image: &[u8], data: &[u8]) -> Result<Self> {
         let img = write_exif(image, data)?;
-        Ok(Self::Recover(img))
+        Ok(Self::Revoke(img))
     }
 }
 
@@ -384,3 +347,46 @@ pub fn read_exif(image: &[u8]) -> Result<Vec<u8>, String> {
 //     let base64_string = base64::encode(decoded_data);
 //     base64_string
 // }
+
+#[cfg(test)]
+mod tests {
+    use super::*; // Import everything from the outer module
+
+    const RECOVER_SPEND: &str = "394ef7757f5bc8cd692337c62abf6fa0ce9932fd4ec6676daddfbe3c1b3b9d11";
+    const REVOKE_SPEND: &str = "821c1a84fa9ee718c02005505fb8315bd479c7b9a878b1eff45929c48dfcaf28";
+    const REVOKE_SCAN: &str = "a0f36cbc380624fa7eef022f39cab2716333451649dd8eb78e86d2e76bdb3f47";
+    const USER_PASSWORD: &str = "correct horse battery staple";
+
+    // Test 1: Create User
+    #[test]
+    fn test_successful_creation() {
+        let result = User::new(
+            SecretKey::from_str(RECOVER_SPEND).unwrap(),
+            SecretKey::from_str(REVOKE_SPEND).unwrap(),
+            SecretKey::from_str(REVOKE_SCAN).unwrap(),
+            USER_PASSWORD.to_owned(),
+        );
+
+        assert!(result.is_ok());
+        let user = result.unwrap();
+    }
+
+    #[test]
+    fn test_login() {
+        let user = User::new(
+            SecretKey::from_str(RECOVER_SPEND).unwrap(),
+            SecretKey::from_str(REVOKE_SPEND).unwrap(),
+            SecretKey::from_str(REVOKE_SCAN).unwrap(),
+            USER_PASSWORD.to_owned(),
+        ).unwrap();
+
+        let recover_data = user.recover_data;
+        let sharding = user.sharding;
+
+        let retrieved_recover_spend = User::login(USER_PASSWORD.to_owned(), &recover_data, sharding);
+
+        assert!(retrieved_recover_spend.is_ok());
+
+        assert!(format!("{}", retrieved_recover_spend.unwrap().display_secret()) == RECOVER_SPEND)
+    }
+}