detee-cli/src/general/grpc.rs

// SPDX-License-Identifier: Apache-2.0

use crate::config::Config;
use crate::snp::grpc::proto::VmContract;
use crate::utils::sign_request;
use detee_shared::general_proto::ReportNodeReq;
use log::{debug, info, warn};
use tokio_stream::StreamExt;
use tonic::transport::Channel;

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)]
    InternalError(#[from] crate::utils::Error),
    #[error(transparent)]
    ConfigError(#[from] crate::config::Error),
    #[error("The Root CA file got corrupted.")]
    CorruptedRootCa(#[from] std::io::Error),
    #[error("Internal app error: could not parse Brain URL")]
    CorruptedBrainUrl,
}

async fn client() -> Result<BrainGeneralCliClient<Channel>, Error> {
    let default_brain_url = Config::get_brain_info().0;
    Ok(BrainGeneralCliClient::new(Config::connect_brain_channel(default_brain_url).await?))
}

pub async fn get_balance(account: &str) -> Result<AccountBalance, Error> {
    let response =
        client().await?.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...");
    client()
        .await?
        .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.");
    Ok(client().await?.inspect_operator(Pubkey { pubkey: wallet }).await?.into_inner())
}

pub async fn list_operators() -> Result<Vec<ListOperatorsResp>, Error> {
    debug!("Getting contracts from brain...");
    let mut operators = Vec::new();
    let mut grpc_stream =
        client().await?.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}");
    Ok(client()
        .await?
        .kick_contract(sign_request(KickReq {
            operator_wallet: Config::get_detee_wallet()?,
            contract_uuid,
            reason,
        })?)
        .await?
        .into_inner()
        .nano_credits)
}

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

pub async fn report_node(
    node_pubkey: String,
    contract: String,
    reason: String,
) -> Result<(), Error> {
    debug!("Getting contracts from brain...");
    client()
        .await?
        .report_node(sign_request(ReportNodeReq {
            admin_pubkey: Config::get_detee_wallet()?,
            node_pubkey,
            contract,
            reason,
        })?)
        .await?;
    Ok(())
}

// super admin

pub async fn admin_list_accounts() -> Result<Vec<Account>, Error> {
    let mut accounts = Vec::new();
    let mut grpc_stream =
        client().await?.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 contracts = Vec::new();
    let mut grpc_stream =
        client().await?.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 req = sign_request(AirdropReq { pubkey, tokens })?;
    let _ = client().await?.airdrop(req).await?;
    Ok(())
}

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