Imix main testing

docs/_docs/user-guide/imix.md

@@ -6,19 +6,23 @@ description: Imix User Guide
 permalink: user-guide/imix
 ---
 ## What is Imix
+
 Imix is the default agent for realm.
 Imix currently only supports http callbacks which interact directly with the graphql API.
 
 ## Configuration
+
 By default Imix is configured using a JSON file at run time.
 
 The config is specified at run time with the `-c` flag.
 For example:
+
 ```bash
 ./imix -c /tmp/imix-config.json
 ```
 
 The imix config is as follows:
+
 ```json
 {
     "service_configs": [],
@@ -31,7 +35,7 @@ The imix config is as follows:
         "c2_configs": [
         {
             "priority": 1,
-            "uri": "http://127.0.0.1/graphql"
+            "uri": "http://127.0.0.1/grpc"
         }
         ]
     }
@@ -42,18 +46,20 @@ The imix config is as follows:
 - `target_forward_connect_ip`: The IP address that you the red teamer would interact with the host through. This is to help keep track of agents when a hosts internal IP is different from the one you interact with in the case of a host behind a proxy.
 - `target_name`: Currently unused.
 - `callback_config`: Define where and when the agent should callback.
-    - `interval`: Number of seconds between callbacks.
-    - `jitter`: Currently unused.
-    - `timeout`: The number of seconds to wait before aborting a connection attempt.
-    - `c2_config` Define where the c2 should callback to.
-        - `priority`: The index that a domain should have.
-        - `uri`: The full URI of the callback endpoint.
+  - `interval`: Number of seconds between callbacks.
+  - `jitter`: Currently unused.
+  - `timeout`: The number of seconds to wait before aborting a connection attempt.
+  - `c2_config` Define where the c2 should callback to.
+    - `priority`: The index that a domain should have.
+    - `uri`: The full URI of the callback endpoint.
 
 ## Functionality
+
 Imix derives all it's functionality from the eldritch language.
 See the [Eldritch User Guide](/user-guide/eldritch) for more information.
 
 ## Task management
+
 Imix can execute up to 127 threads concurrently after that the main imix thread will block behind other threads.
 Every callback interval imix will query each active thread for new output and rely that back to the c2. This means even long running tasks will report their status as new data comes in.
 
@@ -71,15 +77,16 @@ RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --target=x86_64-
 ```
 
 ### MacOS
+
 **MacOS does not support static compilation**
-https://developer.apple.com/forums/thread/706419
+<https://developer.apple.com/forums/thread/706419>
 
 **Cross compilation is more complicated than we'll support**
 Check out this blog a starting point for cross compiling.
-https://wapl.es/rust/2019/02/17/rust-cross-compile-linux-to-macos.html/
-
+<https://wapl.es/rust/2019/02/17/rust-cross-compile-linux-to-macos.html/>
 
 ### Windows
+
 ```bash
 rustup target add x86_64-pc-windows-gnu
 
@@ -88,4 +95,3 @@ sudo apt install gcc-mingw-w64
 
 RUSTFLAGS="-C target-feature=+crt-static" cargo build --release --target=x86_64-pc-windows-gnu
 ```
-

implants/imix/src/exec.rs

@@ -0,0 +1,227 @@
+use anyhow::{Error, Result};
+use c2::pb::Task;
+use chrono::{DateTime, Utc};
+use eldritch::{eldritch_run, EldritchPrintHandler};
+use std::sync::mpsc::Receiver;
+use std::sync::mpsc::Sender;
+use std::thread;
+use tokio::task::JoinHandle;
+use tokio::time::Duration;
+
+pub struct AsyncTask {
+    pub future_join_handle: JoinHandle<Result<(), Error>>,
+    pub start_time: DateTime<Utc>,
+    pub grpc_task: Task,
+    pub print_reciever: Receiver<String>,
+}
+
+async fn handle_exec_tome(
+    task: Task,
+    print_channel_sender: Sender<String>,
+) -> Result<(String, String)> {
+    // TODO: Download auxillary files from CDN
+
+    // Read a tome script
+    // let task_quest = match task.quest {
+    //     Some(quest) => quest,
+    //     None => return Ok(("".to_string(), format!("No quest associated for task ID: {}", task.id))),
+    // };
+
+    let print_handler = EldritchPrintHandler {
+        sender: print_channel_sender,
+    };
+
+    // Execute a tome script
+    let res = match thread::spawn(move || {
+        eldritch_run(
+            task.id.to_string(),
+            task.eldritch.clone(),
+            Some(task.parameters.clone()),
+            &print_handler,
+        )
+    })
+    .join()
+    {
+        Ok(local_thread_res) => local_thread_res,
+        Err(_) => todo!(),
+    };
+    match res {
+        Ok(tome_output) => Ok((tome_output, "".to_string())),
+        Err(tome_error) => Ok(("".to_string(), tome_error.to_string())),
+    }
+}
+
+pub async fn handle_exec_timeout_and_response(
+    task: Task,
+    print_channel_sender: Sender<String>,
+    timeout: Option<Duration>,
+) -> Result<(), Error> {
+    // Tasks will be forcebly stopped after 1 week.
+    let timeout_duration = timeout.unwrap_or_else(|| Duration::from_secs(60 * 60 * 24 * 7));
+
+    // Define a future for our execution task
+    let exec_future = handle_exec_tome(task.clone(), print_channel_sender.clone());
+    // Execute that future with a timeout defined by the timeout argument.
+    let tome_result = match tokio::time::timeout(timeout_duration, exec_future).await {
+        Ok(res) => match res {
+            Ok(tome_result) => tome_result,
+            Err(tome_error) => ("".to_string(), tome_error.to_string()),
+        },
+        Err(timer_elapsed) => (
+            "".to_string(),
+            format!(
+                "Time elapsed task {} has been running for {} seconds",
+                task.id,
+                timer_elapsed.to_string()
+            ),
+        ),
+    };
+
+    print_channel_sender
+        .clone()
+        .send(format!("---[RESULT]----\n{}\n---------", tome_result.0))?;
+    print_channel_sender
+        .clone()
+        .send(format!("---[ERROR]----\n{}\n--------", tome_result.1))?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::handle_exec_tome;
+    use anyhow::Result;
+    use c2::pb::Task;
+    use std::collections::HashMap;
+    use std::sync::mpsc::channel;
+    use std::time::Duration;
+
+    #[test]
+    fn imix_handle_exec_tome() -> Result<()> {
+        let test_tome_input = Task {
+            id: 123,
+            eldritch: r#"
+print(sys.shell(input_params["cmd"])["stdout"])
+1"#
+            .to_string(),
+            parameters: HashMap::from([("cmd".to_string(), "echo hello_from_stdout".to_string())]),
+        };
+
+        let runtime = tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+
+        let (sender, receiver) = channel::<String>();
+
+        let exec_future = handle_exec_tome(test_tome_input, sender.clone());
+        let (eld_output, eld_error) = runtime.block_on(exec_future)?;
+
+        let cmd_output = receiver.recv_timeout(Duration::from_millis(500))?;
+        assert!(cmd_output.contains("hello_from_stdout"));
+        assert_eq!(eld_output, "1".to_string());
+        assert_eq!(eld_error, "".to_string());
+        Ok(())
+    }
+
+    #[test]
+    fn imix_handle_exec_tome_error() -> Result<()> {
+        let test_tome_input = Task {
+            id: 123,
+            eldritch: r#"
+aoeu
+"#
+            .to_string(),
+            parameters: HashMap::new(),
+        };
+
+        let runtime: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
+            .enable_all()
+            .build()
+            .unwrap();
+
+        let (sender, receiver) = channel::<String>();
+
+        let exec_future = handle_exec_tome(test_tome_input, sender.clone());
+        let (eld_output, eld_error) = runtime.block_on(exec_future)?;
+
+        let mut index = 0;
+        loop {
+            let cmd_output = match receiver.recv_timeout(Duration::from_millis(500)) {
+                Ok(local_res_string) => local_res_string,
+                Err(local_err) => {
+                    match local_err.to_string().as_str() {
+                        "channel is empty and sending half is closed" => {
+                            break;
+                        }
+                        "timed out waiting on channel" => break,
+                        _ => eprint!("Error: {}", local_err),
+                    }
+                    break;
+                }
+            };
+            assert_eq!(cmd_output, "".to_string());
+
+            index = index + 1;
+        }
+
+        assert_eq!(eld_output, "".to_string());
+        assert_eq!(eld_error, "[eldritch] Eldritch eval_module failed:\nerror: Variable `aoeu` not found\n --> 123:2:1\n  |\n2 | aoeu\n  | ^^^^\n  |\n".to_string());
+        Ok(())
+    }
+
+    // This test
+    //     #[test]
+    //     fn imix_handle_exec_tome_timeout() -> Result<()> {
+    //         let test_tome_input = Task {
+    //             id: 123,
+    //             eldritch: r#"
+    // print("Hello_world")
+    // time.sleep(5)
+    // "#
+    //             .to_string(),
+    //             parameters: HashMap::new(),
+    //         };
+
+    //         let runtime: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
+    //             .enable_all()
+    //             .build()
+    //             .unwrap();
+
+    //         let (sender, receiver) = channel::<String>();
+
+    //         let start_time = Instant::now();
+    //         let exec_future = handle_exec_timeout_and_response(
+    //             test_tome_input,
+    //             sender.clone(),
+    //             Some(Duration::from_secs(2)),
+    //         );
+    //         runtime.block_on(exec_future)?;
+    //         let end_time = Instant::now();
+    //         let mut index = 0;
+    //         loop {
+    //             let cmd_output = match receiver.recv_timeout(Duration::from_millis(800)) {
+    //                 Ok(local_res_string) => local_res_string,
+    //                 Err(local_err) => {
+    //                     match local_err.to_string().as_str() {
+    //                         "channel is empty and sending half is closed" => {
+    //                             break;
+    //                         }
+    //                         "timed out waiting on channel" => break,
+    //                         _ => eprint!("Error: {}", local_err),
+    //                     }
+    //                     break;
+    //                 }
+    //             };
+    //             println!("eld_output: {}", cmd_output);
+    //             index = index + 1;
+    //         }
+
+    //         println!(
+    //             "Diff {:?}",
+    //             end_time.checked_duration_since(start_time).unwrap()
+    //         );
+    //         assert!(end_time.checked_duration_since(start_time).unwrap() < Duration::from_secs(3));
+
+    //         Ok(())
+    //     }
+}