use std::sync::Arc; use crate::db::{PgRepo, UDB}; use crate::error::Error; use serde::Deserialize; use u_lib::{api::retypes, messaging::Reportable, models::*, types::Id}; use warp::reject::not_found; use warp::Rejection; type EndpResult = Result; #[derive(Deserialize)] pub struct PayloadFlags { brief: BriefMode, } pub struct Endpoints; impl Endpoints { pub async fn get_agents(repo: Arc, id: Option) -> EndpResult { repo.interact(move |mut db| { Ok(match id { Some(id) => { if let Some(agent) = db.get_agent(id)? { vec![agent] } else { vec![] } } None => db.get_agents()?, }) }) .await .map_err(From::from) } pub async fn get_job( repo: Arc, id: Id, params: Option, ) -> EndpResult { let Some(mut job) = repo.interact(move |mut db| db.get_job(id)).await? else { return Err(not_found()) }; Ok(match params.map(|p| p.brief) { Some(BriefMode::Yes) => job, Some(BriefMode::Auto) | None => { if let Some(payload) = &mut job.payload { payload.maybe_join_payload().map_err(Error::from)?; } job } Some(BriefMode::No) => { if let Some(payload) = &mut job.payload { payload.join_payload().map_err(Error::from)?; } job } }) } pub async fn get_jobs(repo: Arc) -> EndpResult { repo.interact(move |mut db| db.get_jobs()) .await .map_err(From::from) } pub async fn get_assigned_jobs( repo: Arc, id: Option, ) -> EndpResult { repo.interact(move |mut db| db.get_assigned_jobs(id, false)) .await .map_err(From::from) } pub async fn get_payloads(repo: Arc) -> EndpResult { repo.interact(move |mut db| db.get_payloads()) .await .map_err(From::from) } pub async fn get_payload( repo: Arc, id: Id, params: Option, ) -> EndpResult { let Some(mut payload) = repo.interact(move |mut db| db.get_payload(id)).await? else { return Err(not_found()) }; Ok(match params.map(|p| p.brief) { Some(BriefMode::Yes) => payload, None | Some(BriefMode::Auto) => { payload.maybe_join_payload().map_err(Error::from)?; payload } _ => { payload.join_payload().map_err(Error::from)?; payload } }) } pub async fn get_personal_jobs( repo: Arc, id: Id, ) -> EndpResult { repo.transaction(move |mut db| { let agent = db.get_agent(id)?; match agent { Some(mut agent) => { agent.touch(); db.upsert_agent(&agent)?; } None => { let mut new_agent = Agent::empty(); new_agent.id = id; db.upsert_agent(&new_agent)?; let job = db .find_job_by_alias("agent_hello")? .expect("agent_hello job not found"); db.set_jobs_for_agent(id, &[job.job.id])?; } } let assigned_jobs = db.get_assigned_jobs(Some(id), true)?; for job in &assigned_jobs { db.update_job_status(job.id, JobState::Running)?; } Ok(assigned_jobs .into_iter() .map(|j| AssignedJobById::from(&j)) .collect()) }) .await .map_err(From::from) } pub async fn upload_jobs(repo: Arc, msg: Vec) -> EndpResult { let jobs = msg .into_iter() .map(|mut job| { if let Some(payload) = &mut job.payload { payload.maybe_split_payload()?; } Ok(job) }) .collect::, Error>>()?; repo.interact(move |mut db| db.insert_jobs(&jobs)) .await .map_err(From::from) } pub async fn del(repo: Arc, id: Id) -> EndpResult<()> { repo.transaction(move |mut db| { [ UDB::del_agents, UDB::del_jobs, UDB::del_results, UDB::del_payloads, ] .iter() .map(|f| f(&mut db, &[id])) .collect::>() }) .await .map_err(From::from) } pub async fn set_jobs( repo: Arc, agent_id: Id, job_idents: Vec, ) -> EndpResult { repo.transaction(move |mut db| { let assigned_job_idents = job_idents .into_iter() .map(|ident| { Id::parse_str(&ident).or_else(|_| { let job_from_db = db.find_job_by_alias(&ident); match job_from_db { Ok(job) => match job { Some(j) => Ok(j.job.id), None => { Err(Error::ProcessingError(format!("unknown ident {ident}"))) } }, Err(e) => Err(e), } }) }) .collect::, Error>>()?; db.set_jobs_for_agent(agent_id, &assigned_job_idents)?; Ok(assigned_job_idents) }) .await .map_err(From::from) } pub async fn report( repo: Arc, msg: Vec, agent_id: Id, ) -> EndpResult { repo.transaction(move |mut db| { for entry in msg { match entry { Reportable::Assigned(mut result) => { let result_agent_id = &result.agent_id; if agent_id != *result_agent_id { warn!("Agent ids are not equal! actual id: {agent_id}, id from job: {result_agent_id}"); continue; } result.touch(); info!("agent {agent_id} updated job {}", result.id); match result.exec_type { JobType::Init => { result.state = JobState::Finished; match &result.result { Some(rbytes) => { let mut agent: Agent = match serde_json::from_slice(&rbytes) { Ok(a) => a, Err(e) => { error!("Error deserializing agent data from {agent_id}: {e}"); continue; } }; agent.state = AgentState::Active; db.upsert_agent(&agent)?; } None => error!("Empty agent data"), }}, JobType::Shell => { result.state = JobState::Finished }, JobType::Service => (), JobType::Terminate => todo!(), JobType::Update => todo!(), } db.update_result(&result)?; } Reportable::Error(e) => { error!("agent {agent_id} reported: {e}"); } Reportable::Dummy => (), }} Ok(()) }) .await .map_err(From::from) } pub async fn update_agent(repo: Arc, agent: Agent) -> EndpResult { repo.interact(move |mut db| db.upsert_agent(&agent)).await?; Ok(()) } pub async fn update_job(repo: Arc, mut job: Job) -> EndpResult { if let Some(payload) = &mut job.payload { payload.maybe_split_payload().map_err(Error::from)?; } repo.interact(move |mut db| db.update_job(&job.job)).await?; Ok(()) } pub async fn update_assigned_job( repo: Arc, assigned: AssignedJob, ) -> EndpResult { repo.interact(move |mut db| db.update_result(&assigned)) .await?; Ok(()) } }