detee-cli/src/general/grpc.rs

use crate::config::Config;
use crate::snp::grpc::proto::VmContract;
use crate::utils::sign_request;
use log::{debug, info, warn};
use tokio_stream::StreamExt;
// use tonic::metadata::errors::InvalidMetadataValue;

pub mod proto {
    pub use detee_shared::common_proto::*;
    pub use detee_shared::general_proto::*;
}

use proto::brain_general_cli_client::BrainGeneralCliClient;
use proto::{
    Account, AccountBalance, AirdropReq, BanUserReq, Empty, InspectOperatorResp, KickReq,
    ListOperatorsResp, Pubkey, RegOperatorReq, SlashReq,
};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Failed to connect to the brain: {0}")]
    BrainConnection(#[from] tonic::transport::Error),
    #[error("Received error from brain: status: {}, message: {}",
        _0.code().to_string(), _0.message())]
    ResponseStatus(#[from] tonic::Status),
    // #[error(transparent)]
    // ConfigError(#[from] crate::config::Error),
    // #[error("Could not find contract {0}")]
    // VmContractNotFound(String),
    // #[error(transparent)]
    // InternalError(#[from] InvalidMetadataValue),
    #[error(transparent)]
    InternalError(#[from] crate::utils::Error),
    #[error(transparent)]
    ConfigError(#[from] crate::config::Error),
}

pub async fn get_balance(account: &str) -> Result<AccountBalance, Error> {
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let response =
        client.get_balance(sign_request(Pubkey { pubkey: account.to_string() })?).await?;
    log::info!("Received account from brain: {response:?}");
    Ok(response.into_inner())
}

pub async fn register_operator(escrow: u64, email: String) -> Result<(), Error> {
    debug!("Connecting to brain to register operator...");
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    client
        .register_operator(sign_request(RegOperatorReq {
            pubkey: Config::get_detee_wallet()?,
            escrow,
            email,
        })?)
        .await?;
    Ok(())
}

pub async fn inspect_operator(wallet: String) -> Result<InspectOperatorResp, Error> {
    debug!("Getting information about operator {wallet} from brain.");
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    Ok(client.inspect_operator(Pubkey { pubkey: wallet }).await?.into_inner())
}

pub async fn list_operators() -> Result<Vec<ListOperatorsResp>, Error> {
    debug!("Getting contracts from brain...");
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let mut operators = Vec::new();
    let mut grpc_stream = client.list_operators(sign_request(Empty {})?).await?.into_inner();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(op) => {
                operators.push(op);
            }
            Err(e) => {
                warn!("Received error instead of operators: {e:?}");
            }
        }
    }
    debug!("Brain terminated list_operators stream.");
    Ok(operators)
}

pub async fn kick_contract(contract_uuid: String, reason: String) -> Result<u64, Error> {
    debug!("gRPC module: connecting to brain and kicking contract {contract_uuid} for reason: {reason}");
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    Ok(client
        .kick_contract(sign_request(KickReq {
            operator_wallet: Config::get_detee_wallet()?,
            contract_uuid,
            reason,
        })?)
        .await?
        .into_inner()
        .nano_lp)
}

pub async fn ban_user(user_wallet: String) -> Result<(), Error> {
    debug!("Connecting to brain to ban user...");
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    client
        .ban_user(sign_request(BanUserReq {
            operator_wallet: Config::get_detee_wallet()?,
            user_wallet,
        })?)
        .await?;
    Ok(())
}

// super admin

pub async fn admin_list_accounts() -> Result<Vec<Account>, Error> {
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let mut accounts = Vec::new();
    let mut grpc_stream = client.list_accounts(sign_request(Empty {})?).await?.into_inner();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(account) => {
                info!("Received account from brain: {account:?}");
                accounts.push(account);
            }
            Err(e) => {
                warn!("Received error instead of contracts: {e:?}");
            }
        }
    }
    debug!("Brain terminated list_contracts stream.");
    Ok(accounts)
}

pub async fn admin_list_contracts() -> Result<Vec<VmContract>, Error> {
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let mut contracts = Vec::new();
    let mut grpc_stream = client.list_all_vm_contracts(sign_request(Empty {})?).await?.into_inner();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(contract) => {
                info!("Received contract from brain: {contract:?}");
                contracts.push(contract);
            }
            Err(e) => {
                warn!("Received error instead of contracts: {e:?}");
            }
        }
    }
    debug!("Brain terminated list_contracts stream.");
    Ok(contracts)
}

pub async fn admin_airdrop(pubkey: String, tokens: u64) -> Result<(), Error> {
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let req = sign_request(AirdropReq { pubkey, tokens })?;
    let _ = client.airdrop(req).await?;
    Ok(())
}

pub async fn admin_slash(pubkey: String, tokens: u64) -> Result<(), Error> {
    let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
    let req = sign_request(SlashReq { pubkey, tokens })?;
    let _ = client.slash(req).await?;
    Ok(())
}