Add transport type to agent config and ent

implants/lib/eldritch/src/runtime/messages/set_callback_interval.rs

@@ -14,7 +14,7 @@ pub struct SetCallbackIntervalMessage {
 }
 
 impl SyncDispatcher for SetCallbackIntervalMessage {
-    fn dispatch(self, _transport: &mut impl Transport, cfg: Config) -> Result<Config> {
+    fn dispatch(self, transport: &mut impl Transport, cfg: Config) -> Result<Config> {
         let mut c = cfg.clone();
         let b = match cfg.info {
             Some(i) => Ok(i),
@@ -28,6 +28,7 @@ impl SyncDispatcher for SetCallbackIntervalMessage {
             host: b.host,
             agent: b.agent,
             interval: self.new_interval,
+            transport: transport.get_type() as i32,
         });
         Ok(c)
     }

implants/lib/pb/Cargo.toml

@@ -4,7 +4,10 @@ version = "0.0.5"
 edition = "2021"
 
 [features]
+default = []
 imix = []
+grpc = []
+http1 = []
 
 
 [dependencies]

implants/lib/pb/src/config.rs

@@ -1,4 +1,7 @@
+use tonic::transport;
 use uuid::Uuid;
+
+use crate::c2::beacon::Transport;
 /// Config holds values necessary to configure an Agent.
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -102,6 +105,13 @@ impl Config {
         let beacon_id =
             std::env::var("IMIX_BEACON_ID").unwrap_or_else(|_| String::from(Uuid::new_v4()));
 
+        #[cfg(feature = "http1")]
+        let transport = crate::c2::beacon::Transport::Http1;
+        #[cfg(feature = "grpc")]
+        let transport = crate::c2::beacon::Transport::Grpc;
+        #[cfg(not(any(feature = "http1", feature = "grpc")))]
+        let transport = crate::c2::beacon::Transport::Unspecified;
+
         let info = crate::c2::Beacon {
             identifier: beacon_id,
             principal: whoami::username(),
@@ -114,6 +124,7 @@ impl Config {
                     5_u64
                 }
             },
+            transport: transport as i32,
             host: Some(host),
             agent: Some(agent),
         };

implants/lib/transport/Cargo.toml

@@ -5,9 +5,9 @@ edition = "2021"
 
 [features]
 default = []
-grpc = []
+grpc = ["pb/grpc"]
 grpc-doh = ["grpc", "dep:hickory-resolver"]
-http1 = []
+http1 = ["pb/http1"]
 mock = ["dep:mockall"]
 
 [dependencies]

implants/lib/transport/src/grpc.rs

@@ -194,6 +194,9 @@ impl Transport for GRPC {
         Ok(())
     }
 
+    fn get_type(&mut self) -> pb::c2::beacon::Transport {
+        return pb::c2::beacon::Transport::Grpc;
+    }
     fn is_active(&self) -> bool {
         self.grpc.is_some()
     }

implants/lib/transport/src/http.rs

@@ -422,6 +422,10 @@ impl Transport for HTTP {
         ))
     }
 
+    fn get_type(&mut self) -> pb::c2::beacon::Transport {
+        return pb::c2::beacon::Transport::Http1;
+    }
+
     fn is_active(&self) -> bool {
         !self.base_url.is_empty()
     }

implants/lib/transport/src/lib.rs

@@ -179,6 +179,18 @@ impl Transport for ActiveTransport {
         }
     }
 
+    fn get_type(&mut self) -> beacon::Transport {
+        match self {
+            #[cfg(feature = "grpc")]
+            Self::Grpc(t) => t.get_type(),
+            #[cfg(feature = "http1")]
+            Self::Http(t) => t.get_type(),
+            #[cfg(feature = "mock")]
+            Self::Mock(t) => t.get_type(),
+            Self::Empty => beacon::Transport::Unspecified,
+        }
+    }
+
     fn is_active(&self) -> bool {
         match self {
             #[cfg(feature = "grpc")]

implants/lib/transport/src/mock.rs

@@ -48,6 +48,9 @@ mock! {
             tx: tokio::sync::mpsc::Sender<ReverseShellResponse>,
         ) -> Result<()>;
 
+        fn get_type(&mut self) -> pb::c2::beacon::Transport {
+            return Transport::Unspecified
+        }
         fn is_active(&self) -> bool;
 
         fn name(&self) -> &'static str;

implants/lib/transport/src/transport.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use pb::c2::*;
+use pb::c2::{beacon, *};
 use std::sync::mpsc::{Receiver, Sender};
 
 #[trait_variant::make(Transport: Send)]
@@ -79,6 +79,8 @@ pub trait UnsafeTransport: Clone + Send {
         tx: tokio::sync::mpsc::Sender<ReverseShellResponse>,
     ) -> Result<()>;
 
+    #[allow(dead_code)]
+    fn get_type(&mut self) -> beacon::Transport;
     /// Returns true if the transport is fully initialized and active
     #[allow(dead_code)]
     fn is_active(&self) -> bool;

tavern/internal/c2/api_claim_tasks.go

@@ -164,6 +164,7 @@ func (srv *Server) ClaimTasks(ctx context.Context, req *c2pb.ClaimTasksRequest)
 		SetLastSeenAt(now).
 		SetNextSeenAt(now.Add(time.Duration(req.Beacon.Interval) * time.Second)).
 		SetInterval(req.Beacon.Interval).
+		SetTransport(req.Beacon.Transport).
 		OnConflict().
 		UpdateNewValues().
 		ID(ctx)