unki/lib/u_lib/src/api.rs

use std::collections::HashMap;

use crate::{
    config::MASTER_PORT,
    messaging::{self, AsMsg, BaseMessage, Empty},
    models::{self, Agent},
    utils::opt_to_string,
    UError, UResult,
};
use reqwest::{header::HeaderMap, Certificate, Client, Identity, Url};
use serde::de::DeserializeOwned;
use uuid::Uuid;

const AGENT_IDENTITY: &[u8] = include_bytes!("../../../certs/alice.p12");
const ROOT_CA_CERT: &[u8] = include_bytes!("../../../certs/ca.crt");

#[derive(Clone)]
pub struct ClientHandler {
    base_url: Url,
    client: Client,
}

impl ClientHandler {
    pub fn new(server: &str, password: Option<String>) -> Self {
        let identity = Identity::from_pkcs12_der(AGENT_IDENTITY, "").unwrap();
        let mut client = Client::builder().identity(identity);
        if let Some(pwd) = password {
            client = client.default_headers(
                HeaderMap::try_from(&HashMap::from([(
                    "Authorization".to_string(),
                    format!("Bearer {pwd}"),
                )]))
                .unwrap(),
            )
        }
        let client = client
            .add_root_certificate(Certificate::from_pem(ROOT_CA_CERT).unwrap())
            .build()
            .unwrap();
        Self {
            client,
            base_url: Url::parse(&format!("https://{}:{}", server, MASTER_PORT)).unwrap(),
        }
    }

    async fn _req<P: AsMsg, M: AsMsg + DeserializeOwned>(
        &self,
        url: impl AsRef<str>,
        payload: P,
    ) -> UResult<M> {
        let request = self
            .client
            .post(self.base_url.join(url.as_ref()).unwrap())
            .json(&payload.as_message());

        let response = request.send().await?;
        let is_success = match response.error_for_status_ref() {
            Ok(_) => Ok(()),
            Err(e) => Err(UError::from(e)),
        };
        let resp = response.text().await?;
        match is_success {
            Ok(_) => serde_json::from_str::<BaseMessage<M>>(&resp)
                .map(|msg| msg.into_inner())
                .or_else(|e| Err(UError::NetError(e.to_string(), resp.clone()))),
            Err(UError::NetError(err, _)) => Err(UError::NetError(err, resp)),
            _ => unreachable!(),
        }
    }

    // get jobs for client
    pub async fn get_personal_jobs(
        &self,
        url_param: Option<Uuid>,
    ) -> UResult<Vec<models::AssignedJob>> {
        self._req(
            format!("get_personal_jobs/{}", opt_to_string(url_param)),
            Empty,
        )
        .await
    }

    // send something to server
    pub async fn report(&self, payload: &[messaging::Reportable]) -> UResult<Empty> {
        self._req("report", payload).await
    }

    // download file
    pub async fn dl(&self, file: String) -> UResult<Vec<u8>> {
        self._req(format!("dl/{file}"), Empty).await
    }
}

//##########// Admin area //##########//
#[cfg(feature = "panel")]
impl ClientHandler {
    /// agent listing
    pub async fn get_agents(&self, agent: Option<Uuid>) -> UResult<Vec<models::Agent>> {
        self._req(format!("get_agents/{}", opt_to_string(agent)), Empty)
            .await
    }

    /// update something
    pub async fn update_item(&self, item: impl AsMsg) -> UResult<Empty> {
        self._req("update_item", item).await
    }

    /// get all available jobs
    pub async fn get_jobs(&self, job: Option<Uuid>) -> UResult<Vec<models::JobMeta>> {
        self._req(format!("get_jobs/{}", opt_to_string(job)), Empty)
            .await
    }

    /// create and upload job
    pub async fn upload_jobs(&self, payload: &[models::JobMeta]) -> UResult<Empty> {
        self._req("upload_jobs", payload).await
    }

    /// delete something
    pub async fn del(&self, item: Uuid) -> UResult<i32> {
        self._req(format!("del/{item}"), Empty).await
    }

    /// set jobs for any agent
    pub async fn set_jobs(&self, agent: Uuid, job_idents: &[String]) -> UResult<Vec<Uuid>> {
        self._req(format!("set_jobs/{agent}"), job_idents).await
    }

    /// get jobs for any agent
    pub async fn get_agent_jobs(&self, agent: Option<Uuid>) -> UResult<Vec<models::AssignedJob>> {
        self._req(format!("set_jobs/{}", opt_to_string(agent)), Empty)
            .await
    }
}