brain-mock-clients-mocks/cli-mock/src/main.rs

#![allow(dead_code)]
pub mod brain {
    tonic::include_proto!("brain");
}

use anyhow::Result;
use brain::{
    brain_cli_service_client::BrainCliServiceClient, DeletedVmUpdate, ListVmContractsReq,
    NewVmRequest, NodeFilters, NodeListResp, VmContract, UpdateVmRequest,
};
use lazy_static::lazy_static;
use log::{debug, info, warn};
use rand::{distributions::Alphanumeric, thread_rng, Rng};

use tokio_stream::StreamExt;
use tonic::transport::Channel;

lazy_static! {
    static ref SECURE_PUBLIC_KEY: String = use_default_string();
}

fn use_default_string() -> String {
    "ThisIsMyEternalClient".to_string()
}

fn generate_random_string() -> String {
    let rng = thread_rng();
    rng.sample_iter(&Alphanumeric)
        .take(16) // Adjust the length as needed
        .map(char::from)
        .collect()
}

async fn get_node_list(mut client: BrainCliServiceClient<Channel>) -> Result<Vec<NodeListResp>> {
    info!("Getting nodes from brain...");
    let mut nodes = Vec::new();
    let mut grpc_stream = client
        .list_nodes(NodeFilters {
            free_ports: 0,
            offers_ipv4: true,
            offers_ipv6: true,
            vcpus: 0,
            memory_mb: 0,
            storage_gb: 0,
            country: String::new(),
        })
        .await?
        .into_inner();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(node) => {
                debug!("Received node from brain: {node:?}");
                nodes.push(node);
            }
            Err(e) => {
                warn!("Received error instead of node list: {e:?}");
            }
        }
    }
    info!("Brain terminated list_nodes stream.");
    Ok(nodes)
}

async fn submit_vm_request(
    mut client: BrainCliServiceClient<Channel>,
    node_pubkey: &str,
) -> Result<()> {
    let req = NewVmRequest {
        uuid: String::new(),
        admin_pubkey: SECURE_PUBLIC_KEY.clone(),
        node_pubkey: node_pubkey.to_string(),
        hostname: "cool-server".to_string(),
        vcpus: 4,
        memory_mb: 4000,
        disk_size_gb: 20,
        public_ipv4: true,
        public_ipv6: true,
        extra_ports: Vec::new(),
        kernel_url: "thisIsMyURL".to_string(),
        kernel_sha: "thisIsMySha".to_string(),
        dtrfs_url: "thisIsMyURL".to_string(),
        dtrfs_sha: "thisIsMySha".to_string(),
    };
    info!("Creating VM {req:?}");
    let result = client.create_vm_contract(req).await;
    match result {
        Ok(confirmation) => {
            let confirmation = confirmation.into_inner();
            if confirmation.error.is_empty() {
                info!("Got VM confirmation: {confirmation:?}");
            } else {
                warn!("Got VM confirmation error: {}", confirmation.error);
            };
        }
        Err(e) => log::error!("Could not create vm: {e:?}"),
    };
    Ok(())
}

async fn list_contracts(mut client: BrainCliServiceClient<Channel>) -> Result<Vec<VmContract>> {
    info!("Getting contracts from brain...");
    let mut contracts = Vec::new();
    let mut grpc_stream = client
        .list_vm_contracts(ListVmContractsReq {
            admin_pubkey: SECURE_PUBLIC_KEY.to_string(),
            node_pubkey: String::new(),
        })
        .await?
        .into_inner();
    while let Some(stream_update) = grpc_stream.next().await {
        match stream_update {
            Ok(node) => {
                debug!("Received contract from brain: {node:?}");
                contracts.push(node);
            }
            Err(e) => {
                warn!("Received error instead of contracts: {e:?}");
            }
        }
    }
    info!("Brain terminated list_contracts stream.");
    Ok(contracts)
}

async fn delete_vm(mut client: BrainCliServiceClient<Channel>, uuid: &str) -> Result<()> {
    let req = DeletedVmUpdate {
        uuid: uuid.to_string(),
    };
    info!("Creating VM {req:?}");
    let result = client.delete_vm(req).await;
    match result {
        Ok(confirmation) => info!("VM deleted: {confirmation:?}"),
        Err(e) => log::error!("Could not delete vm: {e:?}"),
    };
    Ok(())
}

async fn update_vm_request(
    mut client: BrainCliServiceClient<Channel>,
    node_pubkey: &str,
    uuid: &str, 
) -> Result<()> {
    let req = UpdateVmRequest {
        uuid: uuid.to_string(),
        node_pubkey: node_pubkey.to_string(),
        vcpus: 4,
        memory_mb: 4096,
        disk_size_gb: 40,
        kernel_url: "thisIsMyNewURL".to_string(),
        kernel_sha: "thisIsMyNewSha".to_string(),
        dtrfs_url: "thisIsMyNewURL".to_string(),
        dtrfs_sha: "thisIsMyNewSha".to_string(),
    };
    info!("Creating VM {req:?}");
    let result = client.update_vm(req).await;
    match result {
        Ok(confirmation) => {
            let confirmation = confirmation.into_inner();
            if confirmation.error.is_empty() {
                info!("Got VM confirmation: {confirmation:?}");
            } else {
                warn!("Got VM confirmation error: {}", confirmation.error);
            };
        }
        Err(e) => log::error!("Could not create vm: {e:?}"),
    };
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::builder()
        .filter_level(log::LevelFilter::Debug)
        .init();
    info!("Hello! My name is {}", SECURE_PUBLIC_KEY.clone());

    let client = BrainCliServiceClient::connect("http://[::1]:31337").await?;

    let nodes = get_node_list(client.clone()).await?;
    for node in nodes {
        if let Err(e) = submit_vm_request(client.clone(), &node.node_pubkey).await {
            log::error!(
                "Received error when creating VM on node {}: {e:?}",
                &node.node_pubkey
            );
        }
    }

    let contracts = list_contracts(client.clone()).await?;
    for contract in contracts {
        if let Err(e) = update_vm_request(client.clone(), &use_default_string(), &contract.uuid).await {
            log::error!("Received error when updating VM {}: {e:?}", &contract.uuid);
        }
    }

    if std::env::var("DELETE_VMS").is_err() {
        return Ok(());
    }

    let contracts = list_contracts(client.clone()).await?;
    for contract in contracts {
        if let Err(e) = delete_vm(client.clone(), &contract.uuid).await {
            log::error!("Received error when deleting VM {}: {e:?}", &contract.uuid);
        }
    }
    Ok(())
}