lssa/wallet/src/lib.rs

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

use anyhow::Result;
use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
use chain_storage::WalletChainStore;
use common::{
    sequencer_client::SequencerClient,
    transaction::{EncodedTransaction, NSSATransaction},
};
use config::WalletConfig;
use log::info;
use nssa::{Account, AccountId};
use nssa_core::{Commitment, MembershipProof};
use tokio::io::AsyncWriteExt;

use crate::{
    config::PersistentStorage,
    helperfunctions::{fetch_persistent_storage, get_home, produce_data_for_storage},
    poller::TxPoller,
};

pub const HOME_DIR_ENV_VAR: &str = "NSSA_WALLET_HOME_DIR";

pub mod chain_storage;
pub mod cli;
pub mod config;
pub mod helperfunctions;
pub mod poller;
pub mod program_interactions;
pub mod token_transfers;
pub mod transaction_utils;

pub struct WalletCore {
    pub storage: WalletChainStore,
    pub poller: TxPoller,
    pub sequencer_client: Arc<SequencerClient>,
    pub last_synced_block: u64,
}

impl WalletCore {
    pub async fn start_from_config_update_chain(config: WalletConfig) -> Result<Self> {
        let client = Arc::new(SequencerClient::new(config.sequencer_addr.clone())?);
        let tx_poller = TxPoller::new(config.clone(), client.clone());

        let mut storage = WalletChainStore::new(config)?;

        let PersistentStorage {
            accounts: persistent_accounts,
            last_synced_block,
        } = fetch_persistent_storage().await?;
        for pers_acc_data in persistent_accounts {
            storage.insert_account_data(pers_acc_data);
        }

        Ok(Self {
            storage,
            poller: tx_poller,
            sequencer_client: client.clone(),
            last_synced_block,
        })
    }

    /// Store persistent data at home
    pub async fn store_persistent_data(&self) -> Result<PathBuf> {
        let home = get_home()?;
        let storage_path = home.join("storage.json");

        let data = produce_data_for_storage(&self.storage.user_data, self.last_synced_block);
        let storage = serde_json::to_vec_pretty(&data)?;

        let mut storage_file = tokio::fs::File::create(storage_path.as_path()).await?;
        storage_file.write_all(&storage).await?;

        info!("Stored data at {storage_path:#?}");

        Ok(storage_path)
    }

    /// Store persistent data at home
    pub async fn store_config_changes(&self) -> Result<PathBuf> {
        let home = get_home()?;
        let config_path = home.join("wallet_config.json");
        let config = serde_json::to_vec_pretty(&self.storage.wallet_config)?;

        let mut config_file = tokio::fs::File::create(config_path.as_path()).await?;
        config_file.write_all(&config).await?;

        info!("Stored data at {config_path:#?}");

        Ok(config_path)
    }

    pub fn create_new_account_public(&mut self) -> AccountId {
        self.storage
            .user_data
            .generate_new_public_transaction_private_key()
    }

    pub fn create_new_account_private(&mut self) -> AccountId {
        self.storage
            .user_data
            .generate_new_privacy_preserving_transaction_key_chain()
    }

    /// Get account balance
    pub async fn get_account_balance(&self, acc: AccountId) -> Result<u128> {
        Ok(self
            .sequencer_client
            .get_account_balance(acc.to_string())
            .await?
            .balance)
    }

    /// Get accounts nonces
    pub async fn get_accounts_nonces(&self, accs: Vec<AccountId>) -> Result<Vec<u128>> {
        Ok(self
            .sequencer_client
            .get_accounts_nonces(accs.into_iter().map(|acc| acc.to_string()).collect())
            .await?
            .nonces)
    }

    /// Get account
    pub async fn get_account_public(&self, account_id: AccountId) -> Result<Account> {
        let response = self
            .sequencer_client
            .get_account(account_id.to_string())
            .await?;
        Ok(response.account)
    }

    pub fn get_account_private(&self, account_id: &AccountId) -> Option<Account> {
        self.storage
            .user_data
            .user_private_accounts
            .get(account_id)
            .map(|value| value.1.clone())
    }

    pub fn get_private_account_commitment(&self, account_id: &AccountId) -> Option<Commitment> {
        let (keys, account) = self
            .storage
            .user_data
            .user_private_accounts
            .get(account_id)?;
        Some(Commitment::new(&keys.nullifer_public_key, account))
    }

    /// Poll transactions
    pub async fn poll_native_token_transfer(&self, hash: String) -> Result<NSSATransaction> {
        let transaction_encoded = self.poller.poll_tx(hash).await?;
        let tx_base64_decode = BASE64.decode(transaction_encoded)?;
        let pub_tx = borsh::from_slice::<EncodedTransaction>(&tx_base64_decode).unwrap();

        Ok(NSSATransaction::try_from(&pub_tx)?)
    }

    pub async fn check_private_account_initialized(
        &self,
        account_id: &AccountId,
    ) -> Result<Option<MembershipProof>> {
        if let Some(acc_comm) = self.get_private_account_commitment(account_id) {
            self.sequencer_client
                .get_proof_for_commitment(acc_comm)
                .await
                .map_err(anyhow::Error::from)
        } else {
            Ok(None)
        }
    }

    pub fn decode_insert_privacy_preserving_transaction_results(
        &mut self,
        tx: nssa::privacy_preserving_transaction::PrivacyPreservingTransaction,
        acc_decode_data: &[(nssa_core::SharedSecretKey, AccountId)],
    ) -> Result<()> {
        for (output_index, (secret, acc_account_id)) in acc_decode_data.iter().enumerate() {
            let acc_ead = tx.message.encrypted_private_post_states[output_index].clone();
            let acc_comm = tx.message.new_commitments[output_index].clone();

            let res_acc = nssa_core::EncryptionScheme::decrypt(
                &acc_ead.ciphertext,
                secret,
                &acc_comm,
                output_index as u32,
            )
            .unwrap();

            println!("Received new acc {res_acc:#?}");

            self.storage
                .insert_private_account_data(*acc_account_id, res_acc);
        }

        println!("Transaction data is {:?}", tx.message);

        Ok(())
    }
}