Add a Run Once capability to Imix

docs/_docs/user-guide/imix.md

@@ -20,6 +20,13 @@ Imix has compile-time configuration, that may be specified using environment var
 | IMIX_RETRY_INTERVAL | Duration to wait before restarting the agent loop if an error occurs, in seconds. | `5` | No |
 | IMIX_PROXY_URI | Overide system settings for proxy URI over HTTP(S) (must specify a scheme, e.g. `https://`) | No proxy | No |
 | IMIX_HOST_ID | Manually specify the host ID for this beacon. Supersedes the file on disk. | - | No |
+| IMIX_RUN_ONCE | Imix will only do one callback and execution of queued tasks (may want to pair with runtime environment variable `IMIX_BEACON_ID`) | false | No |
+
+Imix has run-time configuration, that may be specified using environment variables during execution.
+
+| Env Var | Description | Default | Required |
+| ------- | ----------- | ------- | -------- |
+| IMIX_BEACON_ID | The identifier to be used during callback (must be globally unique) | Random UUIDv4 | No |
 
 ## Logging
 

implants/imix/src/agent.rs

@@ -116,6 +116,10 @@ impl<T: Transport + 'static> Agent<T> {
                 }
             };
 
+            if self.cfg.run_once {
+                return Ok(());
+            }
+
             let interval = match self.cfg.info.clone() {
                 Some(b) => Ok(b.interval),
                 None => Err(anyhow::anyhow!("beacon info is missing from agent")),

implants/imix/src/run.rs

@@ -24,6 +24,8 @@ pub async fn handle_main() {
         #[cfg(debug_assertions)]
         log::info!("agent config initialized {:#?}", cfg.clone());
 
+        let run_once = cfg.run_once;
+
         match run(cfg).await {
             Ok(_) => {}
             Err(_err) => {
@@ -33,6 +35,10 @@ pub async fn handle_main() {
                 tokio::time::sleep(Duration::from_secs(retry_interval)).await;
             }
         }
+
+        if run_once {
+            break;
+        }
     }
 }
 

implants/lib/pb/src/config.rs

@@ -11,6 +11,8 @@ pub struct Config {
     pub proxy_uri: ::core::option::Option<::prost::alloc::string::String>,
     #[prost(uint64, tag = "4")]
     pub retry_interval: u64,
+    #[prost(bool, tag = "5")]
+    pub run_once: bool,
 }
 
 macro_rules! callback_uri {
@@ -64,6 +66,20 @@ macro_rules! retry_interval {
  */
 pub const RETRY_INTERVAL: &str = retry_interval!();
 
+macro_rules! run_once {
+    () => {
+        match option_env!("IMIX_RUN_ONCE") {
+            Some(_) => true,
+            None => false,
+        }
+    };
+}
+
+/* Compile-time constant for the agent run once flag, derived from the IMIX_RUN_ONCE environment variable during compilation.
+ * Defaults to false if unset.
+ */
+pub const RUN_ONCE: bool = run_once!();
+
 /*
  * Config methods.
  */
@@ -82,8 +98,12 @@ impl Config {
             primary_ip: get_primary_ip(),
         };
 
+        // Try to grab the beacon identitifier from env var, o/w use  a random UUID
+        let beacon_id =
+            std::env::var("IMIX_BEACON_ID").unwrap_or_else(|_| String::from(Uuid::new_v4()));
+
         let info = crate::c2::Beacon {
-            identifier: String::from(Uuid::new_v4()),
+            identifier: beacon_id,
             principal: whoami::username(),
             interval: match CALLBACK_INTERVAL.parse::<u64>() {
                 Ok(i) => i,
@@ -113,6 +133,7 @@ impl Config {
                     5
                 }
             },
+            run_once: RUN_ONCE,
         }
     }
     pub fn refresh_primary_ip(&mut self) {