From 4b4be31e017288d394e363a64099fcefe4e69123 Mon Sep 17 00:00:00 2001 From: plazmoid Date: Fri, 19 Feb 2021 08:14:44 +0500 Subject: [PATCH] f u c k --- README.md | 7 +- bin/u_agent/src/main.rs | 20 ++- bin/u_panel/Cargo.toml | 2 +- bin/u_panel/src/main.rs | 123 +++++++++++-- bin/u_server/Cargo.toml | 2 +- bin/u_server/src/db.rs | 138 ++++++++++++--- bin/u_server/src/errors.rs | 17 +- bin/u_server/src/handlers.rs | 161 +++++++++++------- bin/u_server/src/main.rs | 63 ++++--- lib/u_lib/Cargo.toml | 2 + lib/u_lib/src/api.rs | 35 ++-- lib/u_lib/src/errors.rs | 5 +- lib/u_lib/src/lib.rs | 1 - lib/u_lib/src/messaging.rs | 31 +--- lib/u_lib/src/models/agent.rs | 19 ++- lib/u_lib/src/models/jobs.rs | 90 ++++++++-- lib/u_lib/src/models/mod.rs | 6 +- lib/u_lib/src/utils.rs | 16 ++ .../2020-10-24-111622_create_all/up.sql | 6 +- scripts/exec_bin.sh | 2 +- 20 files changed, 547 insertions(+), 199 deletions(-) diff --git a/README.md b/README.md index 424ad75..3ea0919 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ Ничто не должно нарушать работоспособность и коммуникацию агентов с сервером, поэтому подключение должно быть защищено от прослушки, модификации. -##Установка агента на устройство (u_run) +## Установка агента на устройство (u_run) Утилита u_run осуществляет первичную сборку инфы о платформе, скачивание агента, его установку и запуск. Также она @@ -14,7 +14,7 @@ Исполняемый код шифруется блочным шифром (непопулярным), ключ получает при запуске и подключении к серверу. -##Взаимодействие (u_agent) +## Взаимодействие (u_agent) Агент висит в памяти в виде демона/сервиса/модуля ядра, запуск производится во время старта системы. Раз в 5 секунд агент пингует сервер, показывая свою жизнеспособность, а также запрашивая выставленные инструкции: @@ -23,7 +23,8 @@ * отправить результаты текущих джоб -## Веб-интерфейс (u_panel) +## Управление (u_panel) Панель управления для обменистрирования. Представляет собой u_agent с веб-сервером, транслирующим команды u_server-у. Запускается на localhost +Может быть использован как консольная утилита \ No newline at end of file diff --git a/bin/u_agent/src/main.rs b/bin/u_agent/src/main.rs index e7ae405..f72b211 100644 --- a/bin/u_agent/src/main.rs +++ b/bin/u_agent/src/main.rs @@ -39,23 +39,27 @@ async fn main() { env_logger::init(); let arg_ip = env::args().nth(1); let instance = ClientHandler::new(arg_ip); - info!("Gathering info"); let cli_info = gather().await; info!("Connecting to the server"); retry_until_ok!(instance.init(&cli_info).await); info!("Instanciated! Running main loop"); - loop {/* - let jobs = retry_until_ok!(instance.get_jobs().await).unwrap(); + loop { + let jobs = retry_until_ok!(instance.get_agent_jobs(Some(&*UID)).await); if jobs.len() > 0 { + let job_uids: Vec = jobs.iter() + .map(|j| j.id.to_string()[..8].to_owned()) + .collect(); + info!("Fetched jobs: \n{}", job_uids.join("\n")); let result = build_jobs(jobs) .run_until_complete() - .await; + .await + .into_iter() + .map(|r| r.unwrap())//TODO: panic handler (send info on server) + .collect(); retry_until_ok!(instance.report( - result.into_iter().map(|r| r.unwrap()).collect() + &result ).await) - }*/ - let jobs = retry_until_ok!(instance.get_jobs(&*UID).await); - println!("{:?}", jobs); + } sleep(Duration::from_secs(5)).await; } } diff --git a/bin/u_panel/Cargo.toml b/bin/u_panel/Cargo.toml index 029c533..734d36c 100644 --- a/bin/u_panel/Cargo.toml +++ b/bin/u_panel/Cargo.toml @@ -11,6 +11,6 @@ tokio = { version = "1.2.0", features = ["macros", "rt-multi-thread", "process"] structopt = "0.3.21" log = "^0.4" env_logger = "0.7.1" -uuid = "0.8.1" +uuid = "0.6.5" reqwest = { version = "0.11", features = ["json"] } u_lib = { version = "*", path = "../../lib/u_lib" } diff --git a/bin/u_panel/src/main.rs b/bin/u_panel/src/main.rs index 824abd8..0fecbb8 100644 --- a/bin/u_panel/src/main.rs +++ b/bin/u_panel/src/main.rs @@ -1,28 +1,123 @@ use structopt::StructOpt; use u_lib::{ api::ClientHandler, + models::JobMeta }; -use std::env::args; +use std::path::PathBuf; +use uuid::Uuid; -struct Table; +const DELIM: &'static str = "*************\n"; + +#[derive(StructOpt, Debug)] +struct Args { + #[structopt(subcommand)] + cmd: Cmd +} + +#[derive(StructOpt, Debug)] +enum Cmd { + Agents(LD), + Jobs(JobALD), + Jobmap(JmALD) +} + +#[derive(StructOpt, Debug)] +enum JobALD { + Add { + #[structopt(long, parse(try_from_str = parse_uuid))] + agent: Option, + + #[structopt(subcommand)] + cmd: JobCmd + }, + #[structopt(flatten)] + LD(LD) +} + +#[derive(StructOpt, Debug)] +enum JobCmd { + #[structopt(external_subcommand)] + Cmd(Vec) +} + +#[derive(StructOpt, Debug)] +enum JmALD { + Add { + #[structopt(parse(try_from_str = parse_uuid))] + agent_uid: Uuid, + + #[structopt(parse(try_from_str = parse_uuid))] + job_uids: Vec + }, + List { + #[structopt(parse(try_from_str = parse_uuid))] + uid: Option, + + #[structopt(short, long)] + results: bool + }, + Delete { + #[structopt(parse(try_from_str = parse_uuid))] + uid: Uuid + } +} + +#[derive(StructOpt, Debug)] +enum LD { + List { + #[structopt(parse(try_from_str = parse_uuid))] + uid: Option, + }, + Delete { + #[structopt(parse(try_from_str = parse_uuid))] + uid: Uuid + } +} + +fn parse_uuid(src: &str) -> Result { + Uuid::parse_str(src).map_err(|e| e.to_string()) +} #[tokio::main] async fn main() -> Result<(), &'static str> { - let mut raw_args = args(); - let method = match raw_args.nth(1) { - Some(m) => m, - None => return Err("Method required") - }; + let args: Args = Args::from_args(); let cli_handler = ClientHandler::new(None) .password("123qwe".to_string()); - match method.as_str() { - "ls" => { - let result = cli_handler.get_agents().await; - for cli in result.iter() { - println!("{:?}", cli) + match args.cmd { + Cmd::Agents(action) => match action { + LD::List {uid} => cli_handler.get_agents(uid.as_ref()) + .await.unwrap().into_iter().for_each(|r| println!("{}{}", DELIM, r)), + LD::Delete {uid} => { + println!("{}", cli_handler.del(Some(&uid)).await.unwrap()); } }, - _ => return Err("Unknown method") - }; + Cmd::Jobs(action) => match action { + JobALD::Add {cmd: JobCmd::Cmd(cmd), agent} => { + let job = JobMeta::from_shell(cmd.join(" ")); + let job_uid = job.id; + cli_handler.upload_jobs(&vec![job]).await.unwrap(); + if agent.is_some() { + cli_handler.set_jobs(&vec![job_uid], agent.as_ref()).await.unwrap() + } + }, + JobALD::LD(LD::List {uid}) => cli_handler.get_jobs(uid.as_ref()) + .await.unwrap().into_iter().for_each(|r| println!("{}{}", DELIM, r)), + JobALD::LD(LD::Delete {uid}) => { + println!("{}", cli_handler.del(Some(&uid)).await.unwrap()) + } + } + Cmd::Jobmap(action) => match action { + JmALD::Add {agent_uid, job_uids} => cli_handler.set_jobs(&job_uids, Some(&agent_uid)) + .await.unwrap(), + JmALD::List {uid, results} => if results { + cli_handler.get_results(uid.as_ref()) + .await.unwrap().into_iter().for_each(|r| println!("{}{}", DELIM, r)) + } else { + cli_handler.get_agent_jobs(uid.as_ref()) + .await.unwrap().into_iter().for_each(|r| println!("{}{}", DELIM, r)) + }, + JmALD::Delete {uid} => println!("{}", cli_handler.del(Some(&uid)).await.unwrap()) + } + } Ok(()) } diff --git a/bin/u_server/Cargo.toml b/bin/u_server/Cargo.toml index 6018558..498adc0 100644 --- a/bin/u_server/Cargo.toml +++ b/bin/u_server/Cargo.toml @@ -8,7 +8,7 @@ version = "0.1.0" dotenv = "0.15.0" env_logger = "0.7.1" log = "0.4.11" -anyhow = "*" +thiserror = "*" warp = "0.2.4" uuid = { version = "0.6.5", features = ["serde", "v4"] } diff --git a/bin/u_server/src/db.rs b/bin/u_server/src/db.rs index 06da2b7..797d0e8 100644 --- a/bin/u_server/src/db.rs +++ b/bin/u_server/src/db.rs @@ -1,30 +1,32 @@ use diesel::{ pg::PgConnection, - prelude::* + prelude::*, + result::Error as DslError }; use dotenv::dotenv; use std::{ env, sync::{Arc, Mutex} }; - use crate::{ - errors::USrvResult + errors::* +}; +use u_lib::{ + models::* }; -use u_lib::models::*; use uuid::Uuid; pub type Storage = Arc>; pub struct UDB { - conn: PgConnection + pub conn: PgConnection } impl UDB { pub fn new() -> USrvResult { - dotenv()?; + dotenv().unwrap(); let db_path = env::var("DATABASE_URL").unwrap(); - let conn = PgConnection::establish(&db_path)?; + let conn = PgConnection::establish(&db_path).unwrap(); let instance = UDB { conn }; @@ -39,10 +41,16 @@ impl UDB { Ok(()) } - pub fn get_agents(&self) -> USrvResult> { + pub fn get_agents(&self, uid: Option) -> USrvResult> { use schema::agents; - let result = agents::table - .load::(&self.conn)?; + let result = if uid.is_some() { + agents::table + .filter(agents::id.eq(uid.unwrap())) + .load::(&self.conn)? + } else { + agents::table + .load::(&self.conn)? + }; Ok(result) } @@ -50,35 +58,120 @@ impl UDB { use schema::jobs; let result = if uid.is_some() { jobs::table - .filter(jobs::id.like(uid.unwrap())) + .filter(jobs::id.eq(uid.unwrap())) + .get_results::(&self.conn)? } else { jobs::table - } - .load::(&self.conn)?; + .load::(&self.conn)? + }; Ok(result) } - pub fn get_agent_jobs(&self, uid: Option) -> USrvResult> { + //TODO: belongs_to + pub fn get_agent_jobs(&self, uid: Option, personal: bool) -> USrvResult> { use schema::{results, jobs}; + let mut q = results::table + .inner_join(jobs::table) + .into_boxed(); + if uid.is_some() { + q = q.filter(results::agent_id.eq(uid.unwrap())) + } + if personal { + q = q.filter(results::state.eq(JobState::Queued)) + } + let result = q.select( + (jobs::alias, jobs::id, jobs::exec_type, jobs::platform, jobs::payload) + ) + .get_results::(&self.conn)?; + Ok(result) + } + + pub fn get_results(&self, uid: Option) -> USrvResult> { + use schema::results; let result = if uid.is_some() { - jobs::table - .filter(jobs::id.like(uid.unwrap())) + results::table + .filter(results::agent_id.eq(uid.unwrap())) + .or_filter(results::job_id.eq(uid.unwrap())) + .or_filter(results::id.eq(uid.unwrap())) + .load::(&self.conn)? } else { - jobs::table - } - .load::(&self.conn)?; + results::table + .load::(&self.conn)? + }; Ok(result) } - pub fn add_jobs(&self, jobs: &Vec) -> USrvResult<()> { + pub fn add_jobs(&self, job_metas: &Vec) -> USrvResult<()> { use schema::jobs; diesel::insert_into(jobs::table) - .values(jobs) + .values(job_metas) .execute(&self.conn)?; Ok(()) } -} + pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &Vec) -> USrvResult<()> { + use schema::{agents::dsl::agents, jobs::dsl::jobs, results}; + if let Err(DslError::NotFound) = agents.find(agent_uid).first::(&self.conn) { + return Err(USrvError::NotFound(agent_uid.to_string())) + } + let not_found_jobs = job_uids.iter().filter_map(|job_uid| { + if let Err(DslError::NotFound) = jobs.find(job_uid).first::(&self.conn) { + Some(job_uid.to_string()) + } else { None } + }).collect::>(); + if not_found_jobs.len() > 0 { + return Err(USrvError::NotFound(not_found_jobs.join(", "))); + } + let job_requests = job_uids.iter().map(|job_uid| { + JobResult { + job_id: *job_uid, + agent_id: *agent_uid, + ..Default::default() + } + }).collect::>(); + diesel::insert_into(results::table) + .values(&job_requests) + .execute(&self.conn)?; + Ok(()) + } + + pub fn del_jobs(&self, uids: &Vec) -> USrvResult { + use schema::jobs; + let mut affected = 0; + for &uid in uids { + let deleted = diesel::delete(jobs::table) + .filter(jobs::id.eq(uid)) + .execute(&self.conn)?; + affected += deleted; + } + Ok(affected) + } + + pub fn del_agents(&self, uids: &Vec) -> USrvResult { + use schema::agents; + let mut affected = 0; + for &uid in uids { + let deleted = diesel::delete(agents::table) + .filter(agents::id.eq(uid)) + .execute(&self.conn)?; + affected += deleted; + } + Ok(affected) + } + + pub fn del_results(&self, uids: &Vec) -> USrvResult { + use schema::results; + let mut affected = 0; + for &uid in uids { + let deleted = diesel::delete(results::table) + .filter(results::id.eq(uid)) + .execute(&self.conn)?; + affected += deleted; + } + Ok(affected) + } +} +/* #[cfg(test)] mod tests { use super::*; @@ -109,3 +202,4 @@ mod tests { ) } } +*/ \ No newline at end of file diff --git a/bin/u_server/src/errors.rs b/bin/u_server/src/errors.rs index c836d2f..64f51c3 100644 --- a/bin/u_server/src/errors.rs +++ b/bin/u_server/src/errors.rs @@ -1,3 +1,16 @@ -use anyhow::Result as AnyResult; +use thiserror::Error; +use diesel::result::Error as DslError; -pub type USrvResult = AnyResult; +pub type USrvResult = Result; + +#[derive(Error, Debug)] +pub enum USrvError { + #[error("{0} is not found")] + NotFound(String), + + #[error("Error processing {0}")] + ProcessingError(String), + + #[error(transparent)] + DBError(#[from] DslError) +} diff --git a/bin/u_server/src/handlers.rs b/bin/u_server/src/handlers.rs index ac4627b..0b912d2 100644 --- a/bin/u_server/src/handlers.rs +++ b/bin/u_server/src/handlers.rs @@ -1,16 +1,30 @@ +use std::fmt::Display; use u_lib::{ models::* }; use warp::{ Rejection, Reply, - reply::{with_status, WithStatus}, - http::StatusCode + reply::{with_status}, + http::{StatusCode, Response} }; use crate::db::{ - Storage + Storage, + UDB, }; use uuid::Uuid; +use diesel::SaveChangesDsl; +use crate::errors::USrvError; + +fn build_response(code: StatusCode, body: S) -> Result, Rejection> { + Ok(Response::builder() + .status(code) + .body(format!("{}", body)).unwrap()) +} + +fn build_empty_200() -> Result, Rejection> { + build_response(StatusCode::OK, "") +} pub async fn add_agent( msg: BaseMessage<'_, IAgent>, @@ -18,17 +32,19 @@ pub async fn add_agent( { match db.lock() .unwrap() - .new_agent(&msg.into_item()) { - Ok(_) => Ok(warp::reply()), - //Err(e) => Ok(with_status(s.to_string(), StatusCode::BAD_REQUEST)) TODO - Err(e) => Err(warp::reject()) + .new_agent(&msg.into_inner()) { + Ok(_) => build_empty_200(), + Err(e) => build_response(StatusCode::BAD_REQUEST, e) } } -pub async fn get_agents(db: Storage) -> Result { +pub async fn get_agents( + uid: Option, + db: Storage +) -> Result { match db.lock() .unwrap() - .get_agents() { + .get_agents(uid) { Ok(r) => Ok(warp::reply::json( &r.as_message() )), @@ -40,83 +56,110 @@ pub async fn get_jobs( uid: Option, db: Storage) -> Result { - + match db.lock() + .unwrap() + .get_jobs(uid) { + Ok(r) => Ok(warp::reply::json( + &r.as_message() + )), + Err(e) => Err(warp::reject()) + } } -pub async fn get_agent_jobs( +pub async fn get_results( uid: Option, db: Storage) -> Result { + match db.lock() + .unwrap() + .get_results(uid) { + Ok(r) => Ok(warp::reply::json( + &r.as_message() + )), + Err(e) => Err(warp::reject()) + } +} +pub async fn get_agent_jobs( + uid: Option, + db: Storage, + personal: bool) -> Result +{ + match db.lock() + .unwrap() + .get_agent_jobs(uid, personal) { + Ok(r) => Ok(warp::reply::json( + &r.as_message() + )), + Err(e) => Err(warp::reject()) + } } pub async fn upload_jobs( msg: BaseMessage<'_, Vec>, db: Storage) -> Result { - -} -/* -pub async fn report( - msg: Payload>, - db: Storage) -> Result -{ - let results = msg.item.into_inner(); - let mut storage = db.results().await; - results.into_iter().for_each(|new_result| { - match storage.get_mut(&new_result.id) { - Some(v) => v.push(new_result), - None => storage.insert(new_result.id, vec![new_result]) - } - }); - Ok(with_status(warp::reply(), StatusCode::OK)) + match db.lock() + .unwrap() + .add_jobs(&msg.into_inner()) { + Ok(_) => build_empty_200(), + Err(e) => build_response(StatusCode::BAD_REQUEST, e) + } } -pub async fn get_job_results( +pub async fn del( uid: Uuid, db: Storage) -> Result { - let storage = db.results().await; - match storage.get(&uid) { - Some(v) => Ok(warp::reply::json( - &BaseMessage::new(v.clone()) - )), - None => Err(warp::reject()) + let db = db.lock().unwrap(); + let del_fns = &[ + UDB::del_agents, + UDB::del_jobs, + UDB::del_results, + ]; + for del_fn in del_fns { + let affected = del_fn(&db, &vec![uid]).unwrap(); + if affected > 0 { + return build_response(StatusCode::OK, affected) + } } + build_response(StatusCode::BAD_REQUEST, 0) } -pub async fn get_jobs( +pub async fn set_jobs( + agent_uid: Uuid, + msg: BaseMessage<'_, Vec>, db: Storage) -> Result { - let mut clients = db.clients().await; - let cli = clients.get_mut(&msg.id).unwrap(); - cli.jobs.iter_mut().for_each(|job: &mut JobMeta| { - if job.state == JobState::Queued { - job.state = JobState::Pending - } - }); - Ok(warp::reply::json( - &BaseMessage::new(cli.jobs.clone()) - )) + match db.lock() + .unwrap() + .set_jobs_for_agent(&agent_uid, &msg.into_inner()) { + Ok(_) => build_empty_200(), + Err(e) => build_response(StatusCode::BAD_REQUEST, dbg!(e)) + } } -pub async fn set_jobs( - uid: Option, - msg: BaseMessage<'_, ItemWrap>, +pub async fn report( + msg: BaseMessage<'_, Vec>, db: Storage) -> Result { - let mut clients = db.clients().await; - let cli = clients.get_mut(&uid.unwrap_or(msg.id)).unwrap(); - msg.item.0.into_iter().for_each(|(uuid, job)| { - match cli.jobs.get_mut(&uuid) { - Some(cli_job) => *cli_job = job, - None => cli.jobs.push(job) - }; - }); - Ok(()) - + let db = db.lock().unwrap(); + let id = msg.id; + let mut failed = vec![]; + for res in msg.into_inner() { + if id != res.agent_id { + continue + } + if let Err(e) = res.save_changes::(&db.conn).map_err(USrvError::from) { + failed.push(e.to_string()) + } + } + if failed.len() > 0 { + let err_msg = USrvError::ProcessingError(failed.join(", ")); + return build_response(StatusCode::BAD_REQUEST, err_msg); + } + build_empty_200() } -*/ pub async fn dummy() -> Result { Ok(String::from("ok")) diff --git a/bin/u_server/src/main.rs b/bin/u_server/src/main.rs index 8707671..8b93631 100644 --- a/bin/u_server/src/main.rs +++ b/bin/u_server/src/main.rs @@ -9,7 +9,6 @@ use warp::{ body }; -#[macro_use] extern crate log; extern crate env_logger; @@ -41,8 +40,11 @@ async fn main() { let base_db = UDB::new().unwrap(); let db = warp::any().map(move || base_db.clone()); + let infallible_none = |_| async { + Ok::<(Option,), std::convert::Infallible>((None,)) + }; - let new_client = warp::post() + let new_agent = warp::post() .and(warp::path(Paths::init)) .and(get_content::()) .and(db.clone()) @@ -50,6 +52,7 @@ async fn main() { let get_agents = warp::get() .and(warp::path(Paths::get_agents)) + .and(warp::path::param::().map(Some).or_else(infallible_none)) .and(db.clone()) .and_then(handlers::get_agents); @@ -61,59 +64,67 @@ async fn main() { let get_jobs = warp::get() .and(warp::path(Paths::get_jobs)) - .and(warp::path::param::>()) + .and(warp::path::param::().map(Some).or_else(infallible_none)) .and(db.clone()) .and_then(handlers::get_jobs); let get_agent_jobs = warp::get() .and(warp::path(Paths::get_agent_jobs)) - .and(warp::path::param::>()) + .and(warp::path::param::().map(Some).or_else(infallible_none)) .and(db.clone()) - .and_then(handlers::get_agent_jobs); + .and_then(|uid, db| handlers::get_agent_jobs(uid, db, false)); + + let get_personal_jobs = warp::get() + .and(warp::path(Paths::get_agent_jobs)) + .and(warp::path::param::().map(Some).or_else(infallible_none)) + .and(db.clone()) + .and_then(|uid, db| handlers::get_agent_jobs(uid, db, true)); let del = warp::get() .and(warp::path(Paths::del)) .and(warp::path::param::()) .and(db.clone()) .and_then(handlers::del); -/* + let set_jobs = warp::post() .and(warp::path(Paths::set_jobs)) - .and(warp::path::param::().map(Some)) - .and(get_content::()) - .and(db.clone()) - .and_then(handlers::set_jobs); - - let get_job_results = warp::get() - .and(warp::path(Paths::get_job_results)) .and(warp::path::param::()) + .and(get_content::>()) .and(db.clone()) - .and_then(handlers::get_job_results); + .and_then(handlers::set_jobs); let report = warp::post() .and(warp::path(Paths::report)) .and(get_content::>()) .and(db.clone()) .and_then(handlers::report); -*/ + + let get_results = warp::get() + .and(warp::path(Paths::get_results)) + .and(warp::path::param::().map(Some).or_else(infallible_none)) + .and(db.clone()) + .and_then(handlers::get_results); + let auth_token = warp::header::exact("authorization", "Bearer 123qwe"); - let agent_zone = new_client - .or(get_agent_jobs) -// .or(report) + let agent_zone = new_agent + .or(get_personal_jobs) + .or(report) ; let auth_zone = auth_token - .and(get_agents - .or(get_jobs) - .or(upload_jobs) - .or(del) -// .or(set_jobs) -// .or(get_job_results) + .and( + get_agents + .or(get_jobs) + .or(upload_jobs) + .or(del) + .or(set_jobs) + .or(get_agent_jobs) + .or(get_results) ); - let routes = auth_zone - .or(agent_zone); + let routes = agent_zone + .or(auth_zone); warp::serve(routes.with(warp::log("warp"))) .run(([0,0,0,0], MASTER_PORT)).await; } diff --git a/lib/u_lib/Cargo.toml b/lib/u_lib/Cargo.toml index 124d191..d9a0376 100644 --- a/lib/u_lib/Cargo.toml +++ b/lib/u_lib/Cargo.toml @@ -20,6 +20,8 @@ thiserror = "*" log = "*" env_logger = "0.8.3" diesel-derive-enum = { version = "1", features = ["postgres"] } +chrono = "0.4.19" +strum = { version = "0.20", features = ["derive"] } [dependencies.diesel] version = "1.4.5" diff --git a/lib/u_lib/src/api.rs b/lib/u_lib/src/api.rs index 0d18b8e..0ff71c5 100644 --- a/lib/u_lib/src/api.rs +++ b/lib/u_lib/src/api.rs @@ -5,7 +5,8 @@ use crate::{ MASTER_PORT, models::*, UResult, - UError + UError, + utils::opt_to_string }; use reqwest::{ Client, @@ -32,13 +33,16 @@ macro_rules! build_url_by_method { | instance: &ClientHandler $(, param: &$param_type)? - $(, url: &$url_param)? + $(, url: Option<&$url_param>)? | { let request = ClientHandler::build_post( instance, &format!("{}/{}", stringify!($path), - String::new() $(+ &(url as &$url_param).to_string())? + String::new() + $(+ + &opt_to_string(url as Option<&$url_param>) + )? ) ); request @@ -54,13 +58,16 @@ macro_rules! build_url_by_method { | instance: &ClientHandler $(, param: &$param_type)? - $(, url: &$url_param)? + $(, url: Option<&$url_param>)? | { let request = ClientHandler::build_get( instance, &format!("{}/{}", stringify!($path), - String::new() $(+ &(url as &$url_param).to_string())? + String::new() + $(+ + &opt_to_string(url as Option<&$url_param>) + )? ) ); request @@ -83,18 +90,18 @@ macro_rules! build_handler { pub async fn $path( &self $(, param: &$param_type)? - $(, url_param: &$url_param)? + $(, url_param: Option<&$url_param>)? ) -> UResult<$result> { let request = $crate::build_url_by_method!( $method $path, pname = $($param_name)?, ptype = $($param_type)?, urlparam = $($url_param)? - )(self $(, param as &$param_type)? $(, url_param as &$url_param)?); + )(self $(, param as &$param_type)? $(, url_param as Option<&$url_param>)?); let response = request.send().await?; match response.error_for_status() { - Ok(r) => r.json::>() + Ok(r) => Ok(r.json::>() .await - .map_err(|e| UError::from(e)) - .map(|msg| msg.into_item()), + .map(|msg| msg.into_inner()) + .unwrap_or_default()), Err(e) => Err(UError::from(e)) } } @@ -154,22 +161,22 @@ impl ClientHandler { // A - admin only ////////////////// // client listing (A) -build_handler!(GET get_agents() -> Vec); +build_handler!(GET get_agents/Uuid() -> Vec); // get jobs for client (agent_id=Uuid) build_handler!(GET get_agent_jobs/Uuid() -> Vec); // get all available jobs (A) build_handler!(GET get_jobs/Uuid() -> Vec); // add client to server's db -build_handler!(POST init(IAgent) -> String); +build_handler!(POST init(IAgent) -> ()); // create and upload job (A) build_handler!(POST upload_jobs(Vec) -> ()); // delete something (A) -build_handler!(GET del/Uuid() -> ()); +build_handler!(GET del/Uuid() -> String); // set jobs for client (A) // POST /set_jobs/Uuid json: Vec build_handler!(POST set_jobs/Uuid(Vec) -> ()); // get results (A) // GET /get_job_results/Uuid -build_handler!(GET get_job_results/Uuid() -> Vec); +build_handler!(GET get_results/Uuid() -> Vec); // report job result build_handler!(POST report(Vec) -> ()); diff --git a/lib/u_lib/src/errors.rs b/lib/u_lib/src/errors.rs index ba4442f..9103bae 100644 --- a/lib/u_lib/src/errors.rs +++ b/lib/u_lib/src/errors.rs @@ -29,7 +29,10 @@ pub enum UError { InsuitablePlatform(String, String), #[error("Task {0} doesn't exist")] - NoTask(Uuid) + NoTask(Uuid), + + #[error("Error opening {0}: {1}")] + FilesystemError(String, String) } impl From for UError { diff --git a/lib/u_lib/src/lib.rs b/lib/u_lib/src/lib.rs index 8cf8352..e7a91b6 100644 --- a/lib/u_lib/src/lib.rs +++ b/lib/u_lib/src/lib.rs @@ -21,6 +21,5 @@ extern crate lazy_static; #[macro_use] extern crate diesel; -#[macro_use] extern crate log; extern crate env_logger; \ No newline at end of file diff --git a/lib/u_lib/src/messaging.rs b/lib/u_lib/src/messaging.rs index 6197d2b..e026739 100644 --- a/lib/u_lib/src/messaging.rs +++ b/lib/u_lib/src/messaging.rs @@ -22,39 +22,22 @@ pub trait ToMsg: Clone { pub struct BaseMessage<'cow, I> where I: ToMsg { pub id: Uuid, - item: Cow<'cow, I> + inner: Cow<'cow, I> } impl<'cow, I> BaseMessage<'cow, I> where I: ToMsg { - pub fn new(item: C) -> Self + pub fn new(inner: C) -> Self where C: Into> { - let Moo(item) = item.into(); + let Moo(inner) = inner.into(); Self { id: UID.clone(), - item + inner } } - pub fn into_item(self) -> I { - self.item.into_owned() + pub fn into_inner(self) -> I { + self.inner.into_owned() } -} - -/* -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_create_message_owned() { - let item = String::from("QWEDSA"); - let msg_raw = BaseMessage { - id: *UID, - item: Cow::Owned(item.clone()) - }; - let msg = BaseMessage::new(item); - assert_eq!(msg_raw.item, msg.item); - } -}*/ \ No newline at end of file +} \ No newline at end of file diff --git a/lib/u_lib/src/models/agent.rs b/lib/u_lib/src/models/agent.rs index 6a038e6..599aa55 100644 --- a/lib/u_lib/src/models/agent.rs +++ b/lib/u_lib/src/models/agent.rs @@ -3,6 +3,7 @@ use serde::{ Deserialize }; use std::time::SystemTime; +use std::fmt; use diesel::{ Queryable, Identifiable, @@ -12,7 +13,7 @@ use diesel::{ use crate::{ models::*, UID, - utils::vec_to_string, + utils::*, models::schema::*, }; @@ -36,6 +37,22 @@ pub struct Agent { pub username: String } +impl fmt::Display for Agent { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let mut out = format!("Agent {}", self.id); + if self.alias.is_some() { + out += &format!(" ({})", self.alias.as_ref().unwrap()) + } + out += &format!("\nHostname: {}", self.hostname); + out += &format!("\nIs root: {}", self.is_root); + out += &format!("\nRoot allowed: {}", self.is_root_allowed); + out += &format!("\nLast active: {}", systime_to_string(&self.last_active)); + out += &format!("\nPlatform: {}", self.platform); + out += &format!("\nUsername: {}", self.username); + write!(f, "{}", out) + } +} + #[derive(Clone, Debug, Serialize, Deserialize, Insertable)] #[table_name = "agents"] pub struct IAgent { diff --git a/lib/u_lib/src/models/jobs.rs b/lib/u_lib/src/models/jobs.rs index 0ab2931..f77144f 100644 --- a/lib/u_lib/src/models/jobs.rs +++ b/lib/u_lib/src/models/jobs.rs @@ -1,7 +1,11 @@ use std::{ time::{SystemTime, Duration}, thread, - cmp::PartialEq + cmp::PartialEq, + fmt, + string::ToString, + path::PathBuf, + fs }; use serde::{ Serialize, @@ -11,6 +15,7 @@ use uuid::Uuid; use guess_host_triple::guess_host_triple; use tokio::process::Command; use crate::{ + utils::systime_to_string, models::schema::*, UError, UResult, @@ -23,8 +28,10 @@ use diesel_derive_enum::DbEnum; use diesel::{ Queryable, Identifiable, - Insertable + Insertable, + query_builder::AsChangeset }; +use strum::Display; #[derive(Serialize, Deserialize, Clone, Debug)] @@ -42,7 +49,7 @@ pub enum JobSchedule { //TODO: Scheduled } -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum, Display)] #[PgType = "JobState"] #[DieselType = "Jobstate"] pub enum JobState { @@ -52,9 +59,9 @@ pub enum JobState { Finished, } -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum)] +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum, Display)] #[PgType = "JobType"] -#[DieselType = "Job_type"] +#[DieselType = "Jobtype"] pub enum JobType { Manage, Shell, @@ -151,7 +158,7 @@ impl JobOutput { Insertable )] #[table_name = "jobs"] -pub struct JobMeta { +pub struct JobMeta { // TODO: shell cmd how to exec payload pub alias: String, pub id: Uuid, pub exec_type: JobType, @@ -160,17 +167,49 @@ pub struct JobMeta { pub payload: Option>, } +impl fmt::Display for JobMeta { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let mut out = format!("Job {}", self.id); + out += &format!(" ({})", self.alias); + out += &format!("\nExecutable type: {}", self.exec_type); + out += &format!("\nPlatform: {}", self.platform); + if self.exec_type == JobType::Shell && self.payload.is_some() { + out += &format!("\nPayload: {}", String::from_utf8_lossy(self.payload.as_ref().unwrap())); + } + write!(f, "{}", out) + } +} + impl JobMeta { pub fn from_shell>(shell_cmd: S) -> Self { let shell_cmd = shell_cmd.into(); let job_name = shell_cmd.split(" ").nth(0).unwrap(); Self { - id: Uuid::new_v4(), alias: job_name.to_string(), + payload: Some(shell_cmd.into_bytes()), + ..Default::default() + } + } +/* + pub fn from_file(path: PathBuf) -> UResult { + let data = fs::read(path) + .map_err(|e| UError::FilesystemError( + path.to_string_lossy().to_string(), + e.to_string() + ))?; + let filename = path.file_name().unwrap().to_str().unwrap(); + + }*/ +} + +impl Default for JobMeta { + fn default() -> Self { + Self { + id: Uuid::new_v4(), + alias: String::new(), exec_type: JobType::Shell, - //schedule: JobSchedule::Once, platform: guess_host_triple().unwrap_or("unknown").to_string(), - payload: Some(shell_cmd.into_bytes()) + payload: None } } } @@ -183,17 +222,38 @@ impl JobMeta { Debug, Queryable, Identifiable, - Insertable + Insertable, + AsChangeset )] #[table_name = "results"] pub struct JobResult { pub agent_id: Uuid, + pub created: SystemTime, pub id: Uuid, pub job_id: Uuid, pub result: Option>, pub state: JobState, pub retcode: Option, - pub ts: SystemTime, + pub updated: SystemTime, +} + +impl fmt::Display for JobResult { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let mut out = format!("Result {}", self.id); + out += &format!("\nAgent {}", self.agent_id); + out += &format!("\nJob: {}", self.job_id); + out += &format!("\nUpdated: {}", systime_to_string(&self.updated)); + out += &format!("\nState: {}", self.state); + if self.state == JobState::Finished { + if self.retcode.is_some() { + out += &format!("\nReturn code: {}", self.retcode.unwrap()); + } + if self.result.is_some() { + out += &format!("\nResult: {}", String::from_utf8_lossy(self.result.as_ref().unwrap())); + } + } + write!(f, "{}", out) + } } impl JobResult { @@ -211,12 +271,13 @@ impl Default for JobResult { fn default() -> Self { Self { agent_id: Uuid::nil(), + created: SystemTime::now(), id: Uuid::new_v4(), job_id: Uuid::nil(), result: None, state: JobState::Queued, retcode: None, - ts: SystemTime::now() + updated: SystemTime::now() } } } @@ -254,7 +315,7 @@ impl Job { } None => unimplemented!() }; - let mut cmd_parts = str_payload // WRONG + let mut cmd_parts = str_payload // TODO: WRONG .split(" ") .map(String::from) .collect::>() @@ -285,7 +346,8 @@ impl Job { }; self.result.result = data; self.result.retcode = retcode; - self.result.ts = SystemTime::now(); + self.result.updated = SystemTime::now(); + self.result.state = JobState::Finished; }, _ => todo!() } diff --git a/lib/u_lib/src/models/mod.rs b/lib/u_lib/src/models/mod.rs index 828b59a..98bdcfd 100644 --- a/lib/u_lib/src/models/mod.rs +++ b/lib/u_lib/src/models/mod.rs @@ -9,12 +9,10 @@ pub use crate::{ }, messaging::*, }; - -use std::{ - borrow::Cow -}; use uuid::Uuid; +use std::borrow::Cow; +// with this macro, a type can be used as message (see api) macro_rules! to_message { ($($type:ty),+) => { $( diff --git a/lib/u_lib/src/utils.rs b/lib/u_lib/src/utils.rs index 7e2c361..33b2bbb 100644 --- a/lib/u_lib/src/utils.rs +++ b/lib/u_lib/src/utils.rs @@ -14,6 +14,11 @@ use nix::{ } }; use std::process::exit; +use std::time::SystemTime; +use chrono::{ + DateTime, + offset::Local +}; pub trait OneOrMany { fn into_vec(self) -> Vec; @@ -63,4 +68,15 @@ pub fn setsig(sig: Signal, hnd: SigHandler) { pub fn vec_to_string(v: &[u8]) -> String { String::from_utf8_lossy(v).to_string() +} + +pub fn opt_to_string(item: Option) -> String { + match item { + Some(s) => s.to_string(), + None => String::new() + } +} + +pub fn systime_to_string(time: &SystemTime) -> String { + DateTime::::from(*time).format("%d/%m/%Y %T").to_string() } \ No newline at end of file diff --git a/migrations/2020-10-24-111622_create_all/up.sql b/migrations/2020-10-24-111622_create_all/up.sql index f573623..471cc6e 100644 --- a/migrations/2020-10-24-111622_create_all/up.sql +++ b/migrations/2020-10-24-111622_create_all/up.sql @@ -33,7 +33,7 @@ CREATE TABLE IF NOT EXISTS ip_addrs ( ); CREATE TABLE IF NOT EXISTS jobs ( - alias TEXT + alias TEXT NOT NULL , id UUID NOT NULL DEFAULT uuid_generate_v4() -- Shell, Binary (with program download), -- Python (with program and python download if not exist), Management @@ -49,9 +49,9 @@ CREATE TABLE IF NOT EXISTS results ( , id UUID NOT NULL DEFAULT uuid_generate_v4() , job_id UUID NOT NULL , result BYTEA - , retcode INTEGER , state JobState NOT NULL DEFAULT 'queued' - , ts TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP + , retcode INTEGER + , updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP , FOREIGN KEY(agent_id) REFERENCES agents(id) , FOREIGN KEY(job_id) REFERENCES jobs(id) , PRIMARY KEY(id) diff --git a/scripts/exec_bin.sh b/scripts/exec_bin.sh index c324f6d..48620e5 100755 --- a/scripts/exec_bin.sh +++ b/scripts/exec_bin.sh @@ -4,5 +4,5 @@ source $(dirname $0)/rootdir.sh #set ROOTDIR BIN=$1 [[ $BIN == '' ]] && echo "Bin required" && exit 1 shift -export RUST_LOG=info +[[ $RUST_LOG == '' ]] && export RUST_LOG=info $ROOTDIR/target/x86_64-unknown-linux-gnu/debug/u_$BIN $@