diff --git a/Cargo.toml b/Cargo.toml index ece9d2c..8f5feed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ ] [workspace.dependencies] +anyhow = "1.0.58" reqwest = { version = "0.11", features = ["json"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" diff --git a/Makefile.toml b/Makefile.toml index 63e14e2..18100b2 100644 --- a/Makefile.toml +++ b/Makefile.toml @@ -48,8 +48,12 @@ dependencies = ["build_static_libs", "build_frontend"] command = "${CARGO}" args = ["build", "--target", "${TARGET}", "${@}"] +[tasks.cargo_update] +command = "${CARGO}" +args = ["update"] + [tasks.release_tasks] -condition = { env = { "PROFILE_OVERRIDE" = "release"} } +condition = { env = { PROFILE_OVERRIDE = "release"} } script = ''' BINS=$(ls ./target/${TARGET}/${PROFILE_OVERRIDE}/u_* -1 | grep -v ".d") echo "Stripping..." @@ -59,7 +63,7 @@ upx -9 $BINS ''' [tasks.build] -dependencies = ["cargo_build", "release_tasks"] +dependencies = ["cargo_update", "cargo_build", "release_tasks"] clear = true [tasks.run] @@ -70,6 +74,7 @@ command = "${CARGO}" args = ["test", "--target", "${TARGET}", "--lib", "--", "${@}"] [tasks.integration] +dependencies = ["cargo_update"] script = ''' [[ ! -d "./target/${TARGET}/${PROFILE_OVERRIDE}" ]] && echo 'No target folder. Build project first' && exit 1 cd ./integration diff --git a/bin/u_agent/src/lib.rs b/bin/u_agent/src/lib.rs index 1725e58..e9bc3e9 100644 --- a/bin/u_agent/src/lib.rs +++ b/bin/u_agent/src/lib.rs @@ -115,6 +115,6 @@ pub async fn run_forever() -> ! { // ErrChan::send(UError::Runtime(e.to_string()), "deeeemon").await // } } - info!("Startup"); + info!("Starting agent {}", get_self_uid()); agent_loop(client).await } diff --git a/bin/u_panel/Cargo.toml b/bin/u_panel/Cargo.toml index 618d1bc..e6705f6 100644 --- a/bin/u_panel/Cargo.toml +++ b/bin/u_panel/Cargo.toml @@ -9,7 +9,7 @@ edition = "2021" [dependencies] actix-cors = "0.6.1" actix-web = "4.1" -anyhow = "1.0.44" +anyhow = { workspace = true } futures-util = "0.3.21" mime_guess = "2.0.4" once_cell = "1.8.0" diff --git a/bin/u_panel/src/argparse.rs b/bin/u_panel/src/argparse.rs index ce810ba..80261f4 100644 --- a/bin/u_panel/src/argparse.rs +++ b/bin/u_panel/src/argparse.rs @@ -78,11 +78,11 @@ fn parse_uuid(src: &str) -> Result { pub async fn process_cmd(client: ClientHandler, args: Args) -> UResult { fn to_json(data: AnyResult) -> String { - let data = match data { + let result = match data { Ok(r) => PanelResult::Ok(r), Err(e) => PanelResult::Err(e.downcast().expect("unknown error type")), }; - serde_json::to_string(&data).unwrap() + serde_json::to_string(&result).unwrap() } Ok(match args.cmd { diff --git a/bin/u_server/Cargo.toml b/bin/u_server/Cargo.toml index b825bbf..474ecb1 100644 --- a/bin/u_server/Cargo.toml +++ b/bin/u_server/Cargo.toml @@ -5,6 +5,7 @@ name = "u_server" version = "0.1.0" [dependencies] +anyhow = { workspace = true } diesel = { version = "1.4.5", features = ["postgres", "uuid"] } hyper = "0.14" once_cell = "1.7.2" diff --git a/bin/u_server/src/db.rs b/bin/u_server/src/db.rs index 09dcad3..e72b2f9 100644 --- a/bin/u_server/src/db.rs +++ b/bin/u_server/src/db.rs @@ -1,4 +1,4 @@ -use crate::errors::{Error, SResult}; +use crate::error::Error as ServerError; use diesel::{pg::PgConnection, prelude::*, result::Error as DslError}; use once_cell::sync::OnceCell; use serde::Deserialize; @@ -9,9 +9,10 @@ use u_lib::{ }; use uuid::Uuid; +type Result = std::result::Result; + pub struct UDB { - pub conn: PgConnection, - __p: (), + conn: PgConnection, } static DB: OnceCell> = OnceCell::new(); @@ -34,7 +35,6 @@ impl UDB { ); let instance = UDB { conn: PgConnection::establish(&db_url).unwrap(), - __p: (), }; Mutex::new(instance) }) @@ -42,69 +42,89 @@ impl UDB { .unwrap() } - pub fn insert_jobs(&self, job_metas: &[JobMeta]) -> SResult<()> { + pub fn insert_jobs(&self, job_metas: &[JobMeta]) -> Result<()> { use schema::jobs; + diesel::insert_into(jobs::table) .values(job_metas) - .execute(&self.conn)?; + .execute(&self.conn) + .map_err(with_err_ctx("Can't insert jobs"))?; Ok(()) } - pub fn get_jobs(&self, uid: Option) -> SResult> { + pub fn get_jobs(&self, ouid: Option) -> Result> { use schema::jobs; - let result = if uid.is_some() { - jobs::table - .filter(jobs::id.eq(uid.unwrap())) - .get_results::(&self.conn)? - } else { - jobs::table.load::(&self.conn)? - }; - Ok(result) + + match ouid { + Some(uid) => jobs::table + .filter(jobs::id.eq(uid)) + .get_results::(&self.conn), + None => jobs::table.load::(&self.conn), + } + .map_err(with_err_ctx("Can't get exact jobs")) } - pub fn find_job_by_alias(&self, alias: &str) -> SResult { + pub fn find_job_by_alias(&self, alias: &str) -> Result> { use schema::jobs; + let result = jobs::table .filter(jobs::alias.eq(alias)) - .first::(&self.conn)?; + .first::(&self.conn) + .optional() + .map_err(with_err_ctx(format!("Can't find job by alias {alias}")))?; Ok(result) } - pub fn insert_agent(&self, agent: &Agent) -> SResult<()> { + pub fn insert_agent(&self, agent: &Agent) -> Result<()> { use schema::agents; + diesel::insert_into(agents::table) .values(agent) .on_conflict(agents::id) .do_update() .set(agent) - .execute(&self.conn)?; + .execute(&self.conn) + .map_err(with_err_ctx(format!("Can't insert agent {agent:x?}")))?; + Ok(()) + } + + pub fn insert_result(&self, result: &AssignedJob) -> Result<()> { + use schema::results; + + diesel::insert_into(results::table) + .values(result) + .execute(&self.conn) + .map_err(with_err_ctx(format!("Can't insert result {result:x?}")))?; Ok(()) } - pub fn get_agents(&self, uid: Option) -> SResult> { + pub fn get_agents(&self, ouid: Option) -> Result> { use schema::agents; - 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) + + match ouid { + Some(uid) => agents::table + .filter(agents::id.eq(uid)) + .load::(&self.conn), + None => agents::table.load::(&self.conn), + } + .map_err(with_err_ctx(format!("Can't get agent(s) {ouid:?}"))) } - pub fn update_job_status(&self, uid: Uuid, status: JobState) -> SResult<()> { + pub fn update_job_status(&self, uid: Uuid, status: JobState) -> Result<()> { use schema::results; + diesel::update(results::table) .filter(results::id.eq(uid)) .set(results::state.eq(status)) - .execute(&self.conn)?; + .execute(&self.conn) + .map_err(with_err_ctx(format!("Can't update status of job {uid}")))?; Ok(()) } //TODO: filters possibly could work in a wrong way, check - pub fn get_exact_jobs(&self, uid: Option, personal: bool) -> SResult> { + pub fn get_exact_jobs(&self, uid: Option, personal: bool) -> Result> { use schema::results; + let mut q = results::table.into_boxed(); /*if uid.is_some() { q = q.filter(results::agent_id.eq(uid.unwrap())) @@ -121,32 +141,19 @@ impl UDB { .or_filter(results::job_id.eq(uid.unwrap())) .or_filter(results::id.eq(uid.unwrap())) } - let result = q.load::(&self.conn)?; + let result = q + .load::(&self.conn) + .map_err(with_err_ctx("Can't get exact jobs"))?; Ok(result) } - pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &[Uuid]) -> SResult> { - use schema::{agents::dsl::agents, jobs::dsl::jobs, results}; - if let Err(DslError::NotFound) = agents.find(agent_uid).first::(&self.conn) { - return Err(Error::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.is_empty() { - return Err(Error::NotFound(not_found_jobs.join(", "))); - } + pub fn set_jobs_for_agent(&self, agent_uid: &Uuid, job_uids: &[Uuid]) -> Result> { + use schema::results; + let job_requests = job_uids .iter() .map(|job_uid| { - info!("set_jobs_for_agent: set {} for {}", job_uid, agent_uid); + debug!("set_jobs_for_agent: set {} for {}", job_uid, agent_uid); AssignedJob { job_id: *job_uid, agent_id: *agent_uid, @@ -154,46 +161,84 @@ impl UDB { } }) .collect::>(); + diesel::insert_into(results::table) .values(&job_requests) - .execute(&self.conn)?; - let assigned_uids = job_requests.iter().map(|aj| aj.id).collect(); - Ok(assigned_uids) + .execute(&self.conn) + .map_err(with_err_ctx(format!( + "Can't setup jobs {job_uids:?} for agent {agent_uid:?}" + )))?; + + Ok(job_requests.iter().map(|aj| aj.id).collect()) } - pub fn del_jobs(&self, uids: &[Uuid]) -> SResult { + pub fn del_jobs(&self, uids: &[Uuid]) -> Result { 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)?; + .execute(&self.conn) + .map_err(with_err_ctx("Can't delete jobs"))?; affected += deleted; } Ok(affected) } - pub fn del_results(&self, uids: &[Uuid]) -> SResult { + pub fn del_results(&self, uids: &[Uuid]) -> Result { 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)?; + .execute(&self.conn) + .map_err(with_err_ctx("Can't delete results"))?; affected += deleted; } Ok(affected) } - pub fn del_agents(&self, uids: &[Uuid]) -> SResult { + pub fn del_agents(&self, uids: &[Uuid]) -> Result { 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)?; + .execute(&self.conn) + .map_err(with_err_ctx("Can't delete agents"))?; affected += deleted; } Ok(affected) } + + pub fn update_agent(&self, agent: &Agent) -> Result<()> { + agent + .save_changes::(&self.conn) + .map_err(with_err_ctx(format!("Can't update agent {agent:x?}")))?; + Ok(()) + } + + pub fn update_job(&self, job: &JobMeta) -> Result<()> { + job.save_changes::(&self.conn) + .map_err(with_err_ctx(format!("Can't update job {job:x?}")))?; + Ok(()) + } + + pub fn update_result(&self, result: &AssignedJob) -> Result<()> { + debug!( + "updating result: id = {}, job_id = {}, agent_id = {}", + result.id, result.job_id, result.agent_id + ); + result + .save_changes::(&self.conn) + .map_err(with_err_ctx(format!("Can't update result {result:x?}")))?; + Ok(()) + } +} + +fn with_err_ctx(msg: impl AsRef) -> impl Fn(DslError) -> ServerError { + move |err| ServerError::DBErrorCtx(format!("{}, reason: {err}", msg.as_ref())) } diff --git a/bin/u_server/src/error.rs b/bin/u_server/src/error.rs new file mode 100644 index 0000000..6f07b72 --- /dev/null +++ b/bin/u_server/src/error.rs @@ -0,0 +1,59 @@ +use diesel::result::Error as DslError; +use thiserror::Error; +use warp::{ + http::StatusCode, + reject::Reject, + reply::{with_status, Response}, + Reply, +}; + +#[derive(Error, Debug)] +pub enum Error { + #[error("Error processing {0}")] + ProcessingError(String), + + #[error(transparent)] + DBError(#[from] DslError), + + #[error("DB error: {0}")] + DBErrorCtx(String), + + #[error("General error: {0}")] + Other(String), +} + +impl Reject for Error {} + +pub struct RejResponse { + message: String, + status: StatusCode, +} + +impl RejResponse { + pub fn not_found(msg: impl Into) -> Self { + Self { + message: msg.into(), + status: StatusCode::NOT_FOUND, + } + } + + pub fn bad_request(msg: impl Into) -> Self { + Self { + message: msg.into(), + status: StatusCode::BAD_REQUEST, + } + } + + pub fn internal() -> Self { + Self { + message: "INTERNAL_SERVER_ERROR".to_string(), + status: StatusCode::INTERNAL_SERVER_ERROR, + } + } +} + +impl Reply for RejResponse { + fn into_response(self) -> Response { + with_status(self.message, self.status).into_response() + } +} diff --git a/bin/u_server/src/errors.rs b/bin/u_server/src/errors.rs deleted file mode 100644 index c77472a..0000000 --- a/bin/u_server/src/errors.rs +++ /dev/null @@ -1,22 +0,0 @@ -use diesel::result::Error as DslError; -use thiserror::Error; -use warp::reject::Reject; - -pub type SResult = Result; - -#[derive(Error, Debug)] -pub enum Error { - #[error("{0} is not found")] - NotFound(String), - - #[error("Error processing {0}")] - ProcessingError(String), - - #[error(transparent)] - DBError(#[from] DslError), - - #[error("General error: {0}")] - Other(String), -} - -impl Reject for Error {} diff --git a/bin/u_server/src/handlers.rs b/bin/u_server/src/handlers.rs index 3b0db33..43139ab 100644 --- a/bin/u_server/src/handlers.rs +++ b/bin/u_server/src/handlers.rs @@ -1,8 +1,7 @@ use std::time::SystemTime; use crate::db::UDB; -use crate::errors::Error; -use diesel::SaveChangesDsl; +use crate::error::Error; use u_lib::{ messaging::{AsMsg, BaseMessage, Reportable}, models::*, @@ -39,7 +38,9 @@ impl Endpoints { if agents.is_empty() { let db = UDB::lock_db(); db.insert_agent(&Agent::with_id(uid.unwrap()))?; - let job = db.find_job_by_alias("agent_hello")?; + let job = db + .find_job_by_alias("agent_hello")? + .expect("agent_hello job not found"); db.set_jobs_for_agent(&uid.unwrap(), &[job.id])?; } let result = UDB::lock_db().get_exact_jobs(uid, true)?; @@ -76,8 +77,16 @@ impl Endpoints { msg.into_inner() .into_iter() .map(|ident| { - Uuid::parse_str(&ident) - .or_else(|_| UDB::lock_db().find_job_by_alias(&ident).map(|j| j.id)) + 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)) @@ -117,10 +126,7 @@ impl Endpoints { JobType::Terminate => todo!(), JobType::Update => todo!(), } - let db = UDB::lock_db(); - result - .save_changes::(&db.conn) - .map_err(Error::from)?; + UDB::lock_db().update_result(&result)?; } Reportable::Error(e) => { warn!("{} reported an error: {}", id, e); @@ -132,27 +138,19 @@ impl Endpoints { } pub async fn update_agent(agent: BaseMessage<'static, Agent>) -> EndpResult<()> { - agent - .into_inner() - .save_changes::(&UDB::lock_db().conn) - .map_err(Error::from)?; + UDB::lock_db().update_agent(&agent.into_inner())?; Ok(()) } pub async fn update_job(job: BaseMessage<'static, JobMeta>) -> EndpResult<()> { - job.into_inner() - .save_changes::(&UDB::lock_db().conn) - .map_err(Error::from)?; + UDB::lock_db().update_job(&job.into_inner())?; Ok(()) } pub async fn update_assigned_job( assigned: BaseMessage<'static, AssignedJob>, ) -> EndpResult<()> { - assigned - .into_inner() - .save_changes::(&UDB::lock_db().conn) - .map_err(Error::from)?; + UDB::lock_db().update_result(&assigned.into_inner())?; Ok(()) } diff --git a/bin/u_server/src/u_server.rs b/bin/u_server/src/u_server.rs index 8c11e0a..4417e8f 100644 --- a/bin/u_server/src/u_server.rs +++ b/bin/u_server/src/u_server.rs @@ -12,12 +12,12 @@ extern crate diesel; // in this block mod db; -mod errors; +mod error; mod handlers; -use errors::{Error, SResult}; +use error::{Error as ServerError, RejResponse}; use serde::{de::DeserializeOwned, Deserialize}; -use std::path::PathBuf; +use std::{convert::Infallible, path::PathBuf}; use u_lib::{ config::MASTER_PORT, logging::init_logger, @@ -28,7 +28,8 @@ use u_lib::{ use uuid::Uuid; use warp::{ body, - reply::{json, reply, Json}, + log::{custom, Info}, + reply::{json, reply, Json, Response}, Filter, Rejection, Reply, }; @@ -55,7 +56,7 @@ pub fn init_endpoints( auth_token: &str, ) -> impl Filter + Clone { let path = |p: &'static str| warp::post().and(warp::path(p)); - let infallible_none = |_| async { Ok::<(Option,), std::convert::Infallible>((None,)) }; + let infallible_none = |_| async { Ok::<_, Infallible>((None::,)) }; let get_agents = path("get_agents") .and( @@ -131,7 +132,7 @@ pub fn init_endpoints( .map(ok); let auth_token = format!("Bearer {auth_token}",).into_boxed_str(); - let auth_header = warp::header::exact("Authorization", Box::leak(auth_token)); + let auth_header = warp::header::exact("authorization", Box::leak(auth_token)); let auth_zone = (get_agents .or(get_jobs) @@ -150,32 +151,31 @@ pub fn init_endpoints( auth_zone.or(agent_zone) } -pub fn prefill_jobs() -> SResult<()> { +pub fn prefill_jobs() -> Result<(), ServerError> { let job_alias = "agent_hello"; - let if_job_exists = UDB::lock_db().find_job_by_alias(job_alias); - match if_job_exists { - Ok(_) => Ok(()), - Err(Error::DBError(diesel::result::Error::NotFound)) => { - let agent_hello = JobMeta::builder() - .with_type(JobType::Init) - .with_alias(job_alias) - .build() - .unwrap(); - UDB::lock_db().insert_jobs(&[agent_hello]) - } - Err(e) => Err(e), + let if_job_exists = UDB::lock_db().find_job_by_alias(job_alias)?; + if if_job_exists.is_none() { + let agent_hello = JobMeta::builder() + .with_type(JobType::Init) + .with_alias(job_alias) + .build() + .unwrap(); + UDB::lock_db().insert_jobs(&[agent_hello])? } + Ok(()) } -pub async fn serve() -> SResult<()> { +pub async fn serve() -> Result<(), ServerError> { init_logger(Some("u_server")); prefill_jobs()?; - let env = load_env::().map_err(|e| Error::Other(e.to_string()))?; - let routes = init_endpoints(&env.admin_auth_token); let certs_dir = PathBuf::from("certs"); + let env = load_env::().map_err(|e| ServerError::Other(e.to_string()))?; + let routes = init_endpoints(&env.admin_auth_token) + .recover(handle_rejection) + .with(custom(logger)); - warp::serve(routes.with(warp::log("warp"))) + warp::serve(routes) .tls() .cert_path(certs_dir.join("server.crt")) .key_path(certs_dir.join("server.key")) @@ -185,6 +185,32 @@ pub async fn serve() -> SResult<()> { Ok(()) } +async fn handle_rejection(rej: Rejection) -> Result { + let resp = if let Some(err) = rej.find::() { + error!("{:x?}", err); + RejResponse::bad_request(err.to_string()) + } else if rej.is_not_found() { + RejResponse::not_found("not found placeholder") + } else { + RejResponse::internal() + }; + Ok(resp.into_response()) +} + +fn logger(info: Info<'_>) { + info!(target: "warp", + "{raddr} {agent_uid} \"{path}\"", + raddr = info.remote_addr().unwrap_or(([0, 0, 0, 0], 0).into()), + path = info.path(), + agent_uid = info.user_agent() + .map(|uid: &str| uid.splitn(3, '-') + .take(2) + .collect::() + ) + .unwrap_or_else(|| "NO_AGENT".to_string()) + ); +} + fn ok(_: T) -> impl Reply { reply() } diff --git a/integration/docker-compose.yml b/integration/docker-compose.yml index d148d39..f988509 100644 --- a/integration/docker-compose.yml +++ b/integration/docker-compose.yml @@ -29,7 +29,7 @@ services: - ../.env - ../.env.private environment: - RUST_LOG: info + RUST_LOG: warp=info,u_server_lib=debug healthcheck: test: ss -tlpn | grep 63714 interval: 5s @@ -80,6 +80,7 @@ services: networks: - u_net volumes: + - ${HOME}/.cargo/registry/:/usr/local/cargo/registry/ - ../__Cargo_integration.toml:/tests/Cargo.toml - ./:/tests/integration/ - ../certs:/tests/certs diff --git a/integration/integration_tests.py b/integration/integration_tests.py index 43c54ff..e2b2028 100644 --- a/integration/integration_tests.py +++ b/integration/integration_tests.py @@ -59,7 +59,7 @@ def run_tests(): if not only_setup_cluster: CLUSTER.run('cargo test --test integration') except Exception as e: - # CLUSTER.print_containers_logs() + CLUSTER.print_containers_logs() fail(e) finally: _cleanup() diff --git a/integration/tests/fixtures/agent.rs b/integration/tests/fixtures/agent.rs index 1096aac..5f825d1 100644 --- a/integration/tests/fixtures/agent.rs +++ b/integration/tests/fixtures/agent.rs @@ -17,6 +17,7 @@ impl RegisteredAgent { pub async fn register_agent() -> RegisteredAgent { let cli = ClientHandler::new(&ENV.u_server, None); let agent_uid = Uuid::new_v4(); + println!("registering agent {agent_uid}"); let resp = cli .get_personal_jobs(agent_uid) .await @@ -27,6 +28,7 @@ pub async fn register_agent() -> RegisteredAgent { let job = cli.get_jobs(Some(job_id)).await.unwrap().pop().unwrap(); assert_eq!(job.alias, Some("agent_hello".to_string())); let mut agent_data = AssignedJob::from(&job); + agent_data.agent_id = agent_uid; agent_data.set_result(&Agent { id: agent_uid, ..Default::default() diff --git a/integration/tests/helpers/panel.rs b/integration/tests/helpers/panel.rs index d381898..4caecb2 100644 --- a/integration/tests/helpers/panel.rs +++ b/integration/tests/helpers/panel.rs @@ -43,8 +43,8 @@ impl Panel { .as_ref(), ); match &result { - PanelResult::Ok(r) => eprintln!("<<<+ {r:?}"), - PanelResult::Err(e) => eprintln!("<< eprintln!("<<<+ {r:02x?}"), + PanelResult::Err(e) => eprintln!("<<) -> Self { let identity = Identity::from_pkcs12_der(AGENT_IDENTITY, "").unwrap(); - let mut client = Client::builder().identity(identity); + let mut default_headers = HashMap::from([( + "user-agent".to_string(), + get_self_uid().hyphenated().to_string(), + )]); + if let Some(pwd) = password { - client = client.default_headers( - HeaderMap::try_from(&HashMap::from([( - "Authorization".to_string(), - format!("Bearer {pwd}"), - )])) - .unwrap(), - ) + default_headers.insert("authorization".to_string(), format!("Bearer {pwd}")); } - let client = 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()) .build() .unwrap(); + Self { client, base_url: Url::parse(&format!("https://{}:{}", server, MASTER_PORT)).unwrap(), diff --git a/lib/u_lib/src/models/agent.rs b/lib/u_lib/src/models/agent.rs index b00a677..227ce90 100644 --- a/lib/u_lib/src/models/agent.rs +++ b/lib/u_lib/src/models/agent.rs @@ -96,7 +96,7 @@ impl Agent { hostname: decoder(builder.pop("hostname")), is_root: &decoder(builder.pop("is_root")) == "0", username: decoder(builder.pop("username")), - platform: Platform::current().into_string(), + platform: Platform::current_as_string(), ..Default::default() } } diff --git a/lib/u_lib/src/models/jobs/meta.rs b/lib/u_lib/src/models/jobs/meta.rs index e839f8e..8150313 100644 --- a/lib/u_lib/src/models/jobs/meta.rs +++ b/lib/u_lib/src/models/jobs/meta.rs @@ -31,7 +31,7 @@ pub struct JobMeta { pub exec_type: JobType, ///target triple - #[serde(default)] + #[serde(default = "Platform::current_as_string")] pub platform: String, #[serde(default)] @@ -67,7 +67,7 @@ impl Default for JobMeta { alias: None, argv: String::new(), exec_type: JobType::Shell, - platform: Platform::current().into_string(), + platform: Platform::current_as_string(), payload: None, schedule: None, payload_path: None, diff --git a/lib/u_lib/src/utils/platform.rs b/lib/u_lib/src/utils/platform.rs index 6ecb82b..04aa4ae 100644 --- a/lib/u_lib/src/utils/platform.rs +++ b/lib/u_lib/src/utils/platform.rs @@ -15,6 +15,10 @@ impl Platform { Self(guess_host_triple().unwrap_or("unknown").to_string()) } + pub fn current_as_string() -> String { + Self::current().into_string() + } + pub fn matches(&self, pf: impl AsRef) -> bool { match PlatformReq::from_str(pf.as_ref()) { Ok(p) => p.matches(&_Platform::find(&self.0).unwrap()), diff --git a/migrations/2020-10-24-111622_create_all/up.sql b/migrations/2020-10-24-111622_create_all/up.sql index c2d5ac9..4edc097 100644 --- a/migrations/2020-10-24-111622_create_all/up.sql +++ b/migrations/2020-10-24-111622_create_all/up.sql @@ -4,60 +4,58 @@ CREATE TYPE JobState AS ENUM ('queued', 'running', 'finished'); CREATE TYPE AgentState AS ENUM ('new', 'active', 'banned'); CREATE TABLE IF NOT EXISTS agents ( - alias TEXT - , hostname TEXT NOT NULL - , id UUID NOT NULL DEFAULT uuid_generate_v4() - , ip_gray TEXT - , ip_white TEXT - , is_root BOOLEAN NOT NULL DEFAULT false - , is_root_allowed BOOLEAN NOT NULL DEFAULT false - , last_active TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - -- target triplet - , platform TEXT NOT NULL - , regtime TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - , state AgentState NOT NULL DEFAULT 'new' - -- is needed to processing requests - , token TEXT - , username TEXT NOT NULL - - , PRIMARY KEY(id) + alias TEXT, + hostname TEXT NOT NULL, + id UUID NOT NULL DEFAULT uuid_generate_v4(), + ip_gray TEXT, + ip_white TEXT, + is_root BOOLEAN NOT NULL DEFAULT false, + is_root_allowed BOOLEAN NOT NULL DEFAULT false, + last_active TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + platform TEXT NOT NULL, + regtime TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + state AgentState NOT NULL DEFAULT 'new', + token TEXT, + username TEXT NOT NULL, + + PRIMARY KEY(id) ); CREATE TABLE IF NOT EXISTS jobs ( - alias TEXT - , argv TEXT NOT NULL - , id UUID NOT NULL DEFAULT uuid_generate_v4() - , exec_type JobType NOT NULL DEFAULT 'shell' - , platform TEXT NOT NULL - , payload BYTEA - , payload_path TEXT - , schedule TEXT - - , PRIMARY KEY(id) + alias TEXT, + argv TEXT NOT NULL, + id UUID NOT NULL DEFAULT uuid_generate_v4(), + exec_type JobType NOT NULL DEFAULT 'shell', + platform TEXT NOT NULL, + payload BYTEA, + payload_path TEXT, + schedule TEXT, + + PRIMARY KEY(id) ); CREATE TABLE IF NOT EXISTS results ( - agent_id UUID NOT NULL - , alias TEXT - , created TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - , id UUID NOT NULL DEFAULT uuid_generate_v4() - , job_id UUID NOT NULL - , result BYTEA - , state JobState NOT NULL DEFAULT 'queued' - , exec_type JobType NOT NULL DEFAULT 'shell' - , retcode INTEGER - , updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP - - , FOREIGN KEY(agent_id) REFERENCES agents(id) ON DELETE CASCADE - , FOREIGN KEY(job_id) REFERENCES jobs(id) ON DELETE CASCADE - , PRIMARY KEY(id) + agent_id UUID NOT NULL, + alias TEXT, + created TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + id UUID NOT NULL DEFAULT uuid_generate_v4(), + job_id UUID NOT NULL, + result BYTEA, + state JobState NOT NULL DEFAULT 'queued', + exec_type JobType NOT NULL DEFAULT 'shell', + retcode INTEGER, + updated TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + FOREIGN KEY(agent_id) REFERENCES agents(id) ON DELETE CASCADE, + FOREIGN KEY(job_id) REFERENCES jobs(id) ON DELETE CASCADE, + PRIMARY KEY(id) ); CREATE TABLE IF NOT EXISTS certificates ( - agent_id UUID NOT NULL - , id UUID NOT NULL DEFAULT uuid_generate_v4() - , is_revoked BOOLEAN NOT NULL DEFAULT FALSE + agent_id UUID NOT NULL, + id UUID NOT NULL DEFAULT uuid_generate_v4(), + is_revoked BOOLEAN NOT NULL DEFAULT FALSE, - , PRIMARY KEY(id) - , FOREIGN KEY(agent_id) REFERENCES agents(id) + PRIMARY KEY(id), + FOREIGN KEY(agent_id) REFERENCES agents(id) ); \ No newline at end of file diff --git a/todos.txt b/spec.txt similarity index 60% rename from todos.txt rename to spec.txt index 5be7b2a..53135f6 100644 --- a/todos.txt +++ b/spec.txt @@ -1,7 +1,9 @@ +todos: Upload/download files More tests Agent update (use more JobType's) Erase log macros in release mode Bump wine version to test agent on windows -Store downloaded payload on disk instead of memory -Improve web interface \ No newline at end of file +Store downloaded payload on disk instead of ram +Improve web interface +Migrator binary