sdk_client/crates/sp_client/src/user.rs

use anyhow::{Error, Result};
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 sp_backend::spclient::SpClient;
use tsify::Tsify;
use wasm_bindgen::prelude::*;

use shamir::SecretData;
use std::collections::HashMap;
use std::fs::File;
use std::io::{Cursor, Read, Write};
use std::str::FromStr;
use std::sync::{Mutex, MutexGuard, 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, SpWallet};

use crate::peers::Peer;
use crate::user;
use crate::MutexExt;
use sdk_common::crypto::{
    AeadCore, Aes256Decryption, Aes256Encryption, Aes256Gcm, HalfKey, KeyInit, Purpose,
};

type PreId = String;

const MANAGERS_NUMBER: u8 = 10;
const QUORUM_SHARD: f32 = 0.8;

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

pub fn lock_connected_users() -> Result<MutexGuard<'static, UsersMap>> {
    CONNECTED_USERS
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock_anyhow()
}

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

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

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

    pub(crate) fn get_all_outputs(&self) -> Vec<OutputList> {
        let mut res = Vec::<OutputList>::new();
        if let Some(main) = &self.main {
            res.push(main.get_outputs().clone());
        }
        if let Some(revoke) = &self.revoke {
            res.push(revoke.get_outputs().clone());
        }
        res.push(self.recover.get_outputs().clone());

        res
    }
}

#[derive(Debug, Serialize, Deserialize, Clone, Tsify)]
#[tsify(into_wasm_abi, from_wasm_abi)]
pub struct User {
    pub pre_id: PreId,
    pub processes: Vec<String>,
    pub peers: Vec<Peer>,
    recover_data: Vec<u8>,
    revoke_data: Option<Vec<u8>>,
    shares: Vec<Vec<u8>>,
    outputs: Vec<OutputList>,
}

impl User {
    pub fn new(user_wallets: UserWallets, 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_wallets.try_get_revoke() {
            revoke_data.extend_from_slice(revoke.get_client().get_scan_key().as_ref());
            revoke_data.extend_from_slice(revoke.get_client().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_wallets
            .recover
            .get_client()
            .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 threshold = (MANAGERS_NUMBER as f32 * QUORUM_SHARD).floor();
        debug_assert!(threshold > 0.0 && threshold <= u8::MAX as f32);
        let sharding = shamir::SecretData::with_secret(
            &cipher_recover_part2.to_lower_hex_string(),
            threshold as u8,
        );

        let shares: Vec<Vec<u8>> = (1..MANAGERS_NUMBER)
            .map(|x| {
                sharding.get_share(x).unwrap() // Let's trust it for now
            })
            .collect();

        //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_wallets
                .recover
                .get_client()
                .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);

        Ok(User {
            pre_id: pre_id.to_string(),
            processes: vec![process],
            peers: vec![],
            recover_data,
            revoke_data: Some(revoke_data),
            shares,
            outputs: user_wallets.get_all_outputs(),
        })
    }

    pub fn login(
        pre_id: PreId,
        user_password: String,
        recover_data: &[u8],
        shares: &[Vec<u8>],
        outputs: &[OutputList],
    ) -> 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]; // cipher length == plain.len() + 16 + nonce.len()
        let mut part1_ciphertext = [0u8; 44];

        let mut reader = Cursor::new(recover_data);
        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(),
        )?);

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

        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,
        )?;

        let recover_outputs = outputs
            .iter()
            .find(|o| o.check_fingerprint(&recover_client))
            .cloned();

        let recover_wallet = SpWallet::new(recover_client, recover_outputs)?;

        // 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(), UserWallets::new(None, recover_wallet, None));
            }
        } else {
            let mut user_map = HashMap::new();
            user_map.insert(pre_id, UserWallets::new(None, recover_wallet, 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<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 threshold = (MANAGERS_NUMBER as f32 * QUORUM_SHARD).floor();
        debug_assert!(threshold > 0.0 && threshold <= u8::MAX as f32);

        let part2_key_enc = Vec::from_hex(
            &SecretData::recover_secret(threshold as u8, shares.to_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())
    // }
}

#[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_wallets() -> UserWallets {
        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_wallets = UserWallets::new(
            Some(SpWallet::new(sp_main, None).unwrap()),
            SpWallet::new(sp_recover, None).unwrap(),
            Some(SpWallet::new(sp_revoke, None).unwrap()),
        );

        user_wallets
    }

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

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

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

        let res = User::login(
            user.pre_id.clone(),
            USER_PASSWORD.to_owned(),
            &user.recover_data,
            &user.shares,
            &user_wallets.get_all_outputs(),
        );

        assert!(res.is_ok());

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

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

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