refactor: reorganize super-cli gRPC method to general module

change grpc type import as  proto in snp grpc module
and updated its related imports

src/bin/super-detee-cli.rs

@@ -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:?}"),
     }

src/config.rs

@@ -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;

src/general/grpc.rs

@@ -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(())
+}

src/snp/deploy.rs

@@ -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,

src/snp/grpc.rs

@@ -1,21 +1,16 @@
-pub mod brain {
+pub mod proto {
-    // tonic::include_proto!("vm_proto");
-    pub use detee_shared::common_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> {
+pub async fn get_contract_by_uuid(uuid: &str) -> Result<VmContract, Error> {
-    let req = brain::ListVmContractsReq {
+    let req = ListVmContractsReq {
         wallet: Config::get_detee_wallet()?,
         uuid: uuid.to_string(),
         ..Default::default()

src/snp/injector.rs

@@ -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)]

src/snp/mod.rs

@@ -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 {
+impl From<proto::VmContract> for VmContract {
-    fn from(brain_contract: brain::VmContract) -> Self {
+    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 {
+impl From<proto::VmNodeListResp> for TabledVmNode {
-    fn from(brain_node: brain::VmNodeListResp) -> Self {
+    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> {
+pub fn inspect_node(ip: String) -> Result<proto::VmNodeListResp, Error> {
-    let req = brain::VmNodeFilters { ip, ..Default::default() };
+    let req = proto::VmNodeFilters { ip, ..Default::default() };
     Ok(block_on(grpc::get_one_node(req))?)
 }
 

src/snp/update.rs

@@ -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,