lssa/wallet/src/helperfunctions.rs

use std::{path::PathBuf, str::FromStr, sync::Arc};

use anyhow::Result;
use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
use common::{
    block::HashableBlockData, sequencer_client::SequencerClient, transaction::NSSATransaction,
};
use key_protocol::{
    key_management::key_tree::traits::KeyNode as _, key_protocol_core::NSSAUserData,
};
use nssa::{Account, privacy_preserving_transaction::message::EncryptedAccountData};
use nssa_core::account::Nonce;
use rand::{RngCore, rngs::OsRng};
use serde::Serialize;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use crate::{
    HOME_DIR_ENV_VAR, WalletCore,
    config::{
        InitialAccountData, InitialAccountDataPrivate, InitialAccountDataPublic,
        PersistentAccountDataPrivate, PersistentAccountDataPublic, PersistentStorage, WalletConfig,
    },
};

/// Get home dir for wallet. Env var `NSSA_WALLET_HOME_DIR` must be set before execution to succeed.
pub fn get_home_nssa_var() -> Result<PathBuf> {
    Ok(PathBuf::from_str(&std::env::var(HOME_DIR_ENV_VAR)?)?)
}

/// Get home dir for wallet. Env var `HOME` must be set before execution to succeed.
pub fn get_home_default_path() -> Result<PathBuf> {
    std::env::home_dir()
        .map(|path| path.join(".nssa").join("wallet"))
        .ok_or(anyhow::anyhow!("Failed to get HOME"))
}

/// Get home dir for wallet.
pub fn get_home() -> Result<PathBuf> {
    if let Ok(home) = get_home_nssa_var() {
        Ok(home)
    } else {
        get_home_default_path()
    }
}

/// Fetch config from default home
pub async fn fetch_config() -> Result<WalletConfig> {
    let config_home = get_home()?;
    let mut config_needs_setup = false;

    let config = match tokio::fs::OpenOptions::new()
        .read(true)
        .open(config_home.join("wallet_config.json"))
        .await
    {
        Ok(mut file) => {
            let mut config_contents = vec![];
            file.read_to_end(&mut config_contents).await?;

            serde_json::from_slice(&config_contents)?
        }
        Err(err) => match err.kind() {
            std::io::ErrorKind::NotFound => {
                config_needs_setup = true;

                println!("Config not found, setting up default config");

                WalletConfig::default()
            }
            _ => anyhow::bail!("IO error {err:#?}"),
        },
    };

    if config_needs_setup {
        tokio::fs::create_dir_all(&config_home).await?;

        println!("Created configs dir at path {config_home:#?}");

        let mut file = tokio::fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(config_home.join("wallet_config.json"))
            .await?;

        let default_config_serialized =
            serde_json::to_vec_pretty(&WalletConfig::default()).unwrap();

        file.write_all(&default_config_serialized).await?;

        println!("Configs setted up");
    }

    Ok(config)
}

/// Fetch data stored at home
///
/// File must be created through setup beforehand.
pub async fn fetch_persistent_storage() -> Result<PersistentStorage> {
    let home = get_home()?;
    let accs_path = home.join("storage.json");
    let mut storage_content = vec![];

    match tokio::fs::File::open(accs_path).await {
        Ok(mut file) => {
            file.read_to_end(&mut storage_content).await?;
            Ok(serde_json::from_slice(&storage_content)?)
        }
        Err(err) => match err.kind() {
            std::io::ErrorKind::NotFound => {
                anyhow::bail!("Not found, please setup roots from config command beforehand");
            }
            _ => {
                anyhow::bail!("IO error {err:#?}");
            }
        },
    }
}

/// Produces data for storage
pub fn produce_data_for_storage(
    user_data: &NSSAUserData,
    last_synced_block: u64,
) -> PersistentStorage {
    let mut vec_for_storage = vec![];

    for (account_id, key) in &user_data.public_key_tree.account_id_map {
        if let Some(data) = user_data.public_key_tree.key_map.get(key) {
            vec_for_storage.push(
                PersistentAccountDataPublic {
                    account_id: *account_id,
                    chain_index: key.clone(),
                    data: data.clone(),
                }
                .into(),
            );
        }
    }

    for (account_id, key) in &user_data.private_key_tree.account_id_map {
        if let Some(data) = user_data.private_key_tree.key_map.get(key) {
            vec_for_storage.push(
                PersistentAccountDataPrivate {
                    account_id: *account_id,
                    chain_index: key.clone(),
                    data: data.clone(),
                }
                .into(),
            );
        }
    }

    for (account_id, key) in &user_data.default_pub_account_signing_keys {
        vec_for_storage.push(
            InitialAccountData::Public(InitialAccountDataPublic {
                account_id: account_id.to_string(),
                pub_sign_key: key.clone(),
            })
            .into(),
        )
    }

    for (account_id, (key_chain, account)) in &user_data.default_user_private_accounts {
        vec_for_storage.push(
            InitialAccountData::Private(InitialAccountDataPrivate {
                account_id: account_id.to_string(),
                account: account.clone(),
                key_chain: key_chain.clone(),
            })
            .into(),
        )
    }

    PersistentStorage {
        accounts: vec_for_storage,
        last_synced_block,
    }
}

pub(crate) fn produce_random_nonces(size: usize) -> Vec<Nonce> {
    let mut result = vec![[0; 16]; size];
    result.iter_mut().for_each(|bytes| OsRng.fill_bytes(bytes));
    result.into_iter().map(Nonce::from_le_bytes).collect()
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountPrivacyKind {
    Public,
    Private,
}

pub(crate) fn parse_addr_with_privacy_prefix(
    account_base58: &str,
) -> Result<(String, AccountPrivacyKind)> {
    if account_base58.starts_with("Public/") {
        Ok((
            account_base58.strip_prefix("Public/").unwrap().to_string(),
            AccountPrivacyKind::Public,
        ))
    } else if account_base58.starts_with("Private/") {
        Ok((
            account_base58.strip_prefix("Private/").unwrap().to_string(),
            AccountPrivacyKind::Private,
        ))
    } else {
        anyhow::bail!("Unsupported privacy kind, available variants is Public/ and Private/");
    }
}

/// Human-readable representation of an account.
#[derive(Serialize)]
pub(crate) struct HumanReadableAccount {
    balance: u128,
    program_owner_b64: String,
    data_b64: String,
    nonce: u128,
}

impl From<Account> for HumanReadableAccount {
    fn from(account: Account) -> Self {
        let program_owner_b64 = BASE64.encode(bytemuck::cast_slice(&account.program_owner));
        let data_b64 = BASE64.encode(account.data);
        Self {
            balance: account.balance,
            program_owner_b64,
            data_b64,
            nonce: account.nonce,
        }
    }
}

pub async fn parse_block_range(
    start: u64,
    stop: u64,
    seq_client: Arc<SequencerClient>,
    wallet_core: &mut WalletCore,
) -> Result<()> {
    for block_id in start..(stop + 1) {
        let block =
            borsh::from_slice::<HashableBlockData>(&seq_client.get_block(block_id).await?.block)?;

        for tx in block.transactions {
            let nssa_tx = NSSATransaction::try_from(&tx)?;

            if let NSSATransaction::PrivacyPreserving(tx) = nssa_tx {
                let mut affected_accounts = vec![];

                for (acc_account_id, (key_chain, _)) in
                    &wallet_core.storage.user_data.default_user_private_accounts
                {
                    let view_tag = EncryptedAccountData::compute_view_tag(
                        key_chain.nullifer_public_key.clone(),
                        key_chain.incoming_viewing_public_key.clone(),
                    );

                    for (ciph_id, encrypted_data) in tx
                        .message()
                        .encrypted_private_post_states
                        .iter()
                        .enumerate()
                    {
                        if encrypted_data.view_tag == view_tag {
                            let ciphertext = &encrypted_data.ciphertext;
                            let commitment = &tx.message.new_commitments[ciph_id];
                            let shared_secret = key_chain
                                .calculate_shared_secret_receiver(encrypted_data.epk.clone());

                            let res_acc = nssa_core::EncryptionScheme::decrypt(
                                ciphertext,
                                &shared_secret,
                                commitment,
                                ciph_id as u32,
                            );

                            if let Some(res_acc) = res_acc {
                                println!(
                                    "Received new account for account_id {acc_account_id:#?} with account object {res_acc:#?}"
                                );

                                affected_accounts.push((*acc_account_id, res_acc));
                            }
                        }
                    }
                }

                for keys_node in wallet_core
                    .storage
                    .user_data
                    .private_key_tree
                    .key_map
                    .values()
                {
                    let acc_account_id = keys_node.account_id();
                    let key_chain = &keys_node.value.0;

                    let view_tag = EncryptedAccountData::compute_view_tag(
                        key_chain.nullifer_public_key.clone(),
                        key_chain.incoming_viewing_public_key.clone(),
                    );

                    for (ciph_id, encrypted_data) in tx
                        .message()
                        .encrypted_private_post_states
                        .iter()
                        .enumerate()
                    {
                        if encrypted_data.view_tag == view_tag {
                            let ciphertext = &encrypted_data.ciphertext;
                            let commitment = &tx.message.new_commitments[ciph_id];
                            let shared_secret = key_chain
                                .calculate_shared_secret_receiver(encrypted_data.epk.clone());

                            let res_acc = nssa_core::EncryptionScheme::decrypt(
                                ciphertext,
                                &shared_secret,
                                commitment,
                                ciph_id as u32,
                            );

                            if let Some(res_acc) = res_acc {
                                println!(
                                    "Received new account for account_id {acc_account_id:#?} with account object {res_acc:#?}"
                                );

                                affected_accounts.push((acc_account_id, res_acc));
                            }
                        }
                    }
                }

                for (affected_account_id, new_acc) in affected_accounts {
                    wallet_core
                        .storage
                        .insert_private_account_data(affected_account_id, new_acc);
                }
            }
        }

        wallet_core.last_synced_block = block_id;
        wallet_core.store_persistent_data().await?;

        println!(
            "Block at id {block_id} with timestamp {} parsed",
            block.timestamp
        );
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_get_home_get_env_var() {
        unsafe {
            std::env::set_var(HOME_DIR_ENV_VAR, "/path/to/configs");
        }

        let home = get_home().unwrap();

        assert_eq!(PathBuf::from_str("/path/to/configs").unwrap(), home);

        unsafe {
            std::env::remove_var(HOME_DIR_ENV_VAR);
        }
    }

    #[test]
    fn test_addr_parse_with_privacy() {
        let addr_base58 = "Public/BLgCRDXYdQPMMWVHYRFGQZbgeHx9frkipa8GtpG2Syqy";
        let (_, addr_kind) = parse_addr_with_privacy_prefix(addr_base58).unwrap();

        assert_eq!(addr_kind, AccountPrivacyKind::Public);

        let addr_base58 = "Private/BLgCRDXYdQPMMWVHYRFGQZbgeHx9frkipa8GtpG2Syqy";
        let (_, addr_kind) = parse_addr_with_privacy_prefix(addr_base58).unwrap();

        assert_eq!(addr_kind, AccountPrivacyKind::Private);

        let addr_base58 = "asdsada/BLgCRDXYdQPMMWVHYRFGQZbgeHx9frkipa8GtpG2Syqy";
        assert!(parse_addr_with_privacy_prefix(addr_base58).is_err());
    }
}