diff --git a/docs/_docs/user-guide/eldritch.md b/docs/_docs/user-guide/eldritch.md index 77e269c55..cdba4c4cb 100644 --- a/docs/_docs/user-guide/eldritch.md +++ b/docs/_docs/user-guide/eldritch.md @@ -118,6 +118,19 @@ for user_home_dir in file.list("/home/"): --- +## Agent + +### agent.eval + +`agent.eval(script: str) -> None` + +The agent.eval method takes an arbitrary eldritch payload string and +executes it in the runtime environment of the executing tome. This means that +any `print`s or `eprint`s or output from the script will be merged with that +of the broader tome. + +--- + ## Assets ### assets.copy diff --git a/implants/lib/eldritch/src/agent/eval_impl.rs b/implants/lib/eldritch/src/agent/eval_impl.rs new file mode 100644 index 000000000..1e9b46c41 --- /dev/null +++ b/implants/lib/eldritch/src/agent/eval_impl.rs @@ -0,0 +1,54 @@ +use std::collections::HashMap; + +use crate::runtime::{Environment, Runtime}; +use anyhow::Result; +use pb::eldritch::Tome; + +pub fn eval(env: &Environment, script: String) -> Result<()> { + let tome = Tome { + eldritch: script, + parameters: HashMap::new(), + file_names: Vec::new(), + }; + Runtime::run(env, &tome) +} + +#[cfg(test)] +mod tests { + use std::{sync::mpsc::channel, time::Duration}; + + use crate::runtime::Message; + + use super::*; + use anyhow::Result; + + #[test] + fn test_eval() -> Result<()> { + let (tx, rx) = channel::(); + let test_env = Environment::mock(1, tx); + eval(&test_env, String::from("print(\"hi\")"))?; + let m = rx.recv_timeout(Duration::from_secs(3))?; + match m { + Message::ReportText(r) => { + let expected_output = String::from("hi\n"); + assert!( + r.text == expected_output, + "'{}' did not equal '{}'", + r.text, + expected_output + ); + Ok(()) + } + _ => Err(anyhow::anyhow!("recieved nontext ouput from print")), + }?; + Ok(()) + } + + #[test] + fn test_eval_fail() -> Result<()> { + let (tx, _rx) = channel::(); + let test_env = Environment::mock(1, tx); + assert!(eval(&test_env, String::from("blorp(\"hi\")")).is_err()); + Ok(()) + } +} diff --git a/implants/lib/eldritch/src/agent/mod.rs b/implants/lib/eldritch/src/agent/mod.rs new file mode 100644 index 000000000..2a4643581 --- /dev/null +++ b/implants/lib/eldritch/src/agent/mod.rs @@ -0,0 +1,29 @@ +mod eval_impl; + +use starlark::{ + environment::MethodsBuilder, + eval::Evaluator, + starlark_module, + values::{none::NoneType, starlark_value}, +}; + +/* + * Define our library for this module. + */ +crate::eldritch_lib!(AgentLibrary, "agent_library"); + +/* + * Below, we define starlark wrappers for all of our library methods. + * The functions must be defined here to be present on our library. + */ +#[starlark_module] +#[rustfmt::skip] +#[allow(clippy::needless_lifetimes, clippy::type_complexity, clippy::too_many_arguments)] +fn methods(builder: &mut MethodsBuilder) { + #[allow(unused_variables)] + fn eval(this: &AgentLibrary, starlark_eval: &mut Evaluator<'v, '_>, script: String) -> anyhow::Result { + let env = crate::runtime::Environment::from_extra(starlark_eval.extra)?; + eval_impl::eval(env, script)?; + Ok(NoneType{}) + } +} diff --git a/implants/lib/eldritch/src/lib.rs b/implants/lib/eldritch/src/lib.rs index b1b01d6ce..596b37df5 100644 --- a/implants/lib/eldritch/src/lib.rs +++ b/implants/lib/eldritch/src/lib.rs @@ -1,6 +1,7 @@ #![deny(warnings)] #![allow(clippy::needless_lifetimes) /*Appears necessary for starlark macros*/] +pub mod agent; pub mod assets; pub mod crypto; pub mod file; diff --git a/implants/lib/eldritch/src/runtime/environment.rs b/implants/lib/eldritch/src/runtime/environment.rs index e3c619f4b..364eb0dc0 100644 --- a/implants/lib/eldritch/src/runtime/environment.rs +++ b/implants/lib/eldritch/src/runtime/environment.rs @@ -33,6 +33,11 @@ impl Environment { self.tx.send(msg.into())?; Ok(()) } + + #[cfg(test)] + pub fn mock(id: i64, tx: Sender) -> Self { + Self { id, tx } + } } /* diff --git a/implants/lib/eldritch/src/runtime/eval.rs b/implants/lib/eldritch/src/runtime/eval.rs index 9c5f20d84..29ba9af8d 100644 --- a/implants/lib/eldritch/src/runtime/eval.rs +++ b/implants/lib/eldritch/src/runtime/eval.rs @@ -1,5 +1,6 @@ use super::drain::drain; use crate::{ + agent::AgentLibrary, assets::AssetsLibrary, crypto::CryptoLibrary, file::FileLibrary, @@ -62,7 +63,7 @@ pub async fn start(id: i64, tome: Tome) -> Runtime { log::info!("evaluating tome (task_id={})", id); // Run Tome - match run_impl(&env, &tome) { + match Runtime::run(&env, &tome) { Ok(_) => { #[cfg(debug_assertions)] log::info!("tome evaluation successful (task_id={})", id); @@ -118,20 +119,6 @@ pub async fn start(id: i64, tome: Tome) -> Runtime { } } -fn run_impl(env: &Environment, tome: &Tome) -> Result<()> { - let ast = Runtime::parse(tome).context("failed to parse tome")?; - let module = Runtime::alloc_module(tome).context("failed to allocate module")?; - let globals = Runtime::globals(); - let mut eval: Evaluator = Evaluator::new(&module); - eval.extra = Some(env); - eval.set_print_handler(env); - - match eval.eval_module(ast, &globals) { - Ok(_) => Ok(()), - Err(err) => Err(err.into_anyhow().context("failed to evaluate tome")), - } -} - /* * Eldritch Runtime * @@ -173,6 +160,7 @@ impl Runtime { const report: ReportLibrary = ReportLibrary; const regex: RegexLibrary = RegexLibrary; const http: HTTPLibrary = HTTPLibrary; + const agent: AgentLibrary = AgentLibrary; } GlobalsBuilder::extended_by(&[ @@ -195,6 +183,27 @@ impl Runtime { .build() } + /* + * Parse an Eldritch tome into a starlark Abstract Syntax Tree (AST) Module, + * then allocate a module for the tome, and finally use the passed Environment + * to evaluate the module+AST. + */ + pub fn run(env: &Environment, tome: &Tome) -> Result<()> { + let ast = Runtime::parse(tome).context("failed to parse tome")?; + let module = Runtime::alloc_module(tome).context("failed to allocate module")?; + let globals = Runtime::globals(); + let mut eval: Evaluator = Evaluator::new(&module); + eval.extra = Some(env); + eval.set_print_handler(env); + + match eval.eval_module(ast, &globals) { + Ok(_) => Ok(()), + Err(err) => Err(err + .into_anyhow() + .context("failed to evaluate eldritch script")), + } + } + /* * Parse an Eldritch tome into a starlark Abstract Syntax Tree (AST) Module. */ diff --git a/implants/lib/eldritch/src/runtime/mod.rs b/implants/lib/eldritch/src/runtime/mod.rs index 483024cf7..d3f97d28d 100644 --- a/implants/lib/eldritch/src/runtime/mod.rs +++ b/implants/lib/eldritch/src/runtime/mod.rs @@ -192,6 +192,16 @@ mod tests { want_text: format!("{}\n", r#"["download", "get", "post"]"#), want_error: None, }, + agent_bindings: TestCase { + id: 123, + tome: Tome { + eldritch: String::from("print(dir(agent))"), + parameters: HashMap::new(), + file_names: Vec::new(), + }, + want_text: format!("{}\n", r#"["eval"]"#), + want_error: None, + }, } #[tokio::test(flavor = "multi_thread", worker_threads = 128)]