use crate::db::UDB; use crate::error::Error; use u_lib::{ messaging::{AsMsg, BaseMessage, Reportable}, models::*, utils::OneOrVec, }; use uuid::Uuid; use warp::Rejection; type EndpResult = Result; pub struct Endpoints; impl Endpoints { pub async fn add_agent(msg: Agent) -> EndpResult<()> { UDB::lock_db().insert_agent(&msg).map_err(From::from) } pub async fn get_agents(uid: Option) -> EndpResult> { UDB::lock_db().get_agents(uid).map_err(From::from) } pub async fn get_jobs(uid: Option) -> EndpResult> { UDB::lock_db().get_jobs(uid).map_err(From::from) } pub async fn get_agent_jobs(uid: Option) -> EndpResult> { UDB::lock_db() .get_exact_jobs(uid, false) .map_err(From::from) } pub async fn get_personal_jobs(uid: Uuid) -> EndpResult> { let db = UDB::lock_db(); let mut agents = db.get_agents(Some(uid))?; if agents.is_empty() { let new_agent = Agent::with_id(uid); db.insert_agent(&new_agent)?; let job = db .find_job_by_alias("agent_hello")? .expect("agent_hello job not found"); db.set_jobs_for_agent(&uid, &[job.id])?; } else { let mut agent = agents.pop().unwrap(); agent.touch(); db.update_agent(&agent)?; } let result = db.get_exact_jobs(Some(uid), true)?; for j in result.iter() { db.update_job_status(j.id, JobState::Running)?; } Ok(result) } pub async fn upload_jobs(msg: BaseMessage<'static, Vec>) -> EndpResult> { UDB::lock_db() .insert_jobs(&msg.into_inner()) .map_err(From::from) } pub async fn del(uid: Uuid) -> EndpResult { let db = UDB::lock_db(); let del_fns = &[UDB::del_agents, UDB::del_jobs, UDB::del_results]; for del_fn in del_fns { let affected = del_fn(&db, &[uid])?; if affected > 0 { return Ok(affected); } } Ok(0) } pub async fn set_jobs( agent_uid: Uuid, msg: BaseMessage<'static, Vec>, ) -> EndpResult> { msg.into_inner() .into_iter() .map(|ident| { Uuid::parse_str(&ident).or_else(|_| { let job_from_db = UDB::lock_db().find_job_by_alias(&ident); match job_from_db { Ok(job) => match job { Some(j) => Ok(j.id), None => Err(Error::ProcessingError(format!("unknown ident {ident}"))), }, Err(e) => Err(e), } }) }) .collect::, Error>>() .and_then(|j| UDB::lock_db().set_jobs_for_agent(&agent_uid, &j)) .map_err(From::from) } pub async fn report + AsMsg + 'static>( msg: BaseMessage<'static, Data>, ) -> EndpResult<()> { let id = msg.id; for entry in msg.into_inner().into_vec() { match entry { Reportable::Assigned(mut result) => { let result_agent_id = &result.agent_id; if id != *result_agent_id { warn!("Ids are not equal! actual id: {id}, id from job: {result_agent_id}"); continue; } result.state = JobState::Finished; result.touch(); match result.exec_type { JobType::Init => match &result.result { Some(rbytes) => { let mut agent: Agent = match serde_json::from_slice(&rbytes) { Ok(a) => a, Err(e) => { warn!("Error deserializing agent from {id}: {e}"); continue; } }; agent.state = AgentState::Active; Self::add_agent(agent).await?; } None => warn!("Empty agent data"), }, JobType::Shell => (), JobType::Terminate => todo!(), JobType::Update => todo!(), } UDB::lock_db().update_result(&result)?; } Reportable::Error(e) => { warn!("{} reported an error: {}", id, e); } Reportable::Dummy => (), } } Ok(()) } pub async fn update_agent(agent: BaseMessage<'static, Agent>) -> EndpResult<()> { UDB::lock_db().update_agent(&agent.into_inner())?; Ok(()) } pub async fn update_job(job: BaseMessage<'static, JobMeta>) -> EndpResult<()> { UDB::lock_db().update_job(&job.into_inner())?; Ok(()) } pub async fn update_assigned_job( assigned: BaseMessage<'static, AssignedJob>, ) -> EndpResult<()> { UDB::lock_db().update_result(&assigned.into_inner())?; Ok(()) } pub async fn download(_file_uid: String) -> EndpResult> { todo!() } }