brain-mock/src/grpc.rs

pub mod snp_proto {
    tonic::include_proto!("vm_proto");
}

use crate::data::BrainData;
use crate::grpc::vm_daemon_message;
use log::info;
use snp_proto::brain_cli_server::BrainCli;
use snp_proto::brain_vm_daemon_server::BrainVmDaemon;
use snp_proto::*;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
use tonic::{Request, Response, Status, Streaming};

use detee_shared::sgx::pb::brain::brain_app_cli_server::BrainAppCli;
use detee_shared::sgx::pb::brain::brain_app_daemon_server::BrainAppDaemon;
use detee_shared::sgx::pb::brain::{
    AppContract, AppNodeFilters, AppNodeListResp, BrainMessageApp, DaemonMessageApp, DelAppReq,
    ListAppContractsReq, NewAppReq, NewAppRes, RegisterAppNodeReq,
};
const ADMIN_ACCOUNTS: &[&str] = &[
    "x52w7jARC5erhWWK65VZmjdGXzBK6ZDgfv1A283d8XK",
    "FHuecMbeC1PfjkW2JKyoicJAuiU7khgQT16QUB3Q1XdL",
    "H21Shi4iE7vgfjWEQNvzmpmBMJSaiZ17PYUcdNoAoKNc",
];

pub struct BrainDaemonMock {
    data: Arc<BrainData>,
}

impl BrainDaemonMock {
    pub fn new(data: Arc<BrainData>) -> Self {
        Self { data }
    }
}

pub struct BrainCliMock {
    data: Arc<BrainData>,
}

impl BrainCliMock {
    pub fn new(data: Arc<BrainData>) -> Self {
        Self { data }
    }
}

pub struct BrainAppCliMock {
    data: Arc<BrainData>,
}

impl BrainAppCliMock {
    pub fn new(data: Arc<BrainData>) -> Self {
        Self { data }
    }
}

pub struct BrainAppDaemonMock {
    data: Arc<BrainData>,
}

impl BrainAppDaemonMock {
    pub fn new(data: Arc<BrainData>) -> Self {
        Self { data }
    }
}

#[tonic::async_trait]
impl BrainVmDaemon for BrainDaemonMock {
    type RegisterVmNodeStream = Pin<Box<dyn Stream<Item = Result<VmContract, Status>> + Send>>;
    async fn register_vm_node(
        &self,
        req: Request<RegisterVmNodeReq>,
    ) -> Result<Response<Self::RegisterVmNodeStream>, Status> {
        let req = check_sig_from_req(req)?;
        info!("Starting registration process for {:?}", req);
        let node = crate::data::VmNode {
            public_key: req.node_pubkey.clone(),
            operator_wallet: req.operator_wallet,
            country: req.country,
            region: req.region,
            city: req.city,
            ip: req.main_ip,
            price: req.price,
            ..Default::default()
        };
        self.data.register_node(node);

        info!("Sending existing contracts to {}", req.node_pubkey);
        let contracts = self.data.find_vm_contracts_by_node(&req.node_pubkey);
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for contract in contracts {
                let _ = tx.send(Ok(contract.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::RegisterVmNodeStream
        ))
    }

    type BrainMessagesStream = Pin<Box<dyn Stream<Item = Result<BrainVmMessage, Status>> + Send>>;
    async fn brain_messages(
        &self,
        req: Request<DaemonStreamAuth>,
    ) -> Result<Response<Self::BrainMessagesStream>, Status> {
        let auth = req.into_inner();
        let pubkey = auth.pubkey.clone();
        check_sig_from_parts(
            &pubkey,
            &auth.timestamp,
            &format!("{:?}", auth.contracts),
            &auth.signature,
        )?;
        info!("Daemon {} connected to receive brain messages", pubkey);
        let (tx, rx) = mpsc::channel(6);
        self.data.add_daemon_tx(&pubkey, tx);
        let output_stream = ReceiverStream::new(rx).map(|msg| Ok(msg));
        Ok(Response::new(
            Box::pin(output_stream) as Self::BrainMessagesStream
        ))
    }

    async fn daemon_messages(
        &self,
        req: Request<Streaming<VmDaemonMessage>>,
    ) -> Result<Response<Empty>, Status> {
        let mut req_stream = req.into_inner();
        let pubkey: String;
        if let Some(Ok(msg)) = req_stream.next().await {
            log::debug!(
                "demon_messages received the following auth message: {:?}",
                msg.msg
            );
            if let Some(vm_daemon_message::Msg::Auth(auth)) = msg.msg {
                pubkey = auth.pubkey.clone();
                check_sig_from_parts(
                    &pubkey,
                    &auth.timestamp,
                    &format!("{:?}", auth.contracts),
                    &auth.signature,
                )?;
            } else {
                return Err(Status::unauthenticated(
                    "Could not authenticate the daemon: could not extract auth signature",
                ));
            }
        } else {
            return Err(Status::unauthenticated("Could not authenticate the daemon"));
        }

        // info!("Received a message from daemon {pubkey}: {daemon_message:?}");
        while let Some(daemon_message) = req_stream.next().await {
            match daemon_message {
                Ok(msg) => match msg.msg {
                    Some(vm_daemon_message::Msg::NewVmResp(new_vm_resp)) => {
                        self.data.submit_newvm_resp(new_vm_resp).await;
                    }
                    Some(vm_daemon_message::Msg::UpdateVmResp(update_vm_resp)) => {
                        self.data.submit_updatevm_resp(update_vm_resp).await;
                    }
                    Some(vm_daemon_message::Msg::VmNodeResources(node_resources)) => {
                        self.data.submit_node_resources(node_resources);
                    }
                    _ => {}
                },
                Err(e) => {
                    log::warn!("Daemon disconnected: {e:?}");
                    self.data.del_daemon_tx(&pubkey);
                }
            }
        }
        Ok(Response::new(Empty {}))
    }
}

#[tonic::async_trait]
impl BrainCli for BrainCliMock {
    async fn get_balance(&self, req: Request<Pubkey>) -> Result<Response<AccountBalance>, Status> {
        let req = check_sig_from_req(req)?;
        Ok(Response::new(self.data.get_balance(&req.pubkey).into()))
    }

    async fn new_vm(&self, req: Request<NewVmReq>) -> Result<Response<NewVmResp>, Status> {
        let req = check_sig_from_req(req)?;
        info!("New VM requested via CLI: {req:?}");
        if self
            .data
            .is_user_banned_by_node(&req.admin_pubkey, &req.node_pubkey)
        {
            return Err(Status::permission_denied(
                "This operator banned you. What did you do?",
            ));
        }
        let admin_pubkey = req.admin_pubkey.clone();
        let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
        self.data.submit_newvm_req(req, oneshot_tx).await;
        match oneshot_rx.await {
            Ok(response) => {
                info!("Sending VM confirmation to {admin_pubkey}: {response:?}");
                Ok(Response::new(response))
            }
            Err(e) => {
                log::error!("Something weird happened. Reached error {e:?}");
                Err(Status::unknown(
                    "Request failed due to unknown error. Please try again or contact the DeTEE devs team.",
                ))
            }
        }
    }

    async fn update_vm(&self, req: Request<UpdateVmReq>) -> Result<Response<UpdateVmResp>, Status> {
        let req = check_sig_from_req(req)?;
        info!("Update VM requested via CLI: {req:?}");
        let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
        self.data.submit_updatevm_req(req, oneshot_tx).await;
        match oneshot_rx.await {
            Ok(response) => {
                info!("Sending UpdateVMResp: {response:?}");
                Ok(Response::new(response))
            }
            Err(e) => Err(Status::unknown(format!(
                "Update VM request failed due to error: {e}"
            ))),
        }
    }

    async fn extend_vm(&self, req: Request<ExtendVmReq>) -> Result<Response<Empty>, Status> {
        let req = check_sig_from_req(req)?;
        match self
            .data
            .extend_vm_contract_time(&req.uuid, &req.admin_pubkey, req.locked_nano)
        {
            Ok(()) => Ok(Response::new(Empty {})),
            Err(e) => Err(Status::unknown(format!("Could not extend contract: {e}"))),
        }
    }

    async fn delete_vm(&self, req: Request<DeleteVmReq>) -> Result<Response<Empty>, Status> {
        let req = check_sig_from_req(req)?;
        match self.data.delete_vm(req).await {
            Ok(()) => Ok(Response::new(Empty {})),
            Err(e) => Err(Status::not_found(e.to_string())),
        }
    }

    async fn report_node(&self, req: Request<ReportNodeReq>) -> Result<Response<Empty>, Status> {
        let req = check_sig_from_req(req)?;
        match self.data.find_contract_by_uuid(&req.contract) {
            Ok(contract)
                if contract.admin_pubkey == req.admin_pubkey
                    && contract.node_pubkey == req.node_pubkey =>
            {
                ()
            }
            _ => return Err(Status::unauthenticated("No contract found by this ID.")),
        };
        self.data
            .report_node(req.admin_pubkey, &req.node_pubkey, req.reason);
        Ok(Response::new(Empty {}))
    }

    type ListVmContractsStream = Pin<Box<dyn Stream<Item = Result<VmContract, Status>> + Send>>;
    async fn list_vm_contracts(
        &self,
        req: Request<ListVmContractsReq>,
    ) -> Result<Response<Self::ListVmContractsStream>, Status> {
        let req = check_sig_from_req(req)?;
        info!(
            "CLI {} requested ListVMVmContractsStream. As operator: {}",
            req.wallet, req.as_operator
        );
        let mut contracts = Vec::new();
        if !req.uuid.is_empty() {
            if let Ok(specific_contract) = self.data.find_contract_by_uuid(&req.uuid) {
                if specific_contract.admin_pubkey == req.wallet {
                    contracts.push(specific_contract);
                }
                // TODO: allow operator to inspect contracts
            }
        } else {
            if req.as_operator {
                contracts.append(&mut self.data.find_vm_contracts_by_operator(&req.wallet));
            } else {
                contracts.append(&mut self.data.find_vm_contracts_by_admin(&req.wallet));
            }
        }
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for contract in contracts {
                let _ = tx.send(Ok(contract.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListVmContractsStream
        ))
    }

    type ListVmNodesStream = Pin<Box<dyn Stream<Item = Result<VmNodeListResp, Status>> + Send>>;
    async fn list_vm_nodes(
        &self,
        req: Request<VmNodeFilters>,
    ) -> Result<Response<Self::ListVmNodesStream>, tonic::Status> {
        let req = check_sig_from_req(req)?;
        info!("CLI requested ListVmNodesStream: {req:?}");
        let nodes = self.data.find_vm_nodes_by_filters(&req);
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for node in nodes {
                let _ = tx.send(Ok(node.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListVmNodesStream
        ))
    }

    async fn get_one_vm_node(
        &self,
        req: Request<VmNodeFilters>,
    ) -> Result<Response<VmNodeListResp>, Status> {
        let req = check_sig_from_req(req)?;
        info!("Unknown CLI requested ListVmNodesStream: {req:?}");
        match self.data.get_one_node_by_filters(&req) {
            Some(node) => Ok(Response::new(node.into())),
            None => Err(Status::not_found(
                "Could not find any node based on your search criteria",
            )),
        }
    }

    async fn register_operator(
        &self,
        req: Request<RegOperatorReq>,
    ) -> Result<Response<Empty>, Status> {
        let req = check_sig_from_req(req)?;
        info!("Regitering new operator: {req:?}");
        match self.data.register_operator(req) {
            Ok(()) => Ok(Response::new(Empty {})),
            Err(e) => Err(Status::failed_precondition(e.to_string())),
        }
    }

    async fn kick_contract(&self, req: Request<KickReq>) -> Result<Response<KickResp>, Status> {
        let req = check_sig_from_req(req)?;
        match self
            .data
            .kick_contract(&req.operator_wallet, &req.contract_uuid, &req.reason)
            .await
        {
            Ok(nano_lp) => Ok(Response::new(KickResp { nano_lp })),
            Err(e) => Err(Status::permission_denied(e.to_string())),
        }
    }

    async fn ban_user(&self, req: Request<BanUserReq>) -> Result<Response<Empty>, Status> {
        let req = check_sig_from_req(req)?;
        self.data.ban_user(&req.operator_wallet, &req.user_wallet);
        Ok(Response::new(Empty {}))
    }

    type ListOperatorsStream =
        Pin<Box<dyn Stream<Item = Result<ListOperatorsResp, Status>> + Send>>;
    async fn list_operators(
        &self,
        req: Request<Empty>,
    ) -> Result<Response<Self::ListOperatorsStream>, Status> {
        let _ = check_sig_from_req(req)?;
        let operators = self.data.list_operators();
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for op in operators {
                let _ = tx.send(Ok(op.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListOperatorsStream
        ))
    }

    async fn inspect_operator(
        &self,
        req: Request<Pubkey>,
    ) -> Result<Response<InspectOperatorResp>, Status> {
        match self.data.inspect_operator(&req.into_inner().pubkey) {
            Some(op) => Ok(Response::new(op.into())),
            None => Err(Status::not_found(
                "The wallet you specified is not an operator",
            )),
        }
    }

    async fn airdrop(&self, req: Request<AirdropReq>) -> Result<Response<Empty>, Status> {
        check_admin_key(&req)?;
        let req = check_sig_from_req(req)?;
        self.data.give_airdrop(&req.pubkey, req.tokens);
        Ok(Response::new(Empty {}))
    }

    async fn slash(&self, req: Request<SlashReq>) -> Result<Response<Empty>, Status> {
        check_admin_key(&req)?;
        let req = check_sig_from_req(req)?;
        self.data.slash_account(&req.pubkey, req.tokens);
        Ok(Response::new(Empty {}))
    }

    type ListAllVmContractsStream = Pin<Box<dyn Stream<Item = Result<VmContract, Status>> + Send>>;
    async fn list_all_vm_contracts(
        &self,
        req: Request<Empty>,
    ) -> Result<Response<Self::ListVmContractsStream>, Status> {
        check_admin_key(&req)?;
        let _ = check_sig_from_req(req)?;
        let contracts = self.data.list_all_contracts();
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for contract in contracts {
                let _ = tx.send(Ok(contract.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListVmContractsStream
        ))
    }

    type ListAccountsStream = Pin<Box<dyn Stream<Item = Result<Account, Status>> + Send>>;
    async fn list_accounts(
        &self,
        req: Request<Empty>,
    ) -> Result<Response<Self::ListAccountsStream>, Status> {
        check_admin_key(&req)?;
        let _ = check_sig_from_req(req)?;
        let accounts = self.data.list_accounts();
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for account in accounts {
                let _ = tx.send(Ok(account.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListAccountsStream
        ))
    }
}

trait PubkeyGetter {
    fn get_pubkey(&self) -> Option<String>;
}

#[tonic::async_trait]
impl BrainAppCli for BrainAppCliMock {
    type ListAppContractsStream = Pin<Box<dyn Stream<Item = Result<AppContract, Status>> + Send>>;
    type ListAppNodesStream = Pin<Box<dyn Stream<Item = Result<AppNodeListResp, Status>> + Send>>;

    async fn deploy_app(
        &self,
        req: tonic::Request<NewAppReq>,
    ) -> Result<tonic::Response<NewAppRes>, Status> {
        let req_data = check_sig_from_req(req)?;
        log::info!("Creating new container: {req_data:?}");
        let admin_pubkey = req_data.admin_pubkey.clone();
        let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
        self.data.send_new_container_req(req_data, oneshot_tx).await;

        match oneshot_rx.await {
            Ok(response) => {
                info!("responding container confirmation to {admin_pubkey}: {response:?}");
                Ok(Response::new(response))
            }
            Err(e) => {
                log::error!("Something went wrong. Reached error {e:?}");
                Err(Status::unknown(
                    "Request failed due to unknown error. Please try again or contact the DeTEE devs team.",
                ))
            }
        }
    }

    async fn delete_app(
        &self,
        req: tonic::Request<DelAppReq>,
    ) -> Result<tonic::Response<detee_shared::sgx::pb::brain::Empty>, Status> {
        let req_data = check_sig_from_req(req)?;
        log::info!("deleting container: {}", req_data.uuid.clone());
        if let Err(er) = self.data.send_del_container_req(req_data).await {
            info!("Could not delete container: {er}");
            return Err(Status::not_found("Could not find container"));
        };

        Ok(Response::new(detee_shared::sgx::pb::brain::Empty {}))
    }

    async fn list_app_contracts(
        &self,
        req: tonic::Request<ListAppContractsReq>,
    ) -> Result<tonic::Response<Self::ListAppContractsStream>, Status> {
        let req_data = check_sig_from_req(req)?;
        let app_contracts = self
            .data
            .find_app_contracts_by_admin_pubkey(&req_data.admin_pubkey);

        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for contract in app_contracts {
                let _ = tx.send(contract.into()).await;
            }
        });
        let output_stream = ReceiverStream::new(rx).map(Ok);
        Ok(Response::new(
            Box::pin(output_stream) as Self::ListAppContractsStream
        ))
    }

    async fn list_app_nodes(
        &self,
        req: tonic::Request<AppNodeFilters>,
    ) -> Result<tonic::Response<Self::ListAppNodesStream>, Status> {
        let req = check_sig_from_req(req)?;
        info!("CLI requested ListAppNodes: {req:?}");
        let nodes = self.data.find_app_nodes_by_filters(&req);
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for node in nodes {
                let _ = tx.send(Ok(node.into())).await;
            }
        });
        let output_stream = ReceiverStream::new(rx);
        Ok(Response::new(Box::pin(output_stream)))
    }

    async fn get_one_app_node(
        &self,
        req: tonic::Request<AppNodeFilters>,
    ) -> Result<tonic::Response<AppNodeListResp>, Status> {
        let req = check_sig_from_req(req)?;
        info!("CLI requested GetOneAppNode: {req:?}");
        match self.data.get_one_app_node_by_filters(&req) {
            Some(node) => Ok(Response::new(node.into())),
            None => Err(Status::not_found(
                "Could not find any node based on your search criteria",
            )),
        }
    }
}

#[tonic::async_trait]
impl BrainAppDaemon for BrainAppDaemonMock {
    type RegisterAppNodeStream = Pin<Box<dyn Stream<Item = Result<AppContract, Status>> + Send>>;
    type BrainMessagesStream = Pin<Box<dyn Stream<Item = Result<BrainMessageApp, Status>> + Send>>;

    async fn register_app_node(
        &self,
        req: tonic::Request<RegisterAppNodeReq>,
    ) -> Result<tonic::Response<Self::RegisterAppNodeStream>, Status> {
        let req_data = check_sig_from_req(req)?;
        log::info!(
            "registering app node_key : {}, operator_key: {}",
            &req_data.node_pubkey,
            &req_data.operator_wallet
        );

        let app_node = crate::data::AppNode {
            node_pubkey: req_data.node_pubkey.clone(),
            operator_wallet: req_data.operator_wallet,
            ip: req_data.main_ip,
            city: req_data.city,
            region: req_data.region,
            country: req_data.country,
            ..Default::default()
        };
        self.data.register_app_node(app_node);

        log::info!("Sending existing contracts to {}", &req_data.node_pubkey);
        let app_contracts = self.data.find_app_contracts_by_node(&req_data.node_pubkey);
        let (tx, rx) = mpsc::channel(6);
        tokio::spawn(async move {
            for contract in app_contracts {
                let _ = tx.send(contract.into()).await;
            }
        });
        let output_stream = ReceiverStream::new(rx).map(Ok);
        Ok(Response::new(Box::pin(output_stream)))
    }

    async fn brain_messages(
        &self,
        req: tonic::Request<detee_shared::sgx::pb::brain::DaemonAuth>,
    ) -> Result<tonic::Response<Self::BrainMessagesStream>, Status> {
        let req_data = req.into_inner();
        let pubkey = req_data.pubkey.clone();
        check_sig_from_parts(
            &pubkey,
            &req_data.timestamp,
            &format!("{:?}", req_data.contracts),
            &req_data.signature,
        )?;

        info!(
            "Daemon {} connected to receive brain messages",
            req_data.pubkey
        );
        let (tx, rx) = mpsc::channel(6);
        self.data.add_app_daemon_tx(&req_data.pubkey, tx);
        let output_stream = ReceiverStream::new(rx).map(Ok);
        Ok(Response::new(
            Box::pin(output_stream) as Self::BrainMessagesStream
        ))
    }

    async fn daemon_messages(
        &self,
        req: tonic::Request<Streaming<DaemonMessageApp>>,
    ) -> Result<tonic::Response<detee_shared::sgx::pb::brain::Empty>, Status> {
        let mut req_stream = req.into_inner();
        let mut pubkey;

        if let Some(Ok(msg)) = req_stream.next().await {
            log::debug!(
                "demon_messages received the following auth message: {:?}",
                msg.msg
            );
            if let Some(detee_shared::sgx::pb::brain::daemon_message_app::Msg::Auth(auth)) = msg.msg
            {
                pubkey = auth.pubkey.clone();
                check_sig_from_parts(
                    &pubkey,
                    &auth.timestamp,
                    &format!("{:?}", auth.contracts),
                    &auth.signature,
                )?;
            } else {
                return Err(Status::unauthenticated(
                    "Could not authenticate the daemon: could not extract auth signature",
                ));
            }
        } else {
            return Err(Status::unauthenticated("Could not authenticate the daemon"));
        }

        while let Some(daemon_message) = req_stream.next().await {
            match daemon_message {
                Ok(msg) => match msg.msg {
                    Some(detee_shared::sgx::pb::brain::daemon_message_app::Msg::Auth(
                        daemon_auth,
                    )) => pubkey = daemon_auth.pubkey,
                    Some(detee_shared::sgx::pb::brain::daemon_message_app::Msg::NewAppRes(
                        new_app_res,
                    )) => self.data.send_new_container_resp(new_app_res).await,
                    Some(
                        detee_shared::sgx::pb::brain::daemon_message_app::Msg::AppNodeResources(
                            node_resource,
                        ),
                    ) => self.data.submit_app_node_resources(node_resource),
                    _ => {
                        dbg!("None");
                    }
                },
                Err(e) => {
                    log::warn!("Daemon disconnected: {e:?}");
                    self.data.del_app_daemon_tx(&pubkey);
                }
            }
            //
        }

        Ok(Response::new(detee_shared::sgx::pb::brain::Empty {}))
    }
}

macro_rules! impl_pubkey_getter {
    ($t:ty, $field:ident) => {
        impl PubkeyGetter for $t {
            fn get_pubkey(&self) -> Option<String> {
                Some(self.$field.clone())
            }
        }
    };
    ($t:ty) => {
        impl PubkeyGetter for $t {
            fn get_pubkey(&self) -> Option<String> {
                None
            }
        }
    };
}

impl_pubkey_getter!(Pubkey, pubkey);
impl_pubkey_getter!(NewVmReq, admin_pubkey);
impl_pubkey_getter!(DeleteVmReq, admin_pubkey);
impl_pubkey_getter!(UpdateVmReq, admin_pubkey);
impl_pubkey_getter!(ExtendVmReq, admin_pubkey);
impl_pubkey_getter!(ReportNodeReq, admin_pubkey);
impl_pubkey_getter!(ListVmContractsReq, wallet);
impl_pubkey_getter!(RegisterVmNodeReq, node_pubkey);
impl_pubkey_getter!(RegOperatorReq, pubkey);
impl_pubkey_getter!(KickReq, operator_wallet);
impl_pubkey_getter!(BanUserReq, operator_wallet);

impl_pubkey_getter!(VmNodeFilters);
impl_pubkey_getter!(Empty);
impl_pubkey_getter!(AirdropReq);
impl_pubkey_getter!(SlashReq);

impl_pubkey_getter!(NewAppReq, admin_pubkey);
impl_pubkey_getter!(DelAppReq, admin_pubkey);
impl_pubkey_getter!(ListAppContractsReq, admin_pubkey);

impl_pubkey_getter!(RegisterAppNodeReq);
impl_pubkey_getter!(AppNodeFilters);

fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) -> Result<T, Status> {
    let time = match req.metadata().get("timestamp") {
        Some(t) => t.clone(),
        None => return Err(Status::unauthenticated("Timestamp not found in metadata.")),
    };
    let time = time
        .to_str()
        .map_err(|_| Status::unauthenticated("Timestamp in metadata is not a string"))?;

    let now = chrono::Utc::now();
    let parsed_time = chrono::DateTime::parse_from_rfc3339(time)
        .map_err(|_| Status::unauthenticated("Coult not parse timestamp"))?;
    let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds();
    if seconds_elapsed > 4 || seconds_elapsed < -4 {
        return Err(Status::unauthenticated(format!(
            "Date is not within 4 sec of the time of the server: CLI {} vs Server {}",
            parsed_time, now
        )));
    }

    let signature = match req.metadata().get("request-signature") {
        Some(t) => t,
        None => return Err(Status::unauthenticated("signature not found in metadata.")),
    };
    let signature = bs58::decode(signature)
        .into_vec()
        .map_err(|_| Status::unauthenticated("signature is not a bs58 string"))?;
    let signature = ed25519_dalek::Signature::from_bytes(
        signature
            .as_slice()
            .try_into()
            .map_err(|_| Status::unauthenticated("could not parse ed25519 signature"))?,
    );

    let pubkey_value = match req.metadata().get("pubkey") {
        Some(p) => p.clone(),
        None => return Err(Status::unauthenticated("pubkey not found in metadata.")),
    };
    let pubkey = ed25519_dalek::VerifyingKey::from_bytes(
        &bs58::decode(&pubkey_value)
            .into_vec()
            .map_err(|_| Status::unauthenticated("pubkey is not a bs58 string"))?
            .try_into()
            .map_err(|_| Status::unauthenticated("pubkey does not have the correct size."))?,
    )
    .map_err(|_| Status::unauthenticated("could not parse ed25519 pubkey"))?;

    let req = req.into_inner();
    let message = format!("{time}{req:?}");
    use ed25519_dalek::Verifier;
    pubkey
        .verify(message.as_bytes(), &signature)
        .map_err(|_| Status::unauthenticated("the signature is not valid"))?;
    if let Some(req_pubkey) = req.get_pubkey() {
        if pubkey_value.to_str().unwrap().to_string() != req_pubkey {
            return Err(Status::unauthenticated(
                "pubkey of signature does not match pubkey of request",
            ));
        }
    }
    Ok(req)
}

fn check_sig_from_parts(pubkey: &str, time: &str, msg: &str, sig: &str) -> Result<(), Status> {
    let now = chrono::Utc::now();
    let parsed_time = chrono::DateTime::parse_from_rfc3339(time)
        .map_err(|_| Status::unauthenticated("Coult not parse timestamp"))?;
    let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds();
    if seconds_elapsed > 4 || seconds_elapsed < -4 {
        return Err(Status::unauthenticated(format!(
            "Date is not within 4 sec of the time of the server: CLI {} vs Server {}",
            parsed_time, now
        )));
    }

    let signature = bs58::decode(sig)
        .into_vec()
        .map_err(|_| Status::unauthenticated("signature is not a bs58 string"))?;
    let signature = ed25519_dalek::Signature::from_bytes(
        signature
            .as_slice()
            .try_into()
            .map_err(|_| Status::unauthenticated("could not parse ed25519 signature"))?,
    );

    let pubkey = ed25519_dalek::VerifyingKey::from_bytes(
        &bs58::decode(&pubkey)
            .into_vec()
            .map_err(|_| Status::unauthenticated("pubkey is not a bs58 string"))?
            .try_into()
            .map_err(|_| Status::unauthenticated("pubkey does not have the correct size."))?,
    )
    .map_err(|_| Status::unauthenticated("could not parse ed25519 pubkey"))?;

    let msg = time.to_string() + msg;
    use ed25519_dalek::Verifier;
    pubkey
        .verify(msg.as_bytes(), &signature)
        .map_err(|_| Status::unauthenticated("the signature is not valid"))?;

    Ok(())
}

fn check_admin_key<T>(req: &Request<T>) -> Result<(), Status> {
    let pubkey = match req.metadata().get("pubkey") {
        Some(p) => p.clone(),
        None => return Err(Status::unauthenticated("pubkey not found in metadata.")),
    };
    let pubkey = pubkey
        .to_str()
        .map_err(|_| Status::unauthenticated("could not parse pubkey metadata to str"))?;

    if !ADMIN_ACCOUNTS.contains(&pubkey) {
        return Err(Status::unauthenticated(
            "This operation is reserved to admin accounts",
        ));
    }

    Ok(())
}