sdk_client/crates/sp_client/src/user.rs

use aes::cipher::generic_array::GenericArray;
use aes_gcm::aead::Aead;
use aes_gcm::AeadCore;
use aes_gcm::KeyInit;
use aes_gcm::{aead::Buffer, Aes256Gcm, Key};
use anyhow::{Error, Result};
use bytes::Buf;
use rand::{self, thread_rng, Rng, RngCore};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sp_backend::bitcoin::hashes::Hash;
use sp_backend::bitcoin::hashes::HashEngine;
use sp_backend::bitcoin::hex::{DisplayHex, FromHex};
use sp_backend::bitcoin::secp256k1::SecretKey;
use sp_backend::bitcoin::secp256k1::ThirtyTwoByteHash;
use tsify::Tsify;
use wasm_bindgen::prelude::*;

use bytes::Bytes;
use shamir::SecretData;
use std::collections::HashMap;
use std::fs::File;
use std::io::Read;
use std::io::Write;
use std::str::FromStr;
use std::sync::{Mutex, OnceLock};

use sp_backend::bitcoin::secp256k1::constants::SECRET_KEY_SIZE;
use sp_backend::silentpayments::bitcoin_hashes::sha256;
use sp_backend::silentpayments::sending::SilentPaymentAddress;
use sp_backend::spclient::SpendKey;
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, Purpose};
use crate::user;

type PreId = String;

pub static CONNECTED_USERS: OnceLock<Mutex<HashMap<PreId, UserKeys>>> = OnceLock::new();

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UserKeys {
    pub main: Option<SpClient>,
    pub recover: SpClient,
    pub revoke: Option<SpClient>,
}

impl UserKeys {
    pub fn new(main: Option<SpClient>, recover: SpClient, revoke: Option<SpClient>) -> Self {
        Self {
            main,
            recover,
            revoke,
        }
    }

    pub fn try_get_revoke(&self) -> Option<&SpClient> {
        self.revoke.as_ref()
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct User {
    pub pre_id: String,
    pub process: String,
    recover_data: Vec<u8>,
    revoke_data: Option<Vec<u8>>,
    sharding: Sharding,
}

impl User {
    pub fn new(user_keys: UserKeys, user_password: String, process: String) -> Result<Self> {
        let mut rng = thread_rng();

        // image revoke
        // We just take the 2 revoke keys
        let mut revoke_data = Vec::with_capacity(64);
        if let Some(revoke) = user_keys.try_get_revoke() {
            revoke_data.extend_from_slice(revoke.get_scan_key().as_ref());
            revoke_data.extend_from_slice(revoke.try_get_secret_spend_key()?.as_ref());
        } else {
            return Err(Error::msg("No revoke wallet available"));
        }

        // Take the 2 recover keys
        // split recover spend key
        let recover_spend_key = user_keys.recover.try_get_secret_spend_key()?.clone();
        let (part1_key, part2_key) = recover_spend_key.as_ref().split_at(SECRET_KEY_SIZE / 2);
        let mut recover_data = Vec::<u8>::with_capacity(180); // 32 * 3 + (12+16)*3

        // generate 3 tokens of 32B entropy
        let mut entropy_1: [u8; 32] = Aes256Gcm::generate_key(&mut rng).into();
        let mut entropy_2: [u8; 32] = Aes256Gcm::generate_key(&mut rng).into();
        let mut entropy_3: [u8; 32] = Aes256Gcm::generate_key(&mut rng).into();

        recover_data.extend_from_slice(&entropy_1);
        recover_data.extend_from_slice(&entropy_2);
        recover_data.extend_from_slice(&entropy_3);

        // hash the concatenation
        let mut engine = sha256::HashEngine::default();
        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(
            Purpose::Login,
            part1_key.to_vec(),
            hash1.to_byte_array(),
            Aes256Gcm::generate_nonce(&mut rng).into(),
        )?;

        // encrypt the part1 of the key
        let cipher_recover_part1 = part1_encryption.encrypt_with_aes_key()?;

        recover_data.extend_from_slice(&cipher_recover_part1);

        //Pre ID
        let pre_id: PreId = Self::compute_pre_id(&user_password, &cipher_recover_part1);

        // encrypt the part 2 of the key
        let mut engine = sha256::HashEngine::default();
        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(
            Purpose::Login,
            part2_key.to_vec(),
            hash2.to_byte_array(),
            Aes256Gcm::generate_nonce(&mut rng).into(),
        )?;

        // encrypt the part2 of the key
        let cipher_recover_part2 = part2_encryption.encrypt_with_aes_key()?;

        //create shardings
        let sharding = Sharding::new(&cipher_recover_part2.to_lower_hex_string(), 10u8); //nMembers = 10 for testing, need to recover nmember elsewhere

        //scan key:
        let mut engine = sha256::HashEngine::default();
        engine.write_all(&user_password.as_bytes());
        engine.write_all(&entropy_3);
        let hash3 = sha256::Hash::from_engine(engine);

        let scan_key_encryption = Aes256Encryption::import_key(
            Purpose::ThirtyTwoBytes,
            user_keys.recover.get_scan_key().secret_bytes().to_vec(),
            hash3.to_byte_array(),
            Aes256Gcm::generate_nonce(&mut rng).into(),
        )?;

        // encrypt the scan key
        let cipher_scan_key = scan_key_encryption.encrypt_with_aes_key()?;

        recover_data.extend_from_slice(&cipher_scan_key);

        //Create PRDList
        //@todo
        //Send messages PRDList
        //@todo
        //Receive List Items (PCD)

        Ok(User {
            pre_id: pre_id.to_string(),
            process,
            recover_data,
            revoke_data: Some(revoke_data),
            sharding,
        })
    }

    pub fn login(
        pre_id: PreId,
        user_password: String,
        recover_data: &[u8],
        sharding: Sharding,
    ) -> Result<()> {
        let mut retrieved_spend_key = [0u8; 32];
        let mut retrieved_scan_key = [0u8; 32];
        let mut entropy1 = [0u8; 32];
        let mut entropy2 = [0u8; 32];
        let mut entropy3 = [0u8; 32];
        let mut cipher_scan_key = [0u8; 60];
        let mut part1_ciphertext = [0u8; 44];

        let mut reader = recover_data.reader();
        reader.read_exact(&mut entropy1)?;
        reader.read_exact(&mut entropy2)?;
        reader.read_exact(&mut entropy3)?;
        reader.read_exact(&mut part1_ciphertext)?;
        reader.read_exact(&mut cipher_scan_key)?;

        // We can retrieve the pre_id and check that it matches
        let retrieved_pre_id = Self::compute_pre_id(&user_password, &part1_ciphertext);

        // If pre_id is not the same, password is probably false, or the client is feeding us garbage
        if retrieved_pre_id != pre_id {
            return Err(Error::msg("pre_id and recover_data don't match"));
        }

        // If we already have loaded a user with this pre_id, abort
        if let Some(current_users) = CONNECTED_USERS.get() {
            if current_users
                .to_owned()
                .lock()
                .unwrap()
                .contains_key(&pre_id)
            {
                return Err(Error::msg(format!(
                    "User with pre_id {} already logged in",
                    pre_id
                )));
            }
        }

        retrieved_spend_key[..16].copy_from_slice(&Self::recover_part1(
            &user_password,
            &entropy1,
            part1_ciphertext.to_vec(),
        )?);

        //@todo: get shardings from member managers!
        let shardings = sharding.shares_vec.clone(); // temporary

        retrieved_spend_key[16..].copy_from_slice(&Self::recover_part2(
            &user_password,
            &entropy2,
            shardings,
        )?);

        retrieved_scan_key.copy_from_slice(&Self::recover_key_slice(
            &user_password,
            &entropy3,
            cipher_scan_key.to_vec(),
        )?);

        // we can create the recover sp_client
        let recover_client = SpClient::new(
            "".to_owned(),
            SecretKey::from_slice(&retrieved_scan_key)?,
            SpendKey::Secret(SecretKey::from_slice(&retrieved_spend_key)?),
            None,
            true,
        )?;

        // Adding user to CONNECTED_USERS
        if let Some(current_users) = CONNECTED_USERS.get() {
            let mut lock = current_users.to_owned().lock().unwrap();
            if lock.contains_key(&pre_id) {
                return Err(Error::msg(format!(
                    "User with pre_id {} already exists",
                    pre_id
                )));
            } else {
                lock.insert(pre_id.clone(), UserKeys::new(None, recover_client, None));
            }
        } else {
            let mut user_map = HashMap::new();
            user_map.insert(pre_id, UserKeys::new(None, recover_client, None));
            let new_value = Mutex::new(user_map);
            if let Err(error) = CONNECTED_USERS.set(new_value) {
                return Err(Error::msg(
                    "Failed to set the CONNECTED_USERS static variable",
                ));
            }
        }

        Ok(())
    }

    fn recover_key_slice(password: &str, entropy: &[u8], ciphertext: Vec<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 aes_dec = Aes256Decryption::new(
            Purpose::ThirtyTwoBytes,
            ciphertext,
            hash.to_byte_array().to_vec(),
            None,
        )?;

        aes_dec.decrypt_with_key()
    }

    fn recover_part1(password: &str, entropy: &[u8], ciphertext: Vec<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 aes_dec = Aes256Decryption::new(
            Purpose::Login,
            ciphertext,
            hash.to_byte_array().to_vec(),
            None,
        )?;

        aes_dec.decrypt_with_key()
    }

    fn recover_part2(password: &str, entropy: &[u8], shares_vec: Vec<Vec<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 shares_total: f32 = u8::try_from(shares_vec.len())?.try_into()?;
        let quorum_sharding: u8 = (Sharding::QUORUM_SHARD * shares_total).round() as u8;

        let part2_key_enc = Vec::from_hex(
            &SecretData::recover_secret(quorum_sharding, shares_vec)
                .ok_or_else(|| anyhow::Error::msg("Failed to retrieve the sharded secret"))?,
        )?;

        let aes_dec = Aes256Decryption::new(
            Purpose::Login,
            part2_key_enc,
            hash.to_byte_array().to_vec(),
            None,
        )?;

        aes_dec.decrypt_with_key()
    }

    fn compute_pre_id(user_password: &str, cipher_recover_part1: &[u8]) -> PreId {
        let mut engine = sha256::HashEngine::default();
        engine.write_all(&user_password.as_bytes());
        engine.write_all(&cipher_recover_part1);
        let pre_id = sha256::Hash::from_engine(engine);

        pre_id.to_string()
    }

    //not used
    // pub fn pbkdf2(password: &str, data: &str) -> String {
    //     let data_salt = data.trim_end_matches('=');
    //     let salt = SaltString::from_b64(data_salt)
    //         .map(|s| s)
    //         .unwrap_or_else(|_| panic!("Failed to parse salt value from base64 string"));

    //     let mut password_hash = String::new();
    //     if let Ok(pwd) = Scrypt.hash_password(password.as_bytes(), &salt) {
    //         password_hash.push_str(&pwd.to_string());
    //     }
    //     sha_256(&password_hash)
    // }

    // Test sharing JS side
    pub fn get_shares(&self) -> Vec<String> {
        self.sharding.shares_format_str.clone()
    }

    //Test sharing Js side
    pub fn get_secret(&self, shardings: Vec<String>) -> String {
        let mut shares_vec = Vec::new();

        for s in shardings.iter() {
            let bytes_vec: Vec<u8> = s
                .trim_matches(|c| c == '[' || c == ']')
                .split(',')
                .filter_map(|s| s.trim().parse().ok())
                .collect();
            shares_vec.push(bytes_vec);
        }
        self.sharding.recover_secrete(shares_vec.clone())
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum BackUpImage {
    Recover(Vec<u8>),
    Revoke(Vec<u8>),
}

impl BackUpImage {
    pub fn new_recover(image: &[u8], data: &[u8]) -> Result<Self> {
        let img = write_exif(image, data)?;
        Ok(Self::Recover(img))
    }

    pub fn new_revoke(image: &[u8], data: &[u8]) -> Result<Self> {
        let img = write_exif(image, data)?;
        Ok(Self::Revoke(img))
    }
}

#[derive(Debug, Serialize, Deserialize, Default, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Sharding {
    shares_vec: Vec<Vec<u8>>,
    shares_format_str: Vec<String>,
}

impl Sharding {
    const QUORUM_SHARD: f32 = 0.80_f32;
    pub fn new(part2_key_enc: &str, number_members: u8) -> Self {
        let secret_data = SecretData::with_secret(part2_key_enc, number_members);
        let mut shares_format_str: Vec<String> = Vec::new();
        let shares_vec = (1..=number_members)
            .map(|i| match secret_data.get_share(i) {
                Ok(share) => {
                    let string = format!(
                        "[{}]",
                        share
                            .clone()
                            .iter()
                            .map(|b| format!("{}", b))
                            .collect::<Vec<_>>()
                            .join(",")
                    );
                    shares_format_str.push(string.clone());
                    share
                }
                Err(_) => panic!("Not able to recover the shares!"),
            })
            .collect::<Vec<_>>();

        Sharding {
            shares_vec,
            shares_format_str,
        }
    }

    pub fn recover_secrete(&self, shares: Vec<Vec<u8>>) -> String {
        let quorum_sharding = (Self::QUORUM_SHARD * f32::from(shares.len() as u8)).round() as u8;
        SecretData::recover_secret(quorum_sharding, shares).unwrap()
    }
}

pub fn write_exif(image_to_recover: &[u8], data: &[u8]) -> Result<Vec<u8>> {
    let mut jpeg = Jpeg::from_bytes(Bytes::from(image_to_recover.to_vec()))?;
    let data_bytes = Bytes::from(data.to_owned());
    jpeg.set_exif(Some(data_bytes));
    let output_image_bytes = jpeg.encoder().bytes();
    let output_image = output_image_bytes.to_vec();
    Ok(output_image)
}

pub fn read_exif(image: &[u8]) -> Result<Vec<u8>, String> {
    let image_bytes = Bytes::from(image.to_vec());
    let jpeg = Jpeg::from_bytes(image_bytes).unwrap();

    //exif out
    let mut exif_image = Bytes::new();
    if let Some(ref meta) = jpeg.exif() {
        exif_image = meta.clone();
    } else {
        return Err("No exif data".to_string());
    }
    let exif_bytes = exif_image.as_ref();
    Ok(exif_bytes.to_vec())
}

// //change for return Result?
// pub fn from_hex_to_b58(hex_string: &str) -> String {
//     let decoded_data = hex::decode(hex_string).expect("Failed to decode hex string");
//     let base58_string = bs58::encode(decoded_data).into_string();
//     base58_string
// }
// //change for return Result?
// pub fn from_b58_to_hex(base58_string: &str) -> String {
//     let decoded_data = bs58::decode(base58_string.to_owned()).into_vec().unwrap();
//     let hex_string = decoded_data
//         .iter()
//         .map(|b| format!("{:02x}", b))
//         .collect::<String>();
//     hex_string
// }

// fn from_b64_to_hex(base64_string: &str) -> String {
//     let decoded_data = base64::decode(base64_string).unwrap();
//     let hex_string = decoded_data
//         .iter()
//         .map(|b| format!("{:02x}", b))
//         .collect::<String>();
//     hex_string
// }
// fn from_hex_to_b64(hex_string: &str) -> String {
//     let decoded_data = hex::decode(hex_string).expect("Failed to decode hex 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 RECOVER_SCAN: &str = "3aa8cc570d17ec3a4dc4136e50151cc6de26052d968abfe02a5fea724ce38205";
    const REVOKE_SPEND: &str = "821c1a84fa9ee718c02005505fb8315bd479c7b9a878b1eff45929c48dfcaf28";
    const REVOKE_SCAN: &str = "a0f36cbc380624fa7eef022f39cab2716333451649dd8eb78e86d2e76bdb3f47";
    const MAIN_SPEND: &str = "b9098a6598ac55d8dd0e6b7aab0d1f63eb8792d06143f3c0fb6f5b80476a1c0d";
    const MAIN_SCAN: &str = "79dda4031663ac2cb250c46d896dc92b3c027a48a761b2342fabf1e441ea2857";
    const USER_PASSWORD: &str = "correct horse battery staple";
    const PROCESS: &str = "example";

    fn helper_create_user_keys() -> UserKeys {
        let label = "default".to_owned();
        let sp_main = SpClient::new(
            label.clone(),
            SecretKey::from_str(MAIN_SCAN).unwrap(),
            SpendKey::Secret(SecretKey::from_str(MAIN_SPEND).unwrap()),
            None,
            true,
        )
        .unwrap();
        let sp_recover = SpClient::new(
            label.clone(),
            SecretKey::from_str(RECOVER_SCAN).unwrap(),
            SpendKey::Secret(SecretKey::from_str(RECOVER_SPEND).unwrap()),
            None,
            true,
        )
        .unwrap();
        let sp_revoke = SpClient::new(
            label.clone(),
            SecretKey::from_str(REVOKE_SCAN).unwrap(),
            SpendKey::Secret(SecretKey::from_str(REVOKE_SPEND).unwrap()),
            None,
            true,
        )
        .unwrap();
        let user_keys = UserKeys::new(Some(sp_main), sp_recover, Some(sp_revoke));

        user_keys
    }

    // Test 1: Create User
    #[test]
    fn test_successful_creation() {
        let user_keys = helper_create_user_keys();
        let result = User::new(user_keys, USER_PASSWORD.to_owned(), PROCESS.to_owned());

        assert!(result.is_ok());
        let user = result.unwrap();
    }

    #[test]
    fn test_login() {
        let user_keys = helper_create_user_keys();
        let user = User::new(user_keys, USER_PASSWORD.to_owned(), PROCESS.to_owned()).unwrap();

        let pre_id = user.pre_id;
        let recover_data = user.recover_data;
        let sharding = user.sharding;

        let retrieved_recover_spend = User::login(
            pre_id.clone(),
            USER_PASSWORD.to_owned(),
            &recover_data,
            sharding,
        );

        assert!(retrieved_recover_spend.is_ok());

        let connected = CONNECTED_USERS.get().unwrap().lock().unwrap();

        let recover = &connected.get(&pre_id).unwrap().recover;

        assert!(
            format!(
                "{}",
                recover.try_get_secret_spend_key().unwrap().display_secret()
            ) == RECOVER_SPEND
        )
    }
}