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) -> UResult { 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::(&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(&self, url: impl AsRef) -> Result { self.req_with_payload(url, ()).await } async fn req_with_payload( &self, url: impl AsRef, payload: P, ) -> Result { 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::>(&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> { self.req(format!("get_personal_jobs/{}", url_param)).await } // send something to server pub async fn report(&self, payload: impl OneOrVec) -> Result<()> { self.req_with_payload("report", payload.into_vec()).await } // download file pub async fn dl(&self, file: String) -> Result> { self.req(format!("dl/{file}")).await } /// get exact job pub async fn get_job(&self, job: Uuid) -> Result> { self.req(format!("get_job/{job}")).await } /// get all available jobs pub async fn get_jobs(&self) -> Result> { self.req("get_jobs").await } } //##########// Admin area //##########// #[cfg(feature = "panel")] impl ClientHandler { /// agent listing pub async fn get_agents(&self, agent: Option) -> Result> { 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) -> 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>, ) -> Result> { self.req_with_payload("upload_jobs", payload.into_vec()) .await } /// delete something pub async fn del(&self, item: Uuid) -> Result { self.req(format!("del/{item}")).await } /// set jobs for any agent pub async fn set_jobs( &self, agent: Uuid, job_idents: impl OneOrVec, ) -> Result> { 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) -> Result> { self.req(format!("get_agent_jobs/{}", opt_to_string(agent))) .await } pub async fn ping(&self) -> Result<()> { self.req("ping").await } }