unki/lib/u_lib/src/api.rs

use std::fmt::Debug;
use std::net::SocketAddr;
use std::{collections::HashMap, time::Duration};

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

use crate::{
    config::{get_self_id, MASTER_PORT},
    conv::opt_to_string,
    messaging::{self, AsMsg},
    models::*,
    types::Id,
    UError, UResult,
};

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

pub mod api_types {
    use super::*;

    pub type GetPersonalJobs = Vec<AssignedJobById>;
    pub type Report = ();
    pub type GetJob = Job;
    pub type GetBriefJob = Job;
    pub type GetJobs = Vec<JobMeta>;
    pub type GetAgents = Vec<Agent>;
    pub type UpdateAgent = ();
    pub type UpdateJob = ();
    pub type UpdateResult = ();
    pub type UpdatePayload = ();
    pub type UploadJobs = ();
    pub type UploadPayloads = ();
    pub type Del = ();
    pub type SetJobs = ();
    pub type GetAgentJobs = Vec<AssignedJob>;
    pub type Ping = ();
    pub type GetPayloads = Vec<Payload>;
    pub type GetPayload = Payload;
}

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

impl HttpClient {
    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_id().hyphenated().to_string())]);

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

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

        async fn resolve(domain_name: &str) -> Result<SocketAddr> {
            let dns_response = Client::new()
                .request(
                    Method::GET,
                    format!("https://1.1.1.1/dns-query?name={domain_name}&type=A"),
                )
                .header(header::ACCEPT, "application/dns-json")
                .send()
                .await?
                .text()
                .await?;

            let ip = from_str::<Value>(&dns_response)?
                .get("Answer")
                .and_then(|a| a.get(0))
                .and_then(|a| a.get("data"))
                .map(|ip| ip.to_owned())
                .ok_or_else(|| anyhow!("can't extract dns answer"))?;

            let raw_addr = format!("{}:{MASTER_PORT}", ip.as_str().unwrap());
            Ok(raw_addr.parse().unwrap())
        }

        let client = match resolve(server).await {
            Ok(addr) => client.resolve(server, addr),
            Err(e) => {
                warn!("DNS error: {e}");
                client
            }
        }
        .build()
        .unwrap();

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

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

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

        let response = request
            .send()
            .await
            .context("error while sending request")?;

        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::<R>(&resp).map_err(|e| UError::DeserializeError(e.to_string(), resp))
            }
            Err(UError::NetError(err, url, _)) => Err(UError::NetError(err, url, resp)),
            _ => unreachable!(),
        }
        .map_err(From::from);

        debug!("url = {url}, response = {result:?}");

        result
    }

    /// get jobs for agent
    pub async fn get_personal_jobs(&self, agent_id: Id) -> Result<api_types::GetPersonalJobs> {
        self.req(format!("get_personal_jobs/{}", agent_id)).await
    }

    /// send something to server
    pub async fn report(
        &self,
        payload: impl IntoIterator<Item = messaging::Reportable>,
    ) -> Result<api_types::Report> {
        self.req_with_payload("report", &payload.into_iter().collect::<Vec<_>>())
            .await
    }

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

    /// get exact job
    pub async fn get_job(&self, job: Id, brief: Brief) -> Result<api_types::GetJob> {
        self.req(format!("get_job/{job}?brief={brief}")).await
    }

    pub async fn get_full_job(&self, job: Id) -> Result<api_types::GetJob> {
        self.get_job(job, Brief::No).await
    }

    pub async fn get_brief_job(&self, job: Id) -> Result<api_types::GetJob> {
        self.get_job(job, Brief::Yes).await
    }

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

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

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

    /// update job
    pub async fn update_job(&self, job: &JobMeta) -> Result<api_types::UpdateJob> {
        self.req_with_payload("update_job", job).await
    }

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

    pub async fn update_payload(&self, payload: &Payload) -> Result<api_types::UpdatePayload> {
        self.req_with_payload("update_payload", payload).await
    }

    /// create and upload job
    pub async fn upload_jobs(
        &self,
        jobs: impl IntoIterator<Item = &Job>,
    ) -> Result<api_types::UploadJobs> {
        self.req_with_payload("upload_jobs", &jobs.into_iter().collect::<Vec<_>>())
            .await
    }

    pub async fn upload_payload(&self, payload: &RawPayload) -> Result<api_types::UploadPayloads> {
        self.req_with_payload("upload_payload", payload).await
    }

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

    /// set jobs for any agent
    pub async fn assign_jobs(
        &self,
        assigned: impl IntoIterator<Item = &AssignedJobById>,
    ) -> Result<api_types::SetJobs> {
        self.req_with_payload(
            format!("assign_jobs"),
            &assigned.into_iter().collect::<Vec<_>>(),
        )
        .await
    }

    /// get jobs for any agent by job_id, agent_id or result_id
    pub async fn get_assigned_jobs(&self, id: Option<Id>) -> Result<api_types::GetAgentJobs> {
        self.req(format!("get_assigned_jobs/{}", opt_to_string(id)))
            .await
    }

    pub async fn get_payloads(&self) -> Result<api_types::GetPayloads> {
        self.req("get_payloads").await
    }

    pub async fn get_payload(
        &self,
        payload: impl AsRef<str>,
        brief: Brief,
    ) -> Result<api_types::GetPayload> {
        let payload = payload.as_ref();
        self.req(format!("get_payload/{payload}?brief={brief}"))
            .await
    }

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