unki/lib/u_lib/src/api.rs

use std::collections::HashMap;
use std::fmt::Debug;
use std::net::SocketAddr;

use anyhow::{Context, Result};
use reqwest::{header, header::HeaderMap, Certificate, Client, Identity, Method, Url};
use serde::de::DeserializeOwned;
use serde_json::{from_str, Value};
use uuid::Uuid;

use crate::{
    config::{get_self_uid, MASTER_PORT},
    conv::opt_to_string,
    messaging::{self, AsMsg, BaseMessage},
    misc::OneOrVec,
    models::{self},
    UError, UResult,
};

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

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

impl ClientHandler {
    pub async fn new(server: &str, password: Option<String>) -> UResult<Self> {
        let identity = Identity::from_pkcs12_der(AGENT_IDENTITY, "").unwrap();
        let mut default_headers =
            HashMap::from([(header::USER_AGENT, get_self_uid().hyphenated().to_string())]);

        if let Some(pwd) = password {
            default_headers.insert(header::AUTHORIZATION, format!("Bearer {pwd}"));
        }

        let dns_response = Client::new()
            .request(
                Method::GET,
                format!("https://1.1.1.1/dns-query?name={server}&type=A"),
            )
            .header(header::ACCEPT, "application/dns-json")
            .send()
            .await?
            .text()
            .await?;

        let client = {
            let client = Client::builder()
                .identity(identity)
                .default_headers(HeaderMap::try_from(&default_headers).unwrap())
                .add_root_certificate(Certificate::from_pem(ROOT_CA_CERT).unwrap());

            match from_str::<Value>(&dns_response).unwrap()["Answer"]
                .get(0)
                .and_then(|a| a.get("data"))
            {
                Some(ip) => {
                    let raw_addr = format!("{}:{MASTER_PORT}", ip.as_str().unwrap());
                    let addr: SocketAddr = raw_addr.parse().unwrap();
                    client.resolve(server, addr)
                }
                None => client,
            }
        }
        .build()
        .unwrap();

        Ok(Self {
            client,
            base_url: Url::parse(&format!("https://{}:{}", server, MASTER_PORT)).unwrap(),
        })
    }

    async fn req<R: AsMsg + DeserializeOwned + Default>(&self, url: impl AsRef<str>) -> Result<R> {
        self.req_with_payload(url, ()).await
    }

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

        let response = request
            .send()
            .await
            .context("error while sending request")?;
        let content_len = response.content_length();
        let is_success = match response.error_for_status_ref() {
            Ok(_) => Ok(()),
            Err(e) => Err(UError::from(e)),
        };
        let resp = response.text().await.context("resp")?;

        let result = match is_success {
            Ok(_) => from_str::<BaseMessage<R>>(&resp)
                .map(|msg| msg.into_inner())
                .or_else(|e| match content_len {
                    Some(0) => Ok(Default::default()),
                    _ => Err(UError::NetError(e.to_string(), resp)),
                }),
            Err(UError::NetError(err, _)) => Err(UError::NetError(err, resp)),
            _ => unreachable!(),
        }
        .map_err(From::from);

        debug!("url = {}, resp = {:?}", url.as_ref(), result);

        result
    }

    // get jobs for client
    pub async fn get_personal_jobs(&self, url_param: Uuid) -> Result<Vec<models::AssignedJobById>> {
        self.req(format!("get_personal_jobs/{}", url_param)).await
    }

    // send something to server
    pub async fn report(&self, payload: impl OneOrVec<messaging::Reportable>) -> Result<()> {
        self.req_with_payload("report", payload.into_vec()).await
    }

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

    /// get exact job
    pub async fn get_job(&self, job: Uuid) -> Result<models::FatJobMeta<true>> {
        self.req(format!("get_job/{job}")).await
    }

    /// get all available jobs
    pub async fn get_jobs(&self) -> Result<Vec<models::ThinJobMeta>> {
        self.req("get_jobs").await
    }
}

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

    /// update agent
    pub async fn update_agent(&self, agent: models::Agent) -> Result<()> {
        self.req_with_payload("update_agent", agent).await
    }

    /// update job
    pub async fn update_job(&self, job: models::FatJobMeta<true>) -> Result<()> {
        self.req_with_payload("update_job", job).await
    }

    /// update result
    pub async fn update_result(&self, result: models::AssignedJob) -> Result<()> {
        self.req_with_payload("update_result", result).await
    }

    /// create and upload job
    pub async fn upload_jobs(
        &self,
        payload: impl OneOrVec<models::FatJobMeta<true>>,
    ) -> Result<Vec<Uuid>> {
        self.req_with_payload("upload_jobs", payload.into_vec())
            .await
    }

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

    /// set jobs for any agent
    pub async fn set_jobs(
        &self,
        agent: Uuid,
        job_idents: impl OneOrVec<String>,
    ) -> Result<Vec<Uuid>> {
        self.req_with_payload(format!("set_jobs/{agent}"), job_idents.into_vec())
            .await
    }

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

    pub async fn ping(&self) -> Result<()> {
        self.req("ping").await
    }
}