refactor: reorganize super-cli gRPC method to general module

change grpc type import as  proto in snp grpc module
and updated its related imports
This commit is contained in:
Noor 2025-03-21 18:11:04 +05:30
parent 62999d2469
commit ba872917c2
Signed by: noormohammedb
GPG Key ID: D83EFB8B3B967146
8 changed files with 109 additions and 110 deletions

@ -106,7 +106,7 @@ fn handle_completion(matches: &ArgMatches, mut cmd: Command) {
fn handle_account(matches: &ArgMatches) {
match matches.subcommand() {
Some(("show", _)) => cli_print(Ok(config::Config::get_account_data())),
Some(("list", _)) => match block_on(snp::grpc::admin_list_accounts()) {
Some(("list", _)) => match block_on(general::grpc::admin_list_accounts()) {
Ok(accounts) => {
for a in accounts {
println!("{} {} {}", a.pubkey, a.balance, a.tmp_locked);
@ -123,7 +123,7 @@ fn handle_account(matches: &ArgMatches) {
}
fn handle_contracts(_matches: &ArgMatches) {
match block_on(snp::grpc::admin_list_contracts()) {
match block_on(general::grpc::admin_list_contracts()) {
Ok(contracts) => {
for c in contracts {
println!(
@ -139,7 +139,7 @@ fn handle_contracts(_matches: &ArgMatches) {
fn handle_airdrop(matches: &ArgMatches) {
let wallet: String = matches.get_one::<String>("wallet").unwrap().clone();
let tokens = *matches.get_one::<u64>("amount").unwrap();
match block_on(snp::grpc::admin_airdrop(wallet, tokens)) {
match block_on(general::grpc::admin_airdrop(wallet, tokens)) {
Ok(()) => println!("Success."),
Err(e) => println!("Could not give airdrop due to error: {e:?}"),
}
@ -148,7 +148,7 @@ fn handle_airdrop(matches: &ArgMatches) {
fn handle_slash(matches: &ArgMatches) {
let wallet: String = matches.get_one::<String>("wallet").unwrap().clone();
let tokens = *matches.get_one::<u64>("amount").unwrap();
match block_on(snp::grpc::admin_slash(wallet, tokens)) {
match block_on(general::grpc::admin_slash(wallet, tokens)) {
Ok(()) => println!("Success."),
Err(e) => println!("Could not slash wallet due to error: {e:?}"),
}

@ -1,3 +1,4 @@
use crate::{general, utils::block_on};
use ed25519_dalek::SigningKey;
use log::{debug, info, warn};
use openssl::bn::BigNum;
@ -339,7 +340,7 @@ impl Config {
match Self::get_detee_wallet() {
Ok(key) => {
account_data.wallet_address = key.to_string();
match crate::utils::block_on(crate::snp::grpc::get_balance(&key)) {
match block_on(general::grpc::get_balance(&key)) {
Ok(account) => {
account_data.account_balance = account.balance as f64 / 1_000_000_000.0;
account_data.locked_funds = account.tmp_locked as f64 / 1_000_000_000.0;

@ -1,6 +1,7 @@
use crate::config::Config;
use crate::snp::grpc::proto::VmContract;
use crate::utils::sign_request;
use log::{debug, warn};
use log::{debug, info, warn};
use tokio_stream::StreamExt;
// use tonic::metadata::errors::InvalidMetadataValue;
@ -11,7 +12,8 @@ pub mod proto {
use proto::brain_general_cli_client::BrainGeneralCliClient;
use proto::{
BanUserReq, Empty, InspectOperatorResp, KickReq, ListOperatorsResp, Pubkey, RegOperatorReq,
Account, AccountBalance, AirdropReq, BanUserReq, Empty, InspectOperatorResp, KickReq,
ListOperatorsResp, Pubkey, RegOperatorReq, SlashReq,
};
#[derive(thiserror::Error, Debug)]
@ -33,6 +35,14 @@ pub enum Error {
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?;
@ -96,3 +106,57 @@ pub async fn ban_user(user_wallet: String) -> Result<(), Error> {
.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(())
}

@ -1,5 +1,5 @@
use super::{
grpc::{self, brain},
grpc::{self, proto},
injector, Distro, Dtrfs, Error, VmSshArgs, DEFAULT_ARCHLINUX, DEFAULT_DTRFS,
};
use crate::config::Config;
@ -106,7 +106,7 @@ impl Request {
}
// returns node IP and data regarding the new VM
fn send_vm_request(&self) -> Result<(String, brain::NewVmResp), Error> {
fn send_vm_request(&self) -> Result<(String, proto::NewVmResp), Error> {
let admin_pubkey = Config::get_detee_wallet()?;
let node = self.get_node()?;
let (extra_ports, public_ipv4): (Vec<u32>, bool) = match &self.ipv4 {
@ -130,7 +130,7 @@ impl Request {
self.hours,
self.price,
);
let brain_req = brain::NewVmReq {
let brain_req = proto::NewVmReq {
uuid: String::new(),
hostname: self.hostname.clone(),
admin_pubkey,
@ -155,12 +155,12 @@ impl Request {
Ok((node.ip, new_vm_resp))
}
pub fn get_node(&self) -> Result<grpc::brain::VmNodeListResp, Error> {
pub fn get_node(&self) -> Result<proto::VmNodeListResp, Error> {
let (free_ports, offers_ipv4) = match &self.ipv4 {
IPv4Config::PublishPorts(vec) => (vec.len() as u32, false),
IPv4Config::PublicIPv4 => (0, true),
};
let filters = brain::VmNodeFilters {
let filters = proto::VmNodeFilters {
free_ports,
offers_ipv4,
offers_ipv6: self.public_ipv6,

@ -1,21 +1,16 @@
pub mod brain {
// tonic::include_proto!("vm_proto");
pub use detee_shared::common_proto::*;
pub mod proto {
pub use detee_shared::general_proto::*;
pub use detee_shared::vm_proto::*;
}
use crate::config::Config;
use crate::snp::brain::Account;
use crate::snp::brain::AccountBalance;
use crate::snp::brain::Pubkey;
use brain::{
brain_vm_cli_client::BrainVmCliClient, DeleteVmReq, ExtendVmReq, ListVmContractsReq, NewVmReq,
NewVmResp, ReportNodeReq, UpdateVmReq, UpdateVmResp, VmContract, VmNodeFilters, VmNodeListResp,
};
use detee_shared::general_proto::brain_general_cli_client::BrainGeneralCliClient;
use lazy_static::lazy_static;
use log::{debug, info, warn};
use proto::{
brain_vm_cli_client::BrainVmCliClient, DeleteVmReq, ExtendVmReq, ListVmContractsReq, NewVmReq,
NewVmResp, ReportNodeReq, UpdateVmReq, UpdateVmResp, VmContract, VmNodeFilters, VmNodeListResp,
};
use tokio_stream::StreamExt;
use tonic::metadata::errors::InvalidMetadataValue;
use tonic::metadata::AsciiMetadataValue;
@ -118,28 +113,6 @@ pub async fn get_node_list(req: VmNodeFilters) -> Result<Vec<VmNodeListResp>, Er
Ok(nodes)
}
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 admin_airdrop(pubkey: String, tokens: u64) -> Result<(), Error> {
let mut client = BrainGeneralCliClient::connect(Config::get_brain_url()).await?;
let req = sign_request(brain::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(brain::SlashReq { pubkey, tokens })?;
let _ = client.slash(req).await?;
Ok(())
}
pub async fn get_one_node(req: VmNodeFilters) -> Result<VmNodeListResp, Error> {
let mut client = BrainVmCliClient::connect(Config::get_brain_url()).await?;
let response = client.get_one_vm_node(sign_request(req)?).await?;
@ -193,45 +166,6 @@ pub async fn list_contracts(req: ListVmContractsReq) -> Result<Vec<VmContract>,
Ok(contracts)
}
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(brain::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_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(brain::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 delete_vm(uuid: &str) -> Result<(), Error> {
let mut client = BrainVmCliClient::connect(Config::get_brain_url()).await?;
let req = DeleteVmReq { uuid: uuid.to_string(), admin_pubkey: Config::get_detee_wallet()? };
@ -291,8 +225,8 @@ pub async fn update_vm(req: UpdateVmReq) -> Result<UpdateVmResp, Error> {
}
}
pub async fn get_contract_by_uuid(uuid: &str) -> Result<brain::VmContract, Error> {
let req = brain::ListVmContractsReq {
pub async fn get_contract_by_uuid(uuid: &str) -> Result<VmContract, Error> {
let req = ListVmContractsReq {
wallet: Config::get_detee_wallet()?,
uuid: uuid.to_string(),
..Default::default()

@ -1,4 +1,4 @@
use crate::{config::Config, snp::grpc::brain};
use crate::{config::Config, snp::grpc::proto};
use log::debug;
use std::net::IpAddr;
@ -9,7 +9,7 @@ pub struct Args {
pub vcpus: u32,
pub kernel: String,
pub initrd: String,
pub args: brain::MeasurementArgs,
pub args: proto::MeasurementArgs,
}
#[derive(thiserror::Error, Debug)]

@ -9,7 +9,7 @@ use crate::{
config::{self, Config},
snp,
};
use grpc::brain;
use grpc::proto;
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use tabled::Tabled;
@ -53,10 +53,10 @@ impl crate::HumanOutput for VmSshArgs {
}
}
impl TryFrom<grpc::brain::VmContract> for VmSshArgs {
impl TryFrom<grpc::proto::VmContract> for VmSshArgs {
type Error = Error;
fn try_from(contract: grpc::brain::VmContract) -> Result<Self, Self::Error> {
fn try_from(contract: grpc::proto::VmContract) -> Result<Self, Self::Error> {
let mut args = VmSshArgs { ..Default::default() };
args.uuid = contract.uuid;
args.hostname = contract.hostname;
@ -72,7 +72,7 @@ impl TryFrom<grpc::brain::VmContract> for VmSshArgs {
"This VM does not have a public IP. Getting node IP for node {}",
contract.node_pubkey
);
let node = block_on(snp::grpc::get_one_node(brain::VmNodeFilters {
let node = block_on(snp::grpc::get_one_node(proto::VmNodeFilters {
node_pubkey: contract.node_pubkey.clone(),
..Default::default()
}))?;
@ -200,10 +200,10 @@ fn display_mins(minutes: &u64) -> String {
format!("{hours}h {mins}m")
}
impl From<brain::VmContract> for VmContract {
fn from(brain_contract: brain::VmContract) -> Self {
impl From<proto::VmContract> for VmContract {
fn from(brain_contract: proto::VmContract) -> Self {
let node_pubkey = brain_contract.node_pubkey.clone();
let location = match block_on(snp::grpc::get_one_node(brain::VmNodeFilters {
let location = match block_on(snp::grpc::get_one_node(proto::VmNodeFilters {
node_pubkey: node_pubkey.clone(),
..Default::default()
})) {
@ -241,8 +241,8 @@ pub struct TabledVmNode {
pub reports: usize,
}
impl From<brain::VmNodeListResp> for TabledVmNode {
fn from(brain_node: brain::VmNodeListResp) -> Self {
impl From<proto::VmNodeListResp> for TabledVmNode {
fn from(brain_node: proto::VmNodeListResp) -> Self {
Self {
operator: brain_node.operator,
location: brain_node.city + ", " + &brain_node.region + ", " + &brain_node.country,
@ -255,7 +255,7 @@ impl From<brain::VmNodeListResp> for TabledVmNode {
pub fn ssh(uuid: &str, just_print: bool) -> Result<VmSshArgs, Error> {
log::info!("Getting VM information about {uuid} from brain...");
let req = brain::ListVmContractsReq {
let req = proto::ListVmContractsReq {
wallet: Config::get_detee_wallet()?,
uuid: uuid.to_string(),
..Default::default()
@ -288,9 +288,9 @@ pub fn ssh(uuid: &str, just_print: bool) -> Result<VmSshArgs, Error> {
std::process::exit(2);
}
pub fn get_one_contract(uuid: &str) -> Result<snp::grpc::brain::VmContract, Error> {
pub fn get_one_contract(uuid: &str) -> Result<proto::VmContract, Error> {
log::debug!("Getting contract {uuid} from brain...");
let req = brain::ListVmContractsReq {
let req = proto::ListVmContractsReq {
wallet: Config::get_detee_wallet()?,
uuid: uuid.to_string(),
..Default::default()
@ -302,9 +302,9 @@ pub fn get_one_contract(uuid: &str) -> Result<snp::grpc::brain::VmContract, Erro
Ok(contracts[0].clone())
}
pub fn get_node_by_contract(uuid: &str) -> Result<brain::VmNodeListResp, Error> {
pub fn get_node_by_contract(uuid: &str) -> Result<proto::VmNodeListResp, Error> {
let contract = get_one_contract(uuid)?;
Ok(block_on(snp::grpc::get_one_node(brain::VmNodeFilters {
Ok(block_on(snp::grpc::get_one_node(proto::VmNodeFilters {
node_pubkey: contract.node_pubkey,
..Default::default()
}))?)
@ -315,7 +315,7 @@ pub fn delete_contract(uuid: &str) -> Result<(), Error> {
}
pub fn list_contracts(as_operator: bool) -> Result<Vec<VmContract>, Error> {
let req = brain::ListVmContractsReq {
let req = proto::ListVmContractsReq {
wallet: Config::get_detee_wallet()?,
as_operator,
..Default::default()
@ -346,7 +346,7 @@ pub fn append_uuid_list(uuid: &str, hostname: &str) -> Result<(), Error> {
Ok(())
}
impl super::HumanOutput for Vec<brain::VmNodeListResp> {
impl super::HumanOutput for Vec<proto::VmNodeListResp> {
fn human_cli_print(&self) {
let nodes: Vec<TabledVmNode> = self.iter().map(|n| n.clone().into()).collect();
let style = tabled::settings::Style::rounded();
@ -356,14 +356,14 @@ impl super::HumanOutput for Vec<brain::VmNodeListResp> {
}
}
pub fn print_nodes() -> Result<Vec<brain::VmNodeListResp>, Error> {
pub fn print_nodes() -> Result<Vec<proto::VmNodeListResp>, Error> {
log::debug!("This will support flags in the future, but we have only one node atm.");
let req = brain::VmNodeFilters { ..Default::default() };
let req = proto::VmNodeFilters { ..Default::default() };
Ok(block_on(grpc::get_node_list(req))?)
}
pub fn inspect_node(ip: String) -> Result<brain::VmNodeListResp, Error> {
let req = brain::VmNodeFilters { ip, ..Default::default() };
pub fn inspect_node(ip: String) -> Result<proto::VmNodeListResp, Error> {
let req = proto::VmNodeFilters { ip, ..Default::default() };
Ok(block_on(grpc::get_one_node(req))?)
}

@ -1,5 +1,5 @@
use super::{
grpc::{self, brain},
grpc::{self, proto},
injector, Dtrfs, Error,
};
use crate::config::Config;
@ -71,12 +71,12 @@ impl Request {
}
// returns node IP and data regarding the new VM
fn send_update_vm_request(&self, uuid: &str) -> Result<brain::UpdateVmResp, Error> {
fn send_update_vm_request(&self, uuid: &str) -> Result<proto::UpdateVmResp, Error> {
let (kernel_url, kernel_sha, dtrfs_url, dtrfs_sha) = match self.dtrfs.clone() {
Some(dtrfs) => (dtrfs.kernel_url, dtrfs.kernel_sha, dtrfs.dtrfs_url, dtrfs.dtrfs_sha),
None => (String::new(), String::new(), String::new(), String::new()),
};
Ok(block_on(grpc::update_vm(brain::UpdateVmReq {
Ok(block_on(grpc::update_vm(proto::UpdateVmReq {
uuid: uuid.to_string(),
admin_pubkey: Config::get_detee_wallet()?,
disk_size_gb: self.disk_size_gb,