diff --git a/docs/_docs/dev-guide/imix.md b/docs/_docs/dev-guide/imix.md index 27c206b05..ff4f38206 100644 --- a/docs/_docs/dev-guide/imix.md +++ b/docs/_docs/dev-guide/imix.md @@ -86,3 +86,140 @@ pub use mac_address::MacAddress; ``` - Update the `defaults()` function to include your implementation. N.B. The order from left to right is the order engines will be evaluated. + +## Develop a New Transport + +We've tried to make Imix super extensible for transport development. In fact, all of the transport specific logic is complete abstracted from how Imix operates for callbacks/tome excution. For Imix all Transports live in the `realm/implants/lib/transport/src` directory. + +If creating a new Transport create a new file in the directory and name it after the protocol you plan to use. For example, if writing a DNS Transport then call your file `dns.rs`. Then define your public struct where any connection state/clients will be. For example, + +```rust +#[derive(Debug, Clone)] +pub struct DNS { + dns_client: Option +} +``` + +NOTE: Depending on the struct you build, you may need to derive certain features, see above we derive `Debug` and `Clone`. + +Next, we need to implement the Transport trait for our new struct. This will look like: + +```rust +impl Transport for DNS { + fn init() -> Self { + DNS{ dns_client: None } + } + fn new(callback: String, proxy_uri: Option) -> Result { + // TODO: setup connection/client hook in proxy, anything else needed + // before fuctions get called. + Err(anyhow!("Unimplemented!")) + } + async fn claim_tasks(&mut self, request: ClaimTasksRequest) -> Result { + // TODO: How you wish to handle the `claim_tasks` method. + Err(anyhow!("Unimplemented!")) + } + async fn fetch_asset( + &mut self, + request: FetchAssetRequest, + tx: std::sync::mpsc::Sender, + ) -> Result<()> { + // TODO: How you wish to handle the `fetch_asset` method. + Err(anyhow!("Unimplemented!")) + } + async fn report_credential( + &mut self, + request: ReportCredentialRequest, + ) -> Result { + // TODO: How you wish to handle the `report_credential` method. + Err(anyhow!("Unimplemented!")) + } + async fn report_file( + &mut self, + request: std::sync::mpsc::Receiver, + ) -> Result { + // TODO: How you wish to handle the `report_file` method. + Err(anyhow!("Unimplemented!")) + } + async fn report_process_list( + &mut self, + request: ReportProcessListRequest, + ) -> Result { + // TODO: How you wish to handle the `report_process_list` method. + Err(anyhow!("Unimplemented!")) + } + async fn report_task_output( + &mut self, + request: ReportTaskOutputRequest, + ) -> Result { + // TODO: How you wish to handle the `report_task_output` method. + Err(anyhow!("Unimplemented!")) + } + async fn reverse_shell( + &mut self, + rx: tokio::sync::mpsc::Receiver, + tx: tokio::sync::mpsc::Sender, + ) -> Result<()> { + // TODO: How you wish to handle the `reverse_shell` method. + Err(anyhow!("Unimplemented!")) + } +} +``` + +NOTE: Be Aware that currently `reverse_shell` uses tokio's sender/reciever while the rest of the methods rely on mpsc's. This is an artifact of some implementation details under the hood of Imix. Some day we may wish to move completely over to tokio's but currenlty it would just result in performance loss/less maintainable code. + +After you implement all the functions/write in a decent error message for operators to understad why the function call failed then you need to import the Transport to the broader lib scope. To do this open up `realm/implants/lib/transport/src/lib.rs` and add in your new Transport like so: + +```rust +// more stuff above + +#[cfg(feature = "dns")] +mod dns; +#[cfg(feature = "dns")] +pub use dns::DNS; + +// more stuff below +``` + +Also add your new feature to the Transport Cargo.toml at `realm/implants/lib/transport/Cargo.toml`. + +```toml +# more stuff above + +[features] +default = [] +grpc = [] +dns = [] # <-- see here +mock = ["dep:mockall"] + +# more stuff below +``` + +And that's it! Well, unless you want to _use_ the new transport. In which case you need to swap out the chosen transport being compiled for Imix in it's Cargo.toml (`/workspaces/realm/implants/lib/transport/Cargo.toml`) like so + +```toml +# more stuff above + +[dependencies] +eldritch = { workspace = true, features = ["imix"] } +pb = { workspace = true } +transport = { workspace = true, features = ["dns"] } # <-- see here +host_unique = { workspace = true } + +# more stuff below +``` + +Then just swap which Transport gets intialized on Imix's `run` function in run.rs (`/workspaces/realm/implants/imix/src/run.rs`) accordingly, + +```rust +// more stuff above + +async fn run(cfg: Config) -> anyhow::Result<()> { + let mut agent = Agent::new(cfg, DNS::init())?; // <-- changed this (also imported it) + agent.callback_loop().await?; + Ok(()) +} + +// more stuff below +``` + +And that's all that is needed for Imix to use a new Transport! Now all there is to do is setup some sort of tavern proxy for your new protocol and test! diff --git a/implants/imix/Cargo.toml b/implants/imix/Cargo.toml index 0f1dc3800..5e4960e00 100644 --- a/implants/imix/Cargo.toml +++ b/implants/imix/Cargo.toml @@ -14,7 +14,7 @@ default = [] [dependencies] eldritch = { workspace = true, features = ["imix"] } pb = { workspace = true } -transport = { workspace = true } +transport = { workspace = true, features = ["grpc"] } host_unique = { workspace = true } anyhow = { workspace = true } diff --git a/implants/imix/src/agent.rs b/implants/imix/src/agent.rs index ef22ea855..e3f814084 100644 --- a/implants/imix/src/agent.rs +++ b/implants/imix/src/agent.rs @@ -2,30 +2,32 @@ use crate::{config::Config, task::TaskHandle}; use anyhow::Result; use pb::c2::ClaimTasksRequest; use std::time::{Duration, Instant}; -use transport::{Transport, GRPC}; +use transport::Transport; /* * Agent contains all relevant logic for managing callbacks to a c2 server. * It is responsible for obtaining tasks, executing them, and returning their output. */ -pub struct Agent { +pub struct Agent { cfg: Config, handles: Vec, + t: T, } -impl Agent { +impl Agent { /* * Initialize an agent using the provided configuration. */ - pub fn new(cfg: Config) -> Result { + pub fn new(cfg: Config, t: T) -> Result { Ok(Agent { cfg, handles: Vec::new(), + t, }) } // Claim tasks and start their execution - async fn claim_tasks(&mut self, mut tavern: GRPC) -> Result<()> { + async fn claim_tasks(&mut self, mut tavern: T) -> Result<()> { let tasks = tavern .claim_tasks(ClaimTasksRequest { beacon: Some(self.cfg.info.clone()), @@ -56,7 +58,7 @@ impl Agent { } // Report task output, remove completed tasks - async fn report(&mut self, mut tavern: GRPC) -> Result<()> { + async fn report(&mut self, mut tavern: T) -> Result<()> { // Report output from each handle let mut idx = 0; while idx < self.handles.len() { @@ -79,9 +81,10 @@ impl Agent { * Callback once using the configured client to claim new tasks and report available output. */ pub async fn callback(&mut self) -> Result<()> { - let transport = GRPC::new(self.cfg.callback_uri.clone(), self.cfg.proxy_uri.clone())?; - self.claim_tasks(transport.clone()).await?; - self.report(transport.clone()).await?; + self.t = T::new(self.cfg.callback_uri.clone(), self.cfg.proxy_uri.clone())?; + self.claim_tasks(self.t.clone()).await?; + self.report(self.t.clone()).await?; + self.t = T::init(); // re-init to make sure no active connections during sleep Ok(()) } diff --git a/implants/imix/src/run.rs b/implants/imix/src/run.rs index e613c51b4..3a7ad1bcd 100644 --- a/implants/imix/src/run.rs +++ b/implants/imix/src/run.rs @@ -5,6 +5,8 @@ pub use crate::agent::Agent; pub use crate::config::Config; pub use crate::install::install; +use transport::{Transport, GRPC}; + pub async fn handle_main() { if let Some(("install", _)) = Command::new("imix") .subcommand(Command::new("install").about("Install imix")) @@ -34,7 +36,7 @@ pub async fn handle_main() { } async fn run(cfg: Config) -> anyhow::Result<()> { - let mut agent = Agent::new(cfg)?; + let mut agent = Agent::new(cfg, GRPC::init())?; agent.callback_loop().await?; Ok(()) } diff --git a/implants/lib/transport/Cargo.toml b/implants/lib/transport/Cargo.toml index c64ecf311..cad37d793 100644 --- a/implants/lib/transport/Cargo.toml +++ b/implants/lib/transport/Cargo.toml @@ -4,7 +4,7 @@ version = "0.0.5" edition = "2021" [features] -default = ["grpc"] +default = [] grpc = [] mock = ["dep:mockall"] diff --git a/implants/lib/transport/src/grpc.rs b/implants/lib/transport/src/grpc.rs index 8d8cb478f..d7200d1a4 100644 --- a/implants/lib/transport/src/grpc.rs +++ b/implants/lib/transport/src/grpc.rs @@ -20,10 +20,14 @@ static REVERSE_SHELL_PATH: &str = "/c2.C2/ReverseShell"; #[derive(Debug, Clone)] pub struct GRPC { - grpc: tonic::client::Grpc, + grpc: Option>, } impl Transport for GRPC { + fn init() -> Self { + GRPC { grpc: None } + } + fn new(callback: String, proxy_uri: Option) -> Result { let endpoint = tonic::transport::Endpoint::from_shared(callback)?; @@ -51,7 +55,7 @@ impl Transport for GRPC { }; let grpc = tonic::client::Grpc::new(channel); - Ok(Self { grpc }) + Ok(Self { grpc: Some(grpc) }) } async fn claim_tasks(&mut self, request: ClaimTasksRequest) -> Result { @@ -184,7 +188,13 @@ impl GRPC { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result, tonic::Status> { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -197,7 +207,7 @@ impl GRPC { let mut req = request.into_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ClaimTasks")); - self.grpc.unary(req, path, codec).await + self.grpc.as_mut().unwrap().unary(req, path, codec).await } /// @@ -215,7 +225,13 @@ impl GRPC { tonic::Response>, tonic::Status, > { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -227,7 +243,11 @@ impl GRPC { let mut req = request.into_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "FetchAsset")); - self.grpc.server_streaming(req, path, codec).await + self.grpc + .as_mut() + .unwrap() + .server_streaming(req, path, codec) + .await } /// @@ -236,7 +256,13 @@ impl GRPC { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result, tonic::Status> { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -249,7 +275,7 @@ impl GRPC { let mut req = request.into_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ReportCredential")); - self.grpc.unary(req, path, codec).await + self.grpc.as_mut().unwrap().unary(req, path, codec).await } /// @@ -264,7 +290,13 @@ impl GRPC { &mut self, request: impl tonic::IntoStreamingRequest, ) -> std::result::Result, tonic::Status> { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -276,7 +308,11 @@ impl GRPC { let mut req = request.into_streaming_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ReportFile")); - self.grpc.client_streaming(req, path, codec).await + self.grpc + .as_mut() + .unwrap() + .client_streaming(req, path, codec) + .await } /// @@ -286,7 +322,13 @@ impl GRPC { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result, tonic::Status> { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -298,7 +340,7 @@ impl GRPC { let mut req = request.into_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ReportProcessList")); - self.grpc.unary(req, path, codec).await + self.grpc.as_mut().unwrap().unary(req, path, codec).await } /// @@ -307,7 +349,13 @@ impl GRPC { &mut self, request: impl tonic::IntoRequest, ) -> std::result::Result, tonic::Status> { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -319,7 +367,7 @@ impl GRPC { let mut req = request.into_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ReportTaskOutput")); - self.grpc.unary(req, path, codec).await + self.grpc.as_mut().unwrap().unary(req, path, codec).await } async fn reverse_shell_impl( @@ -329,7 +377,13 @@ impl GRPC { tonic::Response>, tonic::Status, > { - self.grpc.ready().await.map_err(|e| { + if self.grpc.is_none() { + return Err(tonic::Status::new( + tonic::Code::FailedPrecondition, + "grpc client not created".to_string(), + )); + } + self.grpc.as_mut().unwrap().ready().await.map_err(|e| { tonic::Status::new( tonic::Code::Unknown, format!("Service was not ready: {}", e), @@ -341,6 +395,10 @@ impl GRPC { let mut req = request.into_streaming_request(); req.extensions_mut() .insert(GrpcMethod::new("c2.C2", "ReverseShell")); - self.grpc.streaming(req, path, codec).await + self.grpc + .as_mut() + .unwrap() + .streaming(req, path, codec) + .await } } diff --git a/implants/lib/transport/src/mock.rs b/implants/lib/transport/src/mock.rs index 0d32d7666..2a4434e09 100644 --- a/implants/lib/transport/src/mock.rs +++ b/implants/lib/transport/src/mock.rs @@ -10,40 +10,42 @@ mock! { fn clone(&self) -> Self; } impl super::Transport for Transport { - fn new(uri: String, proxy_uri: Option) -> Result; - - async fn claim_tasks(&mut self, request: ClaimTasksRequest) -> Result; - - async fn fetch_asset( - &mut self, - request: FetchAssetRequest, - sender: Sender, - ) -> Result<()>; - - async fn report_credential( - &mut self, - request: ReportCredentialRequest, - ) -> Result; - - async fn report_file( - &mut self, - request: Receiver, - ) -> Result; - - async fn report_process_list( - &mut self, - request: ReportProcessListRequest, - ) -> Result; - - async fn report_task_output( - &mut self, - request: ReportTaskOutputRequest, - ) -> Result; - - async fn reverse_shell( - &mut self, - rx: tokio::sync::mpsc::Receiver, - tx: tokio::sync::mpsc::Sender, - ) -> Result<()>; + fn init() -> Self; + + fn new(uri: String, proxy_uri: Option) -> Result; + + async fn claim_tasks(&mut self, request: ClaimTasksRequest) -> Result; + + async fn fetch_asset( + &mut self, + request: FetchAssetRequest, + sender: Sender, + ) -> Result<()>; + + async fn report_credential( + &mut self, + request: ReportCredentialRequest, + ) -> Result; + + async fn report_file( + &mut self, + request: Receiver, + ) -> Result; + + async fn report_process_list( + &mut self, + request: ReportProcessListRequest, + ) -> Result; + + async fn report_task_output( + &mut self, + request: ReportTaskOutputRequest, + ) -> Result; + + async fn reverse_shell( + &mut self, + rx: tokio::sync::mpsc::Receiver, + tx: tokio::sync::mpsc::Sender, + ) -> Result<()>; } } diff --git a/implants/lib/transport/src/transport.rs b/implants/lib/transport/src/transport.rs index 0ed86bbc1..842004e4f 100644 --- a/implants/lib/transport/src/transport.rs +++ b/implants/lib/transport/src/transport.rs @@ -4,7 +4,11 @@ use std::sync::mpsc::{Receiver, Sender}; #[trait_variant::make(Transport: Send)] pub trait UnsafeTransport: Clone + Send { - // New will initialize a new instance of the transport using the provided URI. + // Init will initialize a new instance of the transport with no active connections. + #[allow(dead_code)] + fn init() -> Self; + + // New will create a new instance of the transport using the provided URI. #[allow(dead_code)] fn new(uri: String, proxy_uri: Option) -> Result;