diff --git a/implants/Cargo.toml b/implants/Cargo.toml index 2e673c141..029ffaee3 100644 --- a/implants/Cargo.toml +++ b/implants/Cargo.toml @@ -24,7 +24,7 @@ members = [ "lib/eldritch/stdlib/eldritch-libtime", "lib/eldritch/eldritch", "lib/eldritch/eldritch-wasm", - "lib/portals/portal-stream", + "lib/portals/portal-stream", "lib/config", ] exclude = [ "lib/eldritch/stdlib/tests", # Excluded to prevent fake_bindings from polluting workspace builds @@ -35,6 +35,7 @@ resolver = "2" transport = { path = "./lib/transport" } host_unique = { path = "./lib/host_unique" } pb = { path = "./lib/pb" } +config = { path = "./lib/config" } netstat = { path = "./lib/netstat" } # Eldritch V2 diff --git a/implants/imix/Cargo.toml b/implants/imix/Cargo.toml index 8ea593d01..f206c2ba0 100644 --- a/implants/imix/Cargo.toml +++ b/implants/imix/Cargo.toml @@ -17,6 +17,7 @@ install = [] tokio-console = ["dep:console-subscriber", "tokio/tracing"] [dependencies] +config = { workspace = true, features = ["imix"] } tokio = { workspace = true, features = [ "rt-multi-thread", "macros", diff --git a/implants/imix/src/agent.rs b/implants/imix/src/agent.rs index 1b8a1b66d..49a370320 100644 --- a/implants/imix/src/agent.rs +++ b/implants/imix/src/agent.rs @@ -1,4 +1,5 @@ use anyhow::{Context as AnyhowContext, Result}; +use config::Config; use eldritch::agent::agent::Agent; use eldritch_agent::Context; use pb::c2::host::Platform; @@ -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; diff --git a/implants/imix/src/main.rs b/implants/imix/src/main.rs index cede74bc8..39310d92d 100644 --- a/implants/imix/src/main.rs +++ b/implants/imix/src/main.rs @@ -16,7 +16,7 @@ mod win_service; #[cfg(all(debug_assertions, feature = "tokio-console"))] use console_subscriber; -pub use pb::config::Config; +pub use config::Config; pub use transport::{ActiveTransport, Transport}; mod agent; @@ -33,6 +33,8 @@ mod version; #[tokio::main] async fn main() -> Result<()> { + pb::xchacha::init_server_pubkey(config::SERVER_PUBKEY); + #[cfg(all(debug_assertions, feature = "tokio-console"))] { console_subscriber::init(); diff --git a/implants/imix/src/run.rs b/implants/imix/src/run.rs index a1ff0c157..d98c9a651 100644 --- a/implants/imix/src/run.rs +++ b/implants/imix/src/run.rs @@ -6,7 +6,7 @@ use std::time::{Duration, Instant}; use crate::agent::ImixAgent; use crate::task::TaskRegistry; use crate::version::VERSION; -use pb::config::Config; +use config::Config; use transport::{ActiveTransport, Transport}; pub static SHUTDOWN: AtomicBool = AtomicBool::new(false); diff --git a/implants/imix/src/shell/manager.rs b/implants/imix/src/shell/manager.rs index ba30e2b35..91731e6f4 100644 --- a/implants/imix/src/shell/manager.rs +++ b/implants/imix/src/shell/manager.rs @@ -375,8 +375,8 @@ impl ShellManager { mod tests { use super::*; use crate::task::TaskRegistry; + use config::Config; use pb::c2::{ReportOutputResponse, ShellTask}; - use pb::config::Config; use transport::MockTransport; #[tokio::test] diff --git a/implants/imix/src/tests/agent_output_aggregation.rs b/implants/imix/src/tests/agent_output_aggregation.rs index d8e7773ef..25bf58b2a 100644 --- a/implants/imix/src/tests/agent_output_aggregation.rs +++ b/implants/imix/src/tests/agent_output_aggregation.rs @@ -1,10 +1,10 @@ use crate::agent::ImixAgent; use crate::task::TaskRegistry; +use config::Config; 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; diff --git a/implants/imix/src/tests/agent_tests.rs b/implants/imix/src/tests/agent_tests.rs index e44255797..ea333abf5 100644 --- a/implants/imix/src/tests/agent_tests.rs +++ b/implants/imix/src/tests/agent_tests.rs @@ -1,7 +1,7 @@ use super::super::agent::ImixAgent; use super::super::task::TaskRegistry; +use config::Config; use eldritch::agent::agent::Agent; -use pb::config::Config; use std::sync::Arc; use transport::MockTransport; diff --git a/implants/imix/src/tests/agent_trait_tests.rs b/implants/imix/src/tests/agent_trait_tests.rs index d4908bdb7..b6e6576c8 100644 --- a/implants/imix/src/tests/agent_trait_tests.rs +++ b/implants/imix/src/tests/agent_trait_tests.rs @@ -1,9 +1,9 @@ use super::super::agent::ImixAgent; use super::super::task::TaskRegistry; +use 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; diff --git a/implants/imix/src/tests/callback_interval_test.rs b/implants/imix/src/tests/callback_interval_test.rs index b3a0ea3de..009fd2b56 100644 --- a/implants/imix/src/tests/callback_interval_test.rs +++ b/implants/imix/src/tests/callback_interval_test.rs @@ -1,7 +1,7 @@ use super::super::agent::ImixAgent; use super::super::task::TaskRegistry; +use config::Config; use eldritch::agent::agent::Agent; -use pb::config::Config; use std::sync::Arc; use transport::MockTransport; diff --git a/implants/imix/src/tests/report_large_file_test.rs b/implants/imix/src/tests/report_large_file_test.rs index c1ebbf76b..bbe6cc0c5 100644 --- a/implants/imix/src/tests/report_large_file_test.rs +++ b/implants/imix/src/tests/report_large_file_test.rs @@ -1,5 +1,6 @@ use crate::agent::ImixAgent; use crate::task::TaskRegistry; +use config::Config; use eldritch::report::std::file_impl; use eldritch_agent::{Agent, Context}; use pb::c2::{ @@ -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; diff --git a/implants/lib/config/Cargo.toml b/implants/lib/config/Cargo.toml new file mode 100644 index 000000000..9265ad850 --- /dev/null +++ b/implants/lib/config/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "config" +version = "0.1.0" +edition = "2024" + +[features] +default = [] +imix = [] + +[dependencies] +pb = { workspace = true } +anyhow = { workspace = true } +host_unique = { workspace = true } +log = { workspace = true } +netdev = { workspace = true } +prost = { workspace = true } +url = { workspace = true } +uuid = { workspace = true, features = ["v4", "fast-rng"] } +whoami = { workspace = true } +const-decoder = { workspace = true } + +[build-dependencies] +reqwest = { workspace = true, features = ["blocking", "json", "rustls-tls", "http2"] } +serde_json = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_yaml = { workspace = true } +urlencoding = { workspace = true } diff --git a/implants/lib/config/build.rs b/implants/lib/config/build.rs new file mode 100644 index 000000000..282d2253e --- /dev/null +++ b/implants/lib/config/build.rs @@ -0,0 +1,307 @@ +use serde::Deserialize; + +#[derive(Debug, Deserialize)] +struct TransportConfig { + #[serde(rename = "URI")] + uri: String, + #[serde(rename = "type")] + transport_type: String, + extra: String, + #[serde(default)] + interval: Option, + #[serde(default)] + jitter: Option, +} + +#[derive(Debug, Deserialize)] +struct YamlConfig { + transports: Vec, + #[serde(default)] + server_pubkey: Option, +} + +/// Result of parsing YAML config, containing values needed by other build steps +struct YamlConfigResult { + /// The first transport URI (used for fetching pubkey) + upstream_uri: Option, + /// Server public key if specified in config + server_pubkey: Option, +} + +fn parse_yaml_config() -> Result, Box> { + // Check if IMIX_CONFIG is set + let config_yaml = match std::env::var("IMIX_CONFIG") { + Ok(yaml_content) => yaml_content, + Err(_) => { + println!("cargo:warning=IMIX_CONFIG not set, skipping YAML config parsing"); + return Ok(None); + } + }; + + // Check that other configuration options are not set + let has_callback_uri = std::env::var("IMIX_CALLBACK_URI").is_ok(); + let has_callback_interval = std::env::var("IMIX_CALLBACK_INTERVAL").is_ok(); + let has_transport_extra = std::env::vars().any(|(k, _)| k.starts_with("IMIX_TRANSPORT_EXTRA_")); + + if has_callback_uri || has_callback_interval || has_transport_extra { + let mut error_msg = String::from( + "Configuration error: Cannot use IMIX_CONFIG with other configuration options.\n", + ); + error_msg.push_str( + "When IMIX_CONFIG is set, all configuration must be done through the YAML file.\n", + ); + error_msg.push_str("Found one or more of:\n"); + + if has_callback_uri { + error_msg.push_str(" - IMIX_CALLBACK_URI\n"); + } + if has_callback_interval { + error_msg.push_str(" - IMIX_CALLBACK_INTERVAL\n"); + } + if has_transport_extra { + error_msg.push_str(" - IMIX_TRANSPORT_EXTRA_*\n"); + } + + error_msg.push_str( + "\nPlease use ONLY the YAML config file OR use environment variables, but not both.", + ); + + return Err(error_msg.into()); + } + + // Parse the YAML config + let config: YamlConfig = serde_yaml::from_str(&config_yaml) + .map_err(|e| format!("Failed to parse YAML config: {}", e))?; + + // Validate that we have at least one transport + if config.transports.is_empty() { + return Err("YAML config must contain at least one transport".into()); + } + + // Build DSN string from transports + let mut dsn_parts = Vec::new(); + + for transport in &config.transports { + // Validate transport type + let transport_type_lower = transport.transport_type.to_lowercase(); + if !["grpc", "http1", "dns"].contains(&transport_type_lower.as_str()) { + return Err(format!( + "Invalid transport type '{}'. Must be one of: GRPC, http1, DNS", + transport.transport_type + ) + .into()); + } + + // Validate that extra is valid JSON + if !transport.extra.is_empty() { + serde_json::from_str::(&transport.extra).map_err(|e| { + format!( + "Invalid JSON in 'extra' field for transport '{}': {}", + transport.uri, e + ) + })?; + } + + // Error if URI already contains query parameters + if transport.uri.contains('?') { + return Err(format!("URI '{}' already contains query parameters. Query parameters should not be present in the URI field.", transport.uri).into()); + } + + // Build DSN part with correct schema and query parameters + let mut dsn_part = transport.uri.clone(); + + // Add query parameters + dsn_part.push('?'); + let mut params = Vec::new(); + + // Add interval if present + if let Some(interval) = transport.interval { + params.push(format!("interval={}", interval)); + } + + // Add jitter if present + if let Some(jitter) = transport.jitter { + params.push(format!("jitter={}", jitter)); + } + + // Add extra as query parameter if not empty + if !transport.extra.is_empty() { + let encoded_extra = urlencoding::encode(&transport.extra); + params.push(format!("extra={}", encoded_extra)); + } + + if !params.is_empty() { + dsn_part.push_str(¶ms.join("&")); + } else { + // Remove the trailing '?' if no params were added + dsn_part.pop(); + } + + dsn_parts.push(dsn_part); + } + + // Join all DSN parts with semicolons + let dsn = dsn_parts.join(";"); + + // Emit the DSN configuration + println!("cargo:rustc-env=IMIX_CALLBACK_URI={}", dsn); + + // Emit server_pubkey if present + if let Some(ref pubkey) = config.server_pubkey { + println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey); + println!("cargo:warning=Using server_pubkey from YAML config"); + } + + println!( + "cargo:warning=Successfully parsed YAML config with {} transport(s)", + config.transports.len() + ); + + // Extract the first transport URI for pubkey fetching + let upstream_uri = config.transports.first().map(|t| t.uri.clone()); + + Ok(Some(YamlConfigResult { + upstream_uri, + server_pubkey: config.server_pubkey, + })) +} + +fn get_pub_key(yaml_config: Option) { + // Check if server pubkey was provided via YAML config + if let Some(ref config) = yaml_config { + if config.server_pubkey.is_some() { + // Already emitted in parse_yaml_config, no need to fetch + println!("cargo:warning=Server pubkey provided via YAML config, skipping fetch"); + return; + } + } + + // Check if IMIX_SERVER_PUBKEY is already set via env var + if std::env::var("IMIX_SERVER_PUBKEY").is_ok() { + println!("cargo:warning=IMIX_SERVER_PUBKEY already set, skipping fetch"); + return; + } + + // Get the callback URI: prefer YAML config upstream, then env var, then default + let callback_uri = yaml_config + .and_then(|c| c.upstream_uri) + .or_else(|| std::env::var("IMIX_CALLBACK_URI").ok()) + .unwrap_or_else(|| "http://127.0.0.1:8000".to_string()); + + // Extract the first URI from semicolon-separated list and strip query parameters + let base_uri = callback_uri + .split(';') + .next() + .unwrap_or(&callback_uri) + .trim() + .split('?') + .next() + .unwrap_or(&callback_uri); + + // Construct the status endpoint URL + let status_url = format!("{}/status", base_uri); + + // Make a GET request to /status + let response = match reqwest::blocking::get(&status_url) { + Ok(resp) => resp, + Err(e) => { + println!("cargo:warning=Failed to connect to {}: {}", status_url, e); + return; + } + }; + + if !response.status().is_success() { + println!( + "cargo:warning=Failed to fetch status from {}: HTTP {}", + status_url, + response.status() + ); + return; + } + + let json = match response.json::() { + Ok(json) => json, + Err(e) => { + println!( + "cargo:warning=Failed to parse JSON response from {}: {}", + status_url, e + ); + return; + } + }; + + let pubkey = match json.get("Pubkey").and_then(|v| v.as_str()) { + Some(key) => key, + None => { + println!( + "cargo:warning=Pubkey field not found in response from {}", + status_url + ); + return; + } + }; + + // Set the IMIX_SERVER_PUBKEY environment variable for the build + println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey); + println!( + "cargo:warning=Successfully fetched server public key from {}", + status_url + ); +} + +fn validate_dsn_config() -> Result<(), Box> { + // Skip validation if YAML config is being used + // (parse_yaml_config already handles validation in that case) + if std::env::var("IMIX_CONFIG").is_ok() { + return Ok(()); + } + + // Check if IMIX_CALLBACK_URI contains query parameters + let callback_uri = + std::env::var("IMIX_CALLBACK_URI").unwrap_or_else(|_| "http://127.0.0.1:8000".to_string()); + let has_query_params = callback_uri.contains('?'); + + // Check if legacy config environment variables are set + let has_callback_interval = std::env::var("IMIX_CALLBACK_INTERVAL").is_ok(); + let has_transport_extra = std::env::vars().any(|(k, _)| k.starts_with("IMIX_TRANSPORT_EXTRA_")); + + // If DSN has query parameters AND legacy config is set, this is an error + if has_query_params && (has_callback_interval || has_transport_extra) { + let mut error_msg = String::from( + "Configuration error: Cannot use both DSN query parameters and legacy environment variables.\n", + ); + error_msg.push_str("Found query parameters in IMIX_CALLBACK_URI and one or more of:\n"); + + if has_callback_interval { + error_msg.push_str(" - IMIX_CALLBACK_INTERVAL\n"); + } + if has_transport_extra { + error_msg.push_str(" - IMIX_TRANSPORT_EXTRA_*\n"); + } + + error_msg.push_str("\nPlease use ONLY DSN query parameters (e.g., https://example.com?interval=10&extra={...})\n"); + error_msg.push_str("OR use legacy environment variables, but not both."); + + return Err(error_msg.into()); + } + + Ok(()) +} + +fn main() -> Result<(), Box> { + // Tell Cargo to rerun this build script if these env vars change + // This fixes the issue where changing IMIX_CONFIG doesn't trigger a rebuild + println!("cargo:rerun-if-env-changed=IMIX_CONFIG"); + println!("cargo:rerun-if-env-changed=IMIX_CALLBACK_URI"); + println!("cargo:rerun-if-env-changed=IMIX_CALLBACK_INTERVAL"); + println!("cargo:rerun-if-env-changed=IMIX_SERVER_PUBKEY"); + + // Parse YAML config if present (this will emit IMIX_CALLBACK_URI if successful) + let yaml_config = parse_yaml_config()?; + + // Validate DSN config (skips if YAML config was used) + validate_dsn_config()?; + + get_pub_key(yaml_config); + Ok(()) +} diff --git a/implants/lib/pb/src/config.rs b/implants/lib/config/src/lib.rs similarity index 87% rename from implants/lib/pb/src/config.rs rename to implants/lib/config/src/lib.rs index aa8dbe9a3..c6f54c921 100644 --- a/implants/lib/pb/src/config.rs +++ b/implants/lib/config/src/lib.rs @@ -3,7 +3,23 @@ use host_unique::HostIDSelector; use url::Url; use uuid::Uuid; -use crate::c2::{AvailableTransports, Transport}; +use pb::c2::{AvailableTransports, Transport}; + +#[cfg(feature = "imix")] +use const_decoder::Decoder as const_decode; + +/* Compile-time constant for the server pubkey, derived from the IMIX_SERVER_PUBKEY environment variable during compilation. + * To find the servers pubkey check the startup messages on the server look for `[INFO] Public key: ` + */ +#[cfg(feature = "imix")] +pub static SERVER_PUBKEY: [u8; 32] = + const_decode::Base64.decode(env!("IMIX_SERVER_PUBKEY").as_bytes()); + +#[cfg(not(feature = "imix"))] +pub static SERVER_PUBKEY: [u8; 32] = [ + 165, 30, 122, 188, 50, 89, 111, 214, 247, 4, 189, 217, 188, 37, 200, 190, 2, 180, 175, 107, + 194, 147, 177, 98, 103, 84, 99, 120, 72, 73, 87, 37, +]; //TODO: Can this struct be removed? /// Config holds values necessary to configure an Agent. @@ -11,7 +27,7 @@ use crate::c2::{AvailableTransports, Transport}; #[derive(Clone, PartialEq, ::prost::Message)] pub struct Config { #[prost(message, optional, tag = "1")] - pub info: ::core::option::Option, + pub info: ::core::option::Option, #[prost(bool, tag = "2")] pub run_once: bool, } @@ -86,14 +102,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 +216,13 @@ fn parse_host_unique_selectors() -> Vec> { */ 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 +242,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 +284,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 +305,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 +346,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] diff --git a/implants/lib/pb/Cargo.toml b/implants/lib/pb/Cargo.toml index b8fc2fa7a..0bdc43b23 100644 --- a/implants/lib/pb/Cargo.toml +++ b/implants/lib/pb/Cargo.toml @@ -38,8 +38,3 @@ url = { workspace = true } tonic-prost-build = { workspace = true } which = { workspace = true } home = "=0.5.11" -reqwest = { workspace = true, features = ["blocking", "json", "rustls-tls", "http2"] } -serde_json = { workspace = true } -serde = { workspace = true, features = ["derive"] } -serde_yaml = { workspace = true } -urlencoding = { workspace = true } diff --git a/implants/lib/pb/build.rs b/implants/lib/pb/build.rs index e009c2589..928f1e47a 100644 --- a/implants/lib/pb/build.rs +++ b/implants/lib/pb/build.rs @@ -1,311 +1,10 @@ -use serde::Deserialize; use std::env; use std::path::PathBuf; use which::which; -#[derive(Debug, Deserialize)] -struct TransportConfig { - #[serde(rename = "URI")] - uri: String, - #[serde(rename = "type")] - transport_type: String, - extra: String, - #[serde(default)] - interval: Option, - #[serde(default)] - jitter: Option, -} - -#[derive(Debug, Deserialize)] -struct YamlConfig { - transports: Vec, - #[serde(default)] - server_pubkey: Option, -} - -/// Result of parsing YAML config, containing values needed by other build steps -struct YamlConfigResult { - /// The first transport URI (used for fetching pubkey) - upstream_uri: Option, - /// Server public key if specified in config - server_pubkey: Option, -} - -fn parse_yaml_config() -> Result, Box> { - // Check if IMIX_CONFIG is set - let config_yaml = match std::env::var("IMIX_CONFIG") { - Ok(yaml_content) => yaml_content, - Err(_) => { - println!("cargo:warning=IMIX_CONFIG not set, skipping YAML config parsing"); - return Ok(None); - } - }; - - // Check that other configuration options are not set - let has_callback_uri = std::env::var("IMIX_CALLBACK_URI").is_ok(); - let has_callback_interval = std::env::var("IMIX_CALLBACK_INTERVAL").is_ok(); - let has_transport_extra = std::env::vars().any(|(k, _)| k.starts_with("IMIX_TRANSPORT_EXTRA_")); - - if has_callback_uri || has_callback_interval || has_transport_extra { - let mut error_msg = String::from( - "Configuration error: Cannot use IMIX_CONFIG with other configuration options.\n", - ); - error_msg.push_str( - "When IMIX_CONFIG is set, all configuration must be done through the YAML file.\n", - ); - error_msg.push_str("Found one or more of:\n"); - - if has_callback_uri { - error_msg.push_str(" - IMIX_CALLBACK_URI\n"); - } - if has_callback_interval { - error_msg.push_str(" - IMIX_CALLBACK_INTERVAL\n"); - } - if has_transport_extra { - error_msg.push_str(" - IMIX_TRANSPORT_EXTRA_*\n"); - } - - error_msg.push_str( - "\nPlease use ONLY the YAML config file OR use environment variables, but not both.", - ); - - return Err(error_msg.into()); - } - - // Parse the YAML config - let config: YamlConfig = serde_yaml::from_str(&config_yaml) - .map_err(|e| format!("Failed to parse YAML config: {}", e))?; - - // Validate that we have at least one transport - if config.transports.is_empty() { - return Err("YAML config must contain at least one transport".into()); - } - - // Build DSN string from transports - let mut dsn_parts = Vec::new(); - - for transport in &config.transports { - // Validate transport type - let transport_type_lower = transport.transport_type.to_lowercase(); - if !["grpc", "http1", "dns"].contains(&transport_type_lower.as_str()) { - return Err(format!( - "Invalid transport type '{}'. Must be one of: GRPC, http1, DNS", - transport.transport_type - ) - .into()); - } - - // Validate that extra is valid JSON - if !transport.extra.is_empty() { - serde_json::from_str::(&transport.extra).map_err(|e| { - format!( - "Invalid JSON in 'extra' field for transport '{}': {}", - transport.uri, e - ) - })?; - } - - // Error if URI already contains query parameters - if transport.uri.contains('?') { - return Err(format!("URI '{}' already contains query parameters. Query parameters should not be present in the URI field.", transport.uri).into()); - } - - // Build DSN part with correct schema and query parameters - let mut dsn_part = transport.uri.clone(); - - // Add query parameters - dsn_part.push('?'); - let mut params = Vec::new(); - - // Add interval if present - if let Some(interval) = transport.interval { - params.push(format!("interval={}", interval)); - } - - // Add jitter if present - if let Some(jitter) = transport.jitter { - params.push(format!("jitter={}", jitter)); - } - - // Add extra as query parameter if not empty - if !transport.extra.is_empty() { - let encoded_extra = urlencoding::encode(&transport.extra); - params.push(format!("extra={}", encoded_extra)); - } - - if !params.is_empty() { - dsn_part.push_str(¶ms.join("&")); - } else { - // Remove the trailing '?' if no params were added - dsn_part.pop(); - } - - dsn_parts.push(dsn_part); - } - - // Join all DSN parts with semicolons - let dsn = dsn_parts.join(";"); - - // Emit the DSN configuration - println!("cargo:rustc-env=IMIX_CALLBACK_URI={}", dsn); - - // Emit server_pubkey if present - if let Some(ref pubkey) = config.server_pubkey { - println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey); - println!("cargo:warning=Using server_pubkey from YAML config"); - } - - println!( - "cargo:warning=Successfully parsed YAML config with {} transport(s)", - config.transports.len() - ); - - // Extract the first transport URI for pubkey fetching - let upstream_uri = config.transports.first().map(|t| t.uri.clone()); - - Ok(Some(YamlConfigResult { - upstream_uri, - server_pubkey: config.server_pubkey, - })) -} - -fn get_pub_key(yaml_config: Option) { - // Check if server pubkey was provided via YAML config - if let Some(ref config) = yaml_config { - if config.server_pubkey.is_some() { - // Already emitted in parse_yaml_config, no need to fetch - println!("cargo:warning=Server pubkey provided via YAML config, skipping fetch"); - return; - } - } - - // Check if IMIX_SERVER_PUBKEY is already set via env var - if std::env::var("IMIX_SERVER_PUBKEY").is_ok() { - println!("cargo:warning=IMIX_SERVER_PUBKEY already set, skipping fetch"); - return; - } - - // Get the callback URI: prefer YAML config upstream, then env var, then default - let callback_uri = yaml_config - .and_then(|c| c.upstream_uri) - .or_else(|| std::env::var("IMIX_CALLBACK_URI").ok()) - .unwrap_or_else(|| "http://127.0.0.1:8000".to_string()); - - // Extract the first URI from semicolon-separated list and strip query parameters - let base_uri = callback_uri - .split(';') - .next() - .unwrap_or(&callback_uri) - .trim() - .split('?') - .next() - .unwrap_or(&callback_uri); - - // Construct the status endpoint URL - let status_url = format!("{}/status", base_uri); - - // Make a GET request to /status - let response = match reqwest::blocking::get(&status_url) { - Ok(resp) => resp, - Err(e) => { - println!("cargo:warning=Failed to connect to {}: {}", status_url, e); - return; - } - }; - - if !response.status().is_success() { - println!( - "cargo:warning=Failed to fetch status from {}: HTTP {}", - status_url, - response.status() - ); - return; - } - - let json = match response.json::() { - Ok(json) => json, - Err(e) => { - println!( - "cargo:warning=Failed to parse JSON response from {}: {}", - status_url, e - ); - return; - } - }; - - let pubkey = match json.get("Pubkey").and_then(|v| v.as_str()) { - Some(key) => key, - None => { - println!( - "cargo:warning=Pubkey field not found in response from {}", - status_url - ); - return; - } - }; - - // Set the IMIX_SERVER_PUBKEY environment variable for the build - println!("cargo:rustc-env=IMIX_SERVER_PUBKEY={}", pubkey); - println!( - "cargo:warning=Successfully fetched server public key from {}", - status_url - ); -} - -fn validate_dsn_config() -> Result<(), Box> { - // Skip validation if YAML config is being used - // (parse_yaml_config already handles validation in that case) - if std::env::var("IMIX_CONFIG").is_ok() { - return Ok(()); - } - - // Check if IMIX_CALLBACK_URI contains query parameters - let callback_uri = - std::env::var("IMIX_CALLBACK_URI").unwrap_or_else(|_| "http://127.0.0.1:8000".to_string()); - let has_query_params = callback_uri.contains('?'); - - // Check if legacy config environment variables are set - let has_callback_interval = std::env::var("IMIX_CALLBACK_INTERVAL").is_ok(); - let has_transport_extra = std::env::vars().any(|(k, _)| k.starts_with("IMIX_TRANSPORT_EXTRA_")); - - // If DSN has query parameters AND legacy config is set, this is an error - if has_query_params && (has_callback_interval || has_transport_extra) { - let mut error_msg = String::from("Configuration error: Cannot use both DSN query parameters and legacy environment variables.\n"); - error_msg.push_str("Found query parameters in IMIX_CALLBACK_URI and one or more of:\n"); - - if has_callback_interval { - error_msg.push_str(" - IMIX_CALLBACK_INTERVAL\n"); - } - if has_transport_extra { - error_msg.push_str(" - IMIX_TRANSPORT_EXTRA_*\n"); - } - - error_msg.push_str("\nPlease use ONLY DSN query parameters (e.g., https://example.com?interval=10&extra={...})\n"); - error_msg.push_str("OR use legacy environment variables, but not both."); - - return Err(error_msg.into()); - } - - Ok(()) -} - fn main() -> Result<(), Box> { - // Tell Cargo to rerun this build script if these env vars change - // This fixes the issue where changing IMIX_CONFIG doesn't trigger a rebuild - println!("cargo:rerun-if-env-changed=IMIX_CONFIG"); - println!("cargo:rerun-if-env-changed=IMIX_CALLBACK_URI"); - println!("cargo:rerun-if-env-changed=IMIX_CALLBACK_INTERVAL"); - println!("cargo:rerun-if-env-changed=IMIX_SERVER_PUBKEY"); println!("cargo:rerun-if-env-changed=PROTOC"); - // Parse YAML config if present (this will emit IMIX_CALLBACK_URI if successful) - let yaml_config = parse_yaml_config()?; - - // Validate DSN config (skips if YAML config was used) - validate_dsn_config()?; - - get_pub_key(yaml_config); - // Skip if no `protoc` can be found match env::var_os("PROTOC") .map(PathBuf::from) diff --git a/implants/lib/pb/src/lib.rs b/implants/lib/pb/src/lib.rs index d32c80814..3e55c8af4 100644 --- a/implants/lib/pb/src/lib.rs +++ b/implants/lib/pb/src/lib.rs @@ -13,5 +13,4 @@ pub mod portal { pub mod trace { include!("generated/trace.rs"); } -pub mod config; pub mod xchacha; diff --git a/implants/lib/pb/src/xchacha.rs b/implants/lib/pb/src/xchacha.rs index d9a8665dd..abe6a5510 100644 --- a/implants/lib/pb/src/xchacha.rs +++ b/implants/lib/pb/src/xchacha.rs @@ -1,37 +1,30 @@ use anyhow::{Context, Result}; use bytes::{Buf, BufMut}; use chacha20poly1305::{aead::generic_array::GenericArray, aead::Aead, AeadCore, KeyInit}; -#[cfg(feature = "imix")] -use const_decoder::Decoder as const_decode; use lru::LruCache; use prost::Message; use rand::rngs::OsRng; use rand_chacha::rand_core::SeedableRng; +use std::sync::OnceLock; use std::{ io::{Read, Write}, marker::PhantomData, num::NonZeroUsize, - sync::{Mutex, OnceLock}, + sync::Mutex, }; use tonic::{ codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder}, Status, }; use x25519_dalek::{EphemeralSecret, PublicKey}; - -/* Compile-time constant for the server pubkey, derived from the IMIX_SERVER_PUBKEY environment variable during compilation. - * To find the servers pubkey check the startup messages on the server look for `[INFO] Public key: ` - */ -#[cfg(feature = "imix")] -static SERVER_PUBKEY: [u8; 32] = const_decode::Base64.decode(env!("IMIX_SERVER_PUBKEY").as_bytes()); - -#[cfg(not(feature = "imix"))] -static SERVER_PUBKEY: [u8; 32] = [ - 165, 30, 122, 188, 50, 89, 111, 214, 247, 4, 189, 217, 188, 37, 200, 190, 2, 180, 175, 107, - 194, 147, 177, 98, 103, 84, 99, 120, 72, 73, 87, 37, -]; // ------------ +static SERVER_PUBKEY: OnceLock<[u8; 32]> = OnceLock::new(); + +pub fn init_server_pubkey(key: [u8; 32]) { + let _ = SERVER_PUBKEY.set(key); +} + const KEY_CACHE_SIZE: usize = 1024; const PUBKEY_LEN: usize = 32; const NONCE_LEN: usize = 24; @@ -60,7 +53,8 @@ fn get_key(pub_key: [u8; 32]) -> Result<[u8; 32]> { fn encrypt_impl(pt_vec: Vec) -> Result> { // Store server pubkey - let server_public: PublicKey = PublicKey::from(SERVER_PUBKEY); + let server_pubkey = SERVER_PUBKEY.get().expect("SERVER_PUBKEY not initialized"); + let server_public: PublicKey = PublicKey::from(*server_pubkey); // Generate ephemeral keys let rng = rand_chacha::ChaCha20Rng::from_entropy(); diff --git a/implants/lib/transport/Cargo.toml b/implants/lib/transport/Cargo.toml index e4779c0f5..11c9fe90a 100644 --- a/implants/lib/transport/Cargo.toml +++ b/implants/lib/transport/Cargo.toml @@ -12,6 +12,7 @@ dns = ["dep:base32", "dep:rand", "dep:hickory-resolver", "dep:url"] mock = ["dep:mockall"] [dependencies] +config = { workspace = true } pb = { workspace = true } serde_json = { workspace = true } diff --git a/implants/lib/transport/src/dns.rs b/implants/lib/transport/src/dns.rs index 569dee331..d39bae8ec 100644 --- a/implants/lib/transport/src/dns.rs +++ b/implants/lib/transport/src/dns.rs @@ -1,8 +1,8 @@ use crate::Transport; use anyhow::{Context, Result}; +use config::Config; use hickory_resolver::system_conf::read_system_conf; use pb::c2::*; -use pb::config::Config; use pb::dns::*; use prost::Message; use std::sync::mpsc::{Receiver, Sender}; diff --git a/implants/lib/transport/src/grpc.rs b/implants/lib/transport/src/grpc.rs index ff128699a..806012929 100644 --- a/implants/lib/transport/src/grpc.rs +++ b/implants/lib/transport/src/grpc.rs @@ -1,7 +1,7 @@ use anyhow::Result; +use config::Config; use http::Uri; use pb::c2::*; -use pb::config::Config; use std::str::FromStr; use std::sync::mpsc::{Receiver, Sender}; use tonic::GrpcMethod; diff --git a/implants/lib/transport/src/http.rs b/implants/lib/transport/src/http.rs index 4c9339db9..30e219c20 100644 --- a/implants/lib/transport/src/http.rs +++ b/implants/lib/transport/src/http.rs @@ -1,9 +1,10 @@ use crate::Transport; use anyhow::{Context, Result}; use bytes::BytesMut; +use config::Config; use hyper_legacy::body::HttpBody; use hyper_legacy::{StatusCode, Uri}; -use pb::{c2::*, config::Config}; +use pb::c2::*; use prost::Message; use std::sync::{ mpsc::{Receiver, Sender}, diff --git a/implants/lib/transport/src/lib.rs b/implants/lib/transport/src/lib.rs index 3b04b5a9c..1d19b2a8f 100644 --- a/implants/lib/transport/src/lib.rs +++ b/implants/lib/transport/src/lib.rs @@ -1,8 +1,14 @@ use anyhow::{anyhow, Result}; +use config::Config; use pb::c2::transport::Type as TransportType; -use pb::{c2::*, config::Config}; +use pb::c2::*; use std::sync::mpsc::{Receiver, Sender}; +#[cfg(test)] +pub fn ensure_test_pubkey() { + pb::xchacha::init_server_pubkey([0u8; 32]); +} + #[cfg(any(feature = "grpc", feature = "http1"))] mod tls_utils; diff --git a/implants/lib/transport/src/mock.rs b/implants/lib/transport/src/mock.rs index 79279fc8a..0ba48d92e 100644 --- a/implants/lib/transport/src/mock.rs +++ b/implants/lib/transport/src/mock.rs @@ -12,7 +12,7 @@ mock! { impl super::Transport for Transport { fn init() -> Self; - fn new(config: pb::config::Config) -> Result; + fn new(config: config::Config) -> Result; async fn claim_tasks(&mut self, request: ClaimTasksRequest) -> Result; diff --git a/implants/lib/transport/src/transport.rs b/implants/lib/transport/src/transport.rs index 2a366a6b3..b5ad827d7 100644 --- a/implants/lib/transport/src/transport.rs +++ b/implants/lib/transport/src/transport.rs @@ -1,5 +1,6 @@ use anyhow::Result; -use pb::{c2::*, config::Config}; +use config::Config; +use pb::c2::*; use std::collections::HashMap; use std::sync::mpsc::{Receiver, Sender};