Small cleanup

MythicAgents · Feb 11, 2024 · 65d3be3 · 65d3be3
1 parent 338f781
commit 65d3be3
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 11 deletions.
diff --git a/Payload_Type/thanatos/agent/config/src/structs.rs b/Payload_Type/thanatos/agent/config/src/structs.rs
@@ -15,6 +15,7 @@ pub enum InitOption {
     Thread = 1,
 
     /// Payload should fork to the background
+    #[cfg(target_os = "linux")]
     Fork = 2,
 }
 
@@ -82,15 +83,15 @@ impl ConfigVars<'_> {
         self.init_option
     }
 
-    pub fn domains(&self) -> &Vec<[u8; 32]> {
+    pub fn domains(&self) -> &[[u8; 32]] {
         &self.domains
     }
 
-    pub fn hostnames(&self) -> &Vec<[u8; 32]> {
+    pub fn hostnames(&self) -> &[[u8; 32]] {
         &self.hostnames
     }
 
-    pub fn usernames(&self) -> &Vec<[u8; 32]> {
+    pub fn usernames(&self) -> &[[u8; 32]] {
         &self.usernames
     }
 }
diff --git a/Payload_Type/thanatos/agent/thanatos_core/src/lib.rs b/Payload_Type/thanatos/agent/thanatos_core/src/lib.rs
@@ -1,14 +1,14 @@
 #![forbid(unsafe_code)]
 
-use config::ConfigVars;
+use config::{ConfigVars, InitOption};
 
 mod guardrails;
 pub mod logging;
 pub mod native;
 
-pub fn initialize_agent<'a, 'b: 'a, F>(f: F, config: &'a ConfigVars<'b>)
+pub fn initialize_agent<F>(f: F, config: ConfigVars<'static>)
 where
-    F: Fn(&'a ConfigVars<'b>),
+    F: Fn(ConfigVars<'static>) + Send + Sync + 'static,
 {
     let domains = config.domains();
     if !domains.is_empty() && !guardrails::check_domain(domains) {
@@ -25,5 +25,14 @@ where
         return;
     }
 
-    f(config)
+    match config.init_option() {
+        InitOption::Thread => {
+            std::thread::spawn(move || f(config));
+        }
+        #[cfg(target_os = "linux")]
+        InitOption::Fork => todo!(),
+        InitOption::None => {
+            f(config);
+        }
+    }
 }
diff --git a/Payload_Type/thanatos/agent/thanatos_http/src/lib.rs b/Payload_Type/thanatos/agent/thanatos_http/src/lib.rs
@@ -6,11 +6,11 @@
 
 use config::ConfigVars;
 
-pub fn entrypoint(config_bytes: &[u8]) {
-    let agent_config: ConfigVars = rmp_serde::from_slice(config_bytes).unwrap();
-    thanatos_core::initialize_agent(run_agent, &agent_config);
+pub fn entrypoint(config_bytes: &'static [u8]) {
+    let config = rmp_serde::from_slice(config_bytes).unwrap();
+    thanatos_core::initialize_agent(run_agent, config);
 }
 
-fn run_agent(config: &ConfigVars) {
+fn run_agent(config: ConfigVars) {
     thanatos_core::debug!(config);
 }