Adding agent.set_callback_uri method

docs/_docs/user-guide/eldritch.md

@@ -138,6 +138,16 @@ The <b>agent.set_callback_interval</b> method takes an unsigned int and changes
 running agent's callback interval to the passed value. This configuration change will
 not persist across agent reboots.
 
+### agent.set_callback_uri
+
+`agent.set_callback_uri(new_uri: str) -> None`
+
+The <b>agent.set_callback_uri</b> method takes an string and changes the
+running agent's callback uri to the passed value. This configuration change will
+not persist across agent reboots. NOTE: please ensure the passed URI path is correct
+for the underlying `Transport` being used, as a URI can take many forms and we make no
+assumptions on `Transport` requirements no gut checks are applied to the passed string.
+
 ---
 
 ## Assets

implants/lib/eldritch/src/agent/mod.rs

@@ -1,5 +1,6 @@
 mod eval_impl;
 mod set_callback_interval_impl;
+mod set_callback_uri_impl;
 
 use starlark::{
     environment::MethodsBuilder,
@@ -33,4 +34,10 @@ fn methods(builder: &mut MethodsBuilder) {
         set_callback_interval_impl::set_callback_interval(env, new_interval)?;
         Ok(NoneType{})
     }
+    #[allow(unused_variables)]
+    fn set_callback_uri(this: &AgentLibrary, starlark_eval: &mut Evaluator<'v, '_>, new_uri: String) -> anyhow::Result<NoneType> {
+        let env = crate::runtime::Environment::from_extra(starlark_eval.extra)?;
+        set_callback_uri_impl::set_callback_uri(env, new_uri)?;
+        Ok(NoneType{})
+    }
 }

implants/lib/eldritch/src/agent/set_callback_uri_impl.rs

@@ -0,0 +1,60 @@
+use crate::runtime::{messages::SetCallbackUriMessage, messages::SyncMessage, Environment};
+use anyhow::Result;
+
+pub fn set_callback_uri(env: &Environment, new_uri: String) -> Result<()> {
+    env.send(SyncMessage::from(SetCallbackUriMessage {
+        id: env.id(),
+        new_uri,
+    }))?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod test {
+    use crate::runtime::{messages::SyncMessage, Message};
+    use pb::eldritch::Tome;
+    use std::collections::HashMap;
+
+    macro_rules! test_cases {
+        ($($name:ident: $value:expr,)*) => {
+        $(
+            #[tokio::test]
+            async fn $name() {
+                let tc: TestCase = $value;
+
+                // Run Eldritch (until finished)
+                let mut runtime = crate::start(tc.id, tc.tome).await;
+                runtime.finish().await;
+
+                // Read Messages
+                let mut found = false;
+                for msg in runtime.messages() {
+                    if let Message::Sync(SyncMessage::SetCallbackUri(m)) = msg {
+                        assert_eq!(tc.new_uri, m.new_uri);
+                        found = true;
+                    }
+                }
+                assert!(found);
+            }
+        )*
+        }
+    }
+
+    struct TestCase {
+        pub id: i64,
+        pub tome: Tome,
+        pub new_uri: String,
+    }
+
+    test_cases! {
+            change_interval: TestCase{
+                id: 123,
+                tome: Tome{
+                    eldritch: String::from(r#"agent.set_callback_uri("https://127.0.0.1")"#),
+                    parameters: HashMap::new(),
+                    file_names: Vec::new(),
+                },
+                new_uri: String::from("https://127.0.0.1"),
+            },
+    }
+}

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

@@ -10,6 +10,7 @@ mod report_start;
 mod report_text;
 mod reverse_shell_pty;
 mod set_callback_interval;
+mod set_callback_uri;
 
 pub use fetch_asset::FetchAssetMessage;
 pub use pb::config::Config;
@@ -23,6 +24,7 @@ pub use report_start::ReportStartMessage;
 pub use report_text::ReportTextMessage;
 pub use reverse_shell_pty::ReverseShellPTYMessage;
 pub use set_callback_interval::SetCallbackIntervalMessage;
+pub use set_callback_uri::SetCallbackUriMessage;
 pub use transport::Transport;
 
 use anyhow::Result;
@@ -121,6 +123,9 @@ impl AsyncDispatcher for AsyncMessage {
 pub enum SyncMessage {
     #[display(fmt = "SetCallbackInterval")]
     SetCallbackInterval(SetCallbackIntervalMessage),
+
+    #[display(fmt = "SetCallbackUri")]
+    SetCallbackUri(SetCallbackUriMessage),
 }
 
 impl SyncDispatcher for SyncMessage {
@@ -130,6 +135,7 @@ impl SyncDispatcher for SyncMessage {
 
         match self {
             Self::SetCallbackInterval(msg) => msg.dispatch(transport, cfg),
+            Self::SetCallbackUri(msg) => msg.dispatch(transport, cfg),
         }
     }
 }

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

@@ -0,0 +1,21 @@
+use super::{SyncDispatcher, Transport};
+use anyhow::Result;
+use pb::config::Config;
+
+/*
+ * SetCallbackUriMessage sets the callback URI in the dispatched config.
+ */
+#[cfg_attr(debug_assertions, derive(Debug, PartialEq))]
+#[derive(Clone)]
+pub struct SetCallbackUriMessage {
+    pub(crate) id: i64,
+    pub(crate) new_uri: String,
+}
+
+impl SyncDispatcher for SetCallbackUriMessage {
+    fn dispatch(self, _transport: &mut impl Transport, cfg: Config) -> Result<Config> {
+        let mut c = cfg.clone();
+        c.callback_uri = self.new_uri;
+        Ok(c)
+    }
+}

implants/lib/eldritch/src/runtime/mod.rs

@@ -204,7 +204,7 @@ mod tests {
                 parameters: HashMap::new(),
                 file_names: Vec::new(),
             },
-            want_text: format!("{}\n", r#"["eval", "set_callback_interval"]"#),
+            want_text: format!("{}\n", r#"["eval", "set_callback_interval", "set_callback_uri"]"#),
             want_error: None,
         },
     }