VirtWebRemote/remote_backend/src/virtweb_client.rs

use crate::app_config::AppConfig;
use crate::utils::time;
use lazy_regex::regex;
use std::collections::HashMap;
use std::fmt::Display;
use std::str::FromStr;
use thiserror::Error;
use uuid::{Error, Uuid};

#[derive(Error, Debug)]
pub enum VirtWebClientError {
    #[error("Invalid status code from VirtWeb: {0}")]
    InvalidStatusCode(u16),
}

#[derive(Eq, PartialEq, Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct GroupID(String);

impl GroupID {
    pub fn route_vm_info(&self) -> String {
        format!("/api/group/{}/vm/info", self.0)
    }

    pub fn route_vm_state(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/state{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_start(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/start{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_shutdown(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/shutdown{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_suspend(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/suspend{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_resume(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/resume{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_kill(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/kill{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_reset(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/reset{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
    pub fn route_vm_screenshot(&self, vm: Option<VMUuid>) -> String {
        format!(
            "/api/group/{}/vm/screenshot{}",
            self.0,
            match vm {
                None => "".to_string(),
                Some(id) => format!("?vm_id={}", id.0),
            }
        )
    }
}

#[derive(Eq, PartialEq, Debug, Copy, Clone, serde::Serialize, serde::Deserialize, Hash)]
pub struct VMUuid(Uuid);

#[derive(Default, serde::Deserialize, serde::Serialize)]
pub struct TreatmentResult {
    ok: usize,
    failed: usize,
}

impl VMUuid {
    pub fn route_info(&self) -> String {
        format!("/api/vm/{}", self.0)
    }

    pub fn route_state(&self) -> String {
        format!("/api/vm/{}/state", self.0)
    }

    pub fn route_start(&self) -> String {
        format!("/api/vm/{}/start", self.0)
    }

    pub fn route_shutdown(&self) -> String {
        format!("/api/vm/{}/shutdown", self.0)
    }
    pub fn route_kill(&self) -> String {
        format!("/api/vm/{}/kill", self.0)
    }

    pub fn route_reset(&self) -> String {
        format!("/api/vm/{}/reset", self.0)
    }
    pub fn route_suspend(&self) -> String {
        format!("/api/vm/{}/suspend", self.0)
    }

    pub fn route_resume(&self) -> String {
        format!("/api/vm/{}/resume", self.0)
    }

    pub fn route_screenshot(&self) -> String {
        format!("/api/vm/{}/screenshot", self.0)
    }
}

impl FromStr for VMUuid {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(VMUuid(Uuid::from_str(s)?))
    }
}

#[derive(serde::Serialize, Debug)]
pub struct TokenClaims {
    pub sub: String,
    pub iat: u64,
    pub exp: u64,
    pub verb: String,
    pub path: String,
    pub nonce: String,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct VMInfo {
    pub uuid: VMUuid,
    pub name: String,
    pub description: Option<String>,
    pub architecture: String,
    pub memory: usize,
    pub number_vcpu: usize,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct VMCaps {
    pub can_get_state: bool,
    pub can_start: bool,
    pub can_shutdown: bool,
    pub can_kill: bool,
    pub can_reset: bool,
    pub can_suspend: bool,
    pub can_resume: bool,
    pub can_screenshot: bool,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct VMState {
    pub state: String,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct LoadAverage {
    one: f64,
    five: f64,
    fifteen: f64,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct SystemSystemInfo {
    physical_core_count: usize,
    uptime: usize,
    used_memory: usize,
    available_memory: usize,
    free_memory: usize,
    load_average: LoadAverage,
}

#[derive(serde::Deserialize, serde::Serialize, Debug)]
pub struct SystemInfo {
    system: SystemSystemInfo,
}

#[derive(serde::Deserialize, Debug)]
pub struct TokenRight {
    verb: String,
    path: String,
}

pub type TokenRights = Vec<TokenRight>;

#[derive(serde::Deserialize)]
pub struct TokenInfo {
    rights: TokenRights,
}

impl TokenInfo {
    /// Check whether a route is allowed or not
    pub fn is_route_allowed(&self, verb: &str, route: &str) -> bool {
        let search_route_split = route.split('/').collect::<Vec<_>>();

        for r in &self.rights {
            if r.verb != verb {
                continue;
            }

            let curr_route_split = r.path.split('/').collect::<Vec<_>>();

            if search_route_split.len() != curr_route_split.len() {
                continue;
            }

            if curr_route_split
                .iter()
                .zip(search_route_split.iter())
                .all(|(curr, search)| curr == &"*" || curr == search)
            {
                return true;
            }
        }

        false
    }

    /// List the groups with access
    pub fn list_groups(&self) -> Vec<GroupID> {
        self.rights
            .iter()
            .filter(|r| r.verb == "GET")
            .filter(|r| regex!("^/api/group/[^/]+/vm/info$").is_match(&r.path))
            .map(|r| GroupID(r.path.split("/").nth(3).unwrap().to_string()))
            .collect::<Vec<_>>()
    }

    /// List the virtual machines with access
    pub fn list_vm(&self) -> Vec<VMUuid> {
        self.rights
            .iter()
            .filter(|r| r.verb == "GET")
            .filter(|r| regex!("^/api/vm/[^/]+$").is_match(&r.path))
            .map(|r| VMUuid::from_str(r.path.rsplit_once('/').unwrap().1).unwrap())
            .collect::<Vec<_>>()
    }

    /// Check if system info can be retrived
    pub fn can_retrieve_system_info(&self) -> bool {
        self.is_route_allowed("GET", "/api/server/info")
    }
}

/// Perform a request on the API
async fn request<D: Display>(uri: D) -> anyhow::Result<reqwest::Response> {
    let url = format!("{}{}", AppConfig::get().virtweb_base_url, uri);
    log::debug!("Will query {uri}...");

    let uri = uri.to_string();
    let jwt = TokenClaims {
        sub: AppConfig::get().virtweb_token_id.to_string(),
        iat: time() - 60 * 2,
        exp: time() + 60 * 3,
        verb: "GET".to_string(),
        path: uri.split_once('?').map(|s| s.0).unwrap_or(&uri).to_string(),
        nonce: Uuid::new_v4().to_string(),
    };
    let jwt = AppConfig::get().token_private_key().sign_jwt(&jwt)?;

    let res = reqwest::Client::new()
        .get(url)
        .header("x-token-id", &AppConfig::get().virtweb_token_id)
        .header("x-token-content", jwt)
        .send()
        .await?;

    if !res.status().is_success() {
        return Err(VirtWebClientError::InvalidStatusCode(res.status().as_u16()).into());
    }

    Ok(res)
}

/// Perform a request on the API
async fn json_request<D: Display, E: serde::de::DeserializeOwned>(uri: D) -> anyhow::Result<E> {
    Ok(request(uri).await?.json().await?)
}

/// Get current token information
pub async fn get_token_info() -> anyhow::Result<TokenInfo> {
    let res: TokenInfo =
        json_request(format!("/api/token/{}", AppConfig::get().virtweb_token_id)).await?;

    Ok(res)
}

/// Get a vm information
pub async fn vm_info(id: VMUuid) -> anyhow::Result<VMInfo> {
    json_request(id.route_info()).await
}

/// Get a vm information
pub async fn vm_state(id: VMUuid) -> anyhow::Result<VMState> {
    json_request(id.route_state()).await
}

/// Start a vm
pub async fn vm_start(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_start()).await?;
    Ok(())
}

/// Shutdown a vm
pub async fn vm_shutdown(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_shutdown()).await?;
    Ok(())
}

/// Kill a vm
pub async fn vm_kill(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_kill()).await?;
    Ok(())
}

/// Reset a vm
pub async fn vm_reset(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_reset()).await?;
    Ok(())
}

/// Suspend a vm
pub async fn vm_suspend(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_suspend()).await?;
    Ok(())
}

/// Resume a vm
pub async fn vm_resume(id: VMUuid) -> anyhow::Result<()> {
    request(id.route_resume()).await?;
    Ok(())
}

/// Grab a screenshot of the VM
pub async fn vm_screenshot(id: VMUuid) -> anyhow::Result<Vec<u8>> {
    Ok(request(id.route_screenshot())
        .await?
        .bytes()
        .await?
        .to_vec())
}

/// Get the VM of a group
pub async fn group_vm_info(id: &GroupID) -> anyhow::Result<Vec<VMInfo>> {
    json_request(id.route_vm_info()).await
}

/// Get the state of one or all VMs of a group
pub async fn group_vm_state(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<HashMap<VMUuid, String>> {
    json_request(id.route_vm_state(vm_id)).await
}

/// Start one or all VMs of a group
pub async fn group_vm_start(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_start(vm_id)).await
}

/// Shutdown one or all VMs of a group
pub async fn group_vm_shutdown(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_shutdown(vm_id)).await
}

/// Kill one or all VMs of a group
pub async fn group_vm_kill(id: &GroupID, vm_id: Option<VMUuid>) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_kill(vm_id)).await
}

/// Reset one or all VMs of a group
pub async fn group_vm_reset(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_reset(vm_id)).await
}

/// Suspend one or all VMs of a group
pub async fn group_vm_suspend(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_suspend(vm_id)).await
}

/// Resume one or all VMs of a group
pub async fn group_vm_resume(
    id: &GroupID,
    vm_id: Option<VMUuid>,
) -> anyhow::Result<TreatmentResult> {
    json_request(id.route_vm_resume(vm_id)).await
}

/// Get the screenshot of one or all VMs of a group
pub async fn group_vm_screenshot(id: &GroupID, vm_id: Option<VMUuid>) -> anyhow::Result<Vec<u8>> {
    Ok(request(id.route_vm_screenshot(vm_id))
        .await?
        .bytes()
        .await?
        .to_vec())
}

/// Get current server information
pub async fn get_server_info() -> anyhow::Result<SystemInfo> {
    json_request("/api/server/info").await
}