diff --git a/docs/_docs/user-guide/imix.md b/docs/_docs/user-guide/imix.md index 7113cbda0..e32511149 100644 --- a/docs/_docs/user-guide/imix.md +++ b/docs/_docs/user-guide/imix.md @@ -145,10 +145,50 @@ The install subcommand executes embedded tomes similar to golem. It will loop through all embedded files looking for main.eldritch. Each main.eldritch will execute in a new thread. This is done to allow imix to install redundantly or install additional (non dependent) tools. -Installation scripts are specified in the `realm/implants/imix/install_scripts` directory. +Installation scripts are specified in the `realm/implants/imixv2/embedded` directory. This feature is currently under active development, and may change. We'll do our best to keep these docs updates in the meantime. +## Events + +Imix supports an event system that allows executing Eldritch scripts when specific internal events occur. This feature is enabled by compiling with the `events` feature flag (enabled by default). + +The system looks for a universal event script at `on_event.eldritch` within the embedded assets. If found, this script is executed for every triggered event. + +The script receives a global variable `input_params` containing: +- `event`: The name of the event (e.g., "on_start", "on_callback_fail"). +- `args`: A dictionary of event-specific arguments. + +### Supported Events + +| Event Name | Description | Arguments | +|String | String | Map | +|---|---|---| +| `on_start` | Triggered when the agent initializes. | None | +| `on_exit` | Triggered when the agent shuts down. | None | +| `on_callback_success` | Triggered after a successful callback to the C2 server. | None | +| `on_callback_fail` | Triggered when a callback fails. | `error`: The error message. | +| `on_sigint` | Triggered on SIGINT (Ctrl+C). | None | +| `on_sigterm` | Triggered on SIGTERM. | None | +| `on_sighup` | Triggered on SIGHUP. | None | +| `on_sigquit` | Triggered on SIGQUIT. | None | +| `on_sigusr1` | Triggered on SIGUSR1. | None | +| `on_sigusr2` | Triggered on SIGUSR2. | None | +| `on_sigchild` | Triggered on SIGCHLD. | None | + +### Example Script + +Create `realm/implants/imixv2/embedded/on_event.eldritch`: + +```python +HANDLED_EVENTS = set(["on_callback_fail", "on_sigint"]) + +evt = input_params.get("event", "???") +if evt in HANDLED_EVENTS: + print(f"[EVENT] {evt} called:", input_params.get("args", {})) +``` + + ## Functionality Imix derives all it's functionality from the eldritch language. diff --git a/implants/Cargo.toml b/implants/Cargo.toml index fe3302bab..6bdc67437 100644 --- a/implants/Cargo.toml +++ b/implants/Cargo.toml @@ -28,7 +28,7 @@ members = [ "lib/eldritchv2/stdlib/tests", "lib/eldritchv2/stdlib/migration", "lib/eldritchv2/eldritchv2", - "lib/portals/portal-stream", + "lib/portals/portal-stream", "lib/eldritchv2/stdlib/eldritch-libevents", ] resolver = "2" @@ -57,6 +57,7 @@ eldritch-libregex = {path = "lib/eldritchv2/stdlib/eldritch-libregex",default-fe eldritch-libreport = {path = "lib/eldritchv2/stdlib/eldritch-libreport",default-features = false } eldritch-libsys = {path = "lib/eldritchv2/stdlib/eldritch-libsys",default-features = false } eldritch-libtime = {path = "lib/eldritchv2/stdlib/eldritch-libtime",default-features = false } +eldritch-libevents = {path = "lib/eldritchv2/stdlib/eldritch-libevents", default-features = false } portal-stream = { path = "lib/portals/portal-stream" } aes = "0.8.3" diff --git a/implants/imixv2/Cargo.toml b/implants/imixv2/Cargo.toml index 79b38cf3c..ba8c5604a 100644 --- a/implants/imixv2/Cargo.toml +++ b/implants/imixv2/Cargo.toml @@ -7,13 +7,14 @@ edition = "2024" crate-type = ["cdylib"] [features] -default = ["install", "grpc", "http1", "dns", "doh"] +default = ["install", "grpc", "http1", "dns", "doh", "events"] grpc = ["transport/grpc"] http1 = ["transport/http1"] dns = ["transport/dns"] doh = ["transport/doh"] win_service = [] install = [] +events = [] tokio-console = ["dep:console-subscriber", "tokio/tracing"] [dependencies] @@ -25,6 +26,7 @@ tokio = { workspace = true, features = [ "net", "io-util", "tracing", + "signal", ] } portal-stream = { workspace = true } anyhow = { workspace = true } @@ -40,6 +42,11 @@ pb = { workspace = true, features = ["imix"] } portable-pty = { workspace = true } rust-embed = { workspace = true } console-subscriber = { workspace = true, optional = true } +libc = { workspace = true } +nix = { workspace = true } + +[target.'cfg(target_os = "linux")'.dependencies] +signal-hook = { version = "0.3", features = ["extended-siginfo"] } [target.'cfg(target_os = "windows")'.dependencies] windows-service = { workspace = true } diff --git a/implants/imixv2/embedded/.gitkeep b/implants/imixv2/embedded/.gitkeep new file mode 100644 index 000000000..e69de29bb diff --git a/implants/imixv2/embedded/on_event.eldritch b/implants/imixv2/embedded/on_event.eldritch new file mode 100644 index 000000000..dae5dc683 --- /dev/null +++ b/implants/imixv2/embedded/on_event.eldritch @@ -0,0 +1,8 @@ +def on_callback_fail(args={}): + print("[on_callback_fail]", args) + +def on_callback_success(args={}): + print("[on_callback_success]", args) + +def on_start(args={}): + print("started") \ No newline at end of file diff --git a/implants/imixv2/src/assets.rs b/implants/imixv2/src/assets.rs index 24c7f952c..d9ac97ed5 100644 --- a/implants/imixv2/src/assets.rs +++ b/implants/imixv2/src/assets.rs @@ -1,5 +1,5 @@ use rust_embed::RustEmbed; #[derive(RustEmbed)] -#[folder = "./install_scripts"] +#[folder = "./embedded"] pub struct Asset; diff --git a/implants/imixv2/src/event.rs b/implants/imixv2/src/event.rs new file mode 100644 index 000000000..1a281ce30 --- /dev/null +++ b/implants/imixv2/src/event.rs @@ -0,0 +1,279 @@ +/* + * Event callbacks that let eldritch functions run when the implant does certain tasks + */ +#[cfg(feature = "events")] +use eldritchv2::{ + Interpreter, Value, + assets::std::{EmbeddedAssets, StdAssetsLibrary}, + conversion::ToValue, +}; +#[cfg(feature = "events")] +use std::{ + collections::{BTreeMap,BTreeSet}, + sync::{Arc, OnceLock}, +}; + +#[cfg(feature = "events")] +static EVENT_SCRIPT: OnceLock> = OnceLock::new(); +#[cfg(feature = "events")] +static EVENT_CALLBACKS: OnceLock> = OnceLock::new(); + +#[cfg(all(feature = "events", target_os = "linux"))] +#[derive(Debug, Clone)] +struct Sender { + pid: i32, + command: String, + args: Vec, + tty: String, + parent: Option>, +} + +#[cfg(all(feature = "events", target_os = "linux"))] +pub async fn catch_signals() { + use signal_hook::consts::{ + SIGCHLD, SIGHUP, SIGINT, SIGQUIT, SIGTERM, SIGUSR1, SIGUSR2, SIGWINCH, + }; + use signal_hook::iterator::SignalsInfo; + use signal_hook::iterator::exfiltrator::WithOrigin; + use signal_hook::iterator::exfiltrator::origin::Origin; + use tokio::sync::mpsc; + + // We send (Origin, Option) to capture proc info immediately in the signal thread + let (tx, mut rx) = mpsc::channel::<(Origin, Option)>(32); + + // Spawn a dedicated thread for signal handling because the iterator is blocking + std::thread::spawn(move || { + let mut signals = match SignalsInfo::::new(&[ + SIGINT, SIGTERM, SIGHUP, SIGQUIT, SIGUSR1, SIGUSR2, SIGCHLD, SIGWINCH, + ]) { + Ok(s) => s, + Err(e) => { + #[cfg(debug_assertions)] + eprintln!("Failed to register signals: {}", e); + return; + } + }; + + for info in &mut signals { + let mut sender = None; + // Attempt to read process info immediately to avoid race conditions + if let Some(p) = info.process { + sender = get_pid_info(p.pid, 15); + } + + if tx.blocking_send((info, sender)).is_err() { + break; + } + } + }); + + while let Some((info, sender)) = rx.recv().await { + let sig = info.signal as i32; + + let sig_name = match sig { + SIGINT => "sigint", + SIGTERM => "sigterm", + SIGHUP => "sighup", + SIGQUIT => "sigquit", + SIGUSR1 => "sigusr1", + SIGUSR2 => "sigusr2", + SIGCHLD => "sigchild", + SIGWINCH => "sigwinch", + _ => "unknown", + }; + + // Args that can passed to the event callback. "sender" + let mut args: BTreeMap = BTreeMap::new(); + args.insert("signal".into(), sig_name.to_string().to_value()); + if let Some(s) = sender { + args.insert("sender".into(), sender_to_value(&s)); + } + + // Actually spawn the event + tokio::spawn(async move { + on_event("on_signal", args).await; + }); + } +} + +#[cfg(all(feature = "events", not(target_os = "linux")))] +pub async fn catch_signals() {} + +#[cfg(all(feature = "events", target_os = "linux"))] +fn sender_to_value(sender: &Sender) -> Value { + let mut map = BTreeMap::new(); + map.insert("pid".to_string(), (sender.pid as i64).to_value()); + map.insert("command".to_string(), sender.command.clone().to_value()); + map.insert("args".to_string(), sender.args.clone().to_value()); + map.insert("tty".to_string(), sender.tty.clone().to_value()); + + if let Some(parent) = &sender.parent { + map.insert("parent".to_string(), sender_to_value(parent)); + } + + map.to_value() +} + +// Read the process calling tree for the kill command. Important if we want to take action +// against rough signals +#[cfg(all(feature = "events", target_os = "linux"))] +fn get_pid_info(pid: i32, max_rec: u32) -> Option { + if max_rec <= 0 { + return None; + } // Recursion loop + + // Read stat for comm, ppid, tty + use std::io::Read; + let path = format!("/proc/{}/stat", pid); + let mut buffer = [0u8; 512]; + let mut file = std::fs::File::open(&path).ok()?; + let n = file.read(&mut buffer).ok()?; + let content = &buffer[..n]; + + // Parse stat + let end_of_comm = content.iter().rposition(|&b| b == b')')?; + // comm is between the first '(' and the last ')' + let start_of_comm = content.iter().position(|&b| b == b'(')? + 1; + let command = String::from_utf8_lossy(&content[start_of_comm..end_of_comm]).into_owned(); + + if end_of_comm + 2 >= content.len() { + return None; + } + let rest = &content[end_of_comm + 2..]; + let mut iter = rest.split(|&b| b == b' '); + + iter.next(); // state + let ppid_bytes = iter.next()?; + let ppid_str = std::str::from_utf8(ppid_bytes).ok()?; + let ppid: i32 = ppid_str.parse().ok()?; + + iter.next(); // pgrp + iter.next(); // session + + let mut tty_str = String::new(); + if let Ok(link) = std::fs::read_link(format!("/proc/{}/fd/0", pid)) { + tty_str = link.to_string_lossy().into_owned() + } + + // Read cmdline for all args if the file still exists + let mut args = Vec::new(); + if let Ok(mut cmdline_file) = std::fs::File::open(format!("/proc/{}/cmdline", pid)) { + let mut cmdline_content = Vec::new(); + if cmdline_file.read_to_end(&mut cmdline_content).is_ok() { + args = cmdline_content + .split(|&b| b == 0) + .filter(|s| !s.is_empty()) + .map(|s| String::from_utf8_lossy(s).into_owned()) + .collect(); + } + } + + // Recurse for parent if depth allows + let parent = if ppid > 0 { + get_pid_info(ppid, max_rec - 1).map(Box::new) + } else { + None + }; + + Some(Sender { + pid, + command, + args, + tty: tty_str, + parent, + }) +} + +#[cfg(feature = "events")] +pub fn load_event_script() -> bool { + let script = EVENT_SCRIPT.get_or_init(|| { + let path = "on_event.eldritch"; + crate::assets::Asset::get(path).map(|f| String::from_utf8_lossy(&f.data).into_owned()) + }); + + // Use the lexer to get all the function definitions. This allows us + // to only load the interpreter if the callback is supported + if let Some(content) = script { + use eldritch_core::{ + Lexer, + TokenKind, + }; + + let mut lexer = Lexer::new(content.to_string()); + let tokens = lexer.scan_tokens(); + let mut is_func = false; + EVENT_CALLBACKS.get_or_init(|| { + let mut set = BTreeSet::new(); + for t in tokens { + if is_func && let TokenKind::Identifier(name) = &t.kind { + // Add this callback to to the set. It will be called later + if name.starts_with("on_") && name.len() > 3 { + set.insert(name.to_string()); + } + continue; + }; + // Track if the previous token was a func def + if let TokenKind::Def = &t.kind { + is_func = true; + } else { + is_func = false + } + } + #[cfg(debug_assertions)] + log::info!("loaded events from '{}': {:?}", "on_event.eldritch", set); + set + }); + return true; + }; + return false; +} + +#[cfg(feature = "events")] +pub async fn on_event(name: &str, args: BTreeMap) { + use eldritchv2::conversion::ToValue; + + // Dont run if there isnt a function for this callback + if let Some(callbacks) = EVENT_CALLBACKS.get() { + if !callbacks.contains(name) { + #[cfg(debug_assertions)] + log::info!("no callback registered for '{}'. Skipping", name); + return; + } + } + + // Check for the universal event script + let script_content = EVENT_SCRIPT.get().cloned().flatten(); + + let mut content = match script_content { + Some(s) => s, + None => return, + }; + + #[cfg(debug_assertions)] + log::info!( + "Running event script 'on_event.eldritch' for event: {}", + name + ); + + // Add event name to args + let mut input_params: BTreeMap = BTreeMap::new(); + input_params.insert("event".to_string(), name.to_string().to_value()); + input_params.insert("args".to_owned(), args.to_value()); + // Embedded assets are exposed to the callbacks + let asset_backend = Arc::new(EmbeddedAssets::::new()); + let mut locker = StdAssetsLibrary::new(); + let _ = locker.add(asset_backend); + // Execute using Eldritch V2 Interpreter + let mut interpreter = Interpreter::new().with_default_libs(); + interpreter.register_lib(locker); + interpreter.define_variable("input_params", input_params.to_value()); + + content = content + "\n" + name + "(input_params['args']);"; + match interpreter.interpret(&content) { + Ok(_) => {} + Err(_e) => { + #[cfg(debug_assertions)] + log::error!("Failed to execute event script: events/on_event.eldritch: {_e}"); + } + } +} diff --git a/implants/imixv2/src/lib.rs b/implants/imixv2/src/lib.rs index e2d176e7c..b64266474 100644 --- a/implants/imixv2/src/lib.rs +++ b/implants/imixv2/src/lib.rs @@ -2,6 +2,7 @@ extern crate alloc; pub mod agent; pub mod assets; +pub mod event; pub mod portal; pub mod run; pub mod shell; diff --git a/implants/imixv2/src/main.rs b/implants/imixv2/src/main.rs index 1b22f417a..824e77ec3 100644 --- a/implants/imixv2/src/main.rs +++ b/implants/imixv2/src/main.rs @@ -21,6 +21,8 @@ pub use transport::{ActiveTransport, Transport}; mod agent; mod assets; +#[cfg(feature = "events")] +mod event; mod install; mod portal; mod run; diff --git a/implants/imixv2/src/run.rs b/implants/imixv2/src/run.rs index 1f2441898..444a1b86f 100644 --- a/implants/imixv2/src/run.rs +++ b/implants/imixv2/src/run.rs @@ -1,4 +1,6 @@ use anyhow::Result; +#[cfg(feature = "events")] +use std::collections::BTreeMap; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; @@ -6,9 +8,14 @@ use std::time::{Duration, Instant}; use crate::agent::ImixAgent; use crate::task::TaskRegistry; use crate::version::VERSION; +#[cfg(feature = "events")] +use eldritchv2::conversion::ToValue; use pb::config::Config; use transport::{ActiveTransport, Transport}; +#[cfg(feature = "events")] +use crate::event; + pub static SHUTDOWN: AtomicBool = AtomicBool::new(false); pub async fn run_agent() -> Result<()> { @@ -39,6 +46,15 @@ pub async fn run_agent() -> Result<()> { #[cfg(debug_assertions)] log::info!("Agent initialized"); + // Run the onstart event script + #[cfg(feature = "events")] + { + if event::load_event_script() { + tokio::spawn(crate::event::catch_signals()); + tokio::spawn(event::on_event("on_start", BTreeMap::new())); + } + } + while !SHUTDOWN.load(Ordering::Relaxed) { let start = Instant::now(); let agent_ref = agent.clone(); @@ -65,6 +81,10 @@ pub async fn run_agent() -> Result<()> { } } + // Run the on_exit event script + #[cfg(feature = "events")] + event::on_event("on_exit", BTreeMap::new()).await; + #[cfg(debug_assertions)] log::info!("Agent shutting down"); @@ -90,9 +110,34 @@ async fn run_agent_cycle(agent: Arc>, registry: Arc 0 { + let t = &available_transports.transports + [available_transports.active_index as usize]; + uri = t.uri.clone(); + } + } + } + let mut args = BTreeMap::new(); + args.insert("uri".to_string(), uri.to_value()); + event::on_event("on_callback_start", args).await; + } let transport = match ActiveTransport::new(config) { Ok(t) => t, Err(_e) => { + #[cfg(feature = "events")] + { + let mut map = BTreeMap::new(); + map.insert("error".to_string(), _e.root_cause().to_string().to_value()); + tokio::spawn(event::on_event("on_callback_fail", map)); + } #[cfg(debug_assertions)] log::error!("Failed to create transport: {_e:#}"); agent.rotate_callback_uri().await; @@ -118,11 +163,21 @@ async fn process_tasks(agent: &ImixAgent, _registry: &TaskRegis Ok(_) => { #[cfg(debug_assertions)] log::info!("Callback success"); + + #[cfg(feature = "events")] + tokio::spawn(event::on_event("on_callback_success", BTreeMap::new())); } Err(_e) => { #[cfg(debug_assertions)] log::error!("Callback failed: {_e:#}"); agent.rotate_callback_uri().await; + + #[cfg(feature = "events")] + { + let mut map = BTreeMap::new(); + map.insert("error".to_string(), _e.root_cause().to_string().to_value()); + tokio::spawn(event::on_event("on_callback_fail", map)); + } } } } diff --git a/implants/lib/eldritchv2/eldritch-core/src/ast.rs b/implants/lib/eldritchv2/eldritch-core/src/ast.rs index fd0f94aa0..de1cfe813 100644 --- a/implants/lib/eldritchv2/eldritch-core/src/ast.rs +++ b/implants/lib/eldritchv2/eldritch-core/src/ast.rs @@ -66,6 +66,9 @@ pub trait ForeignValue: fmt::Debug + Send + Sync { args: &[Value], kwargs: &BTreeMap, ) -> Result; + fn get_attr(&self, _name: &str) -> Option { + None + } } #[derive(Clone)] diff --git a/implants/lib/eldritchv2/eldritch-core/src/interpreter/core.rs b/implants/lib/eldritchv2/eldritch-core/src/interpreter/core.rs index 8e412d246..a348cd7bb 100644 --- a/implants/lib/eldritchv2/eldritch-core/src/interpreter/core.rs +++ b/implants/lib/eldritchv2/eldritch-core/src/interpreter/core.rs @@ -1,4 +1,4 @@ -use super::super::ast::{BuiltinFn, Environment, Value}; +use super::super::ast::{Argument, BuiltinFn, Environment, Expr, ExprKind, Value}; use super::super::lexer::Lexer; use super::super::parser::Parser; use super::super::token::{Span, TokenKind}; @@ -300,6 +300,39 @@ impl Interpreter { self.env.write().values.insert(name.to_string(), value); } + pub fn call_value( + &mut self, + func: &Value, + args: &[Value], + kwargs: &BTreeMap, + ) -> Result { + // We use Argument::Positional for everything for now to simplify + let mut expr_args = Vec::new(); + for arg in args { + expr_args.push(Argument::Positional(Expr { + kind: ExprKind::Literal(arg.clone()), + span: Span::new(0, 0, 0), + })); + } + for (name, val) in kwargs { + expr_args.push(Argument::Keyword( + name.clone(), + Expr { + kind: ExprKind::Literal(val.clone()), + span: Span::new(0, 0, 0), + }, + )); + } + + let callee = Expr { + kind: ExprKind::Literal(func.clone()), + span: Span::new(0, 0, 0), + }; + + eval::functions::call_function(self, &callee, &expr_args, Span::new(0, 0, 0)) + .map_err(|e| e.to_string()) + } + pub fn lookup_variable(&self, name: &str, span: Span) -> Result { let mut current_env = Some(self.env.clone()); while let Some(env_arc) = current_env { diff --git a/implants/lib/eldritchv2/eldritch-core/src/interpreter/eval/access.rs b/implants/lib/eldritchv2/eldritch-core/src/interpreter/eval/access.rs index 84afa9202..7deb9c79e 100644 --- a/implants/lib/eldritchv2/eldritch-core/src/interpreter/eval/access.rs +++ b/implants/lib/eldritchv2/eldritch-core/src/interpreter/eval/access.rs @@ -273,7 +273,10 @@ pub(crate) fn evaluate_getattr( } // Support Foreign Objects - if let Value::Foreign(_) = &obj_val { + if let Value::Foreign(f) = &obj_val { + if let Some(val) = f.get_attr(&name) { + return Ok(val); + } // Return a bound method where the receiver is the foreign object return Ok(Value::BoundMethod(Box::new(obj_val), name)); } diff --git a/implants/lib/eldritchv2/eldritch-macros/src/impls.rs b/implants/lib/eldritchv2/eldritch-macros/src/impls.rs index d21cac271..b08dd3937 100644 --- a/implants/lib/eldritchv2/eldritch-macros/src/impls.rs +++ b/implants/lib/eldritchv2/eldritch-macros/src/impls.rs @@ -97,6 +97,22 @@ pub fn expand_eldritch_library( } }); + let has_get_attr = trait_def.items.iter().any(|item| { + if let TraitItem::Method(m) = item { + m.sig.ident == "_eldritch_get_attr" + } else { + false + } + }); + + if !has_get_attr { + trait_def.items.push(parse_quote! { + fn _eldritch_get_attr(&self, _name: &str) -> Option { + None + } + }); + } + trait_def.items.push(parse_quote! { fn _eldritch_call_method( &self, @@ -158,6 +174,10 @@ pub fn expand_eldritch_library_impl( ) -> Result { ::_eldritch_call_method(self, interp, name, args, kwargs) } + + fn get_attr(&self, name: &str) -> Option { + ::_eldritch_get_attr(self, name) + } } }) } diff --git a/implants/lib/eldritchv2/eldritchv2/Cargo.toml b/implants/lib/eldritchv2/eldritchv2/Cargo.toml index b52fb31e8..4067931ff 100644 --- a/implants/lib/eldritchv2/eldritchv2/Cargo.toml +++ b/implants/lib/eldritchv2/eldritchv2/Cargo.toml @@ -22,6 +22,7 @@ eldritch-libregex = { workspace = true, default-features = false } eldritch-libreport = { workspace = true, default-features = false } eldritch-libsys = { workspace = true, default-features = false } eldritch-libtime = { workspace = true, default-features = false } +eldritch-libevents = { workspace = true, default-features = false } pb = { workspace = true } eldritch-repl = { workspace = true, default-features = false } @@ -67,6 +68,7 @@ stdlib = [ "eldritch-libreport/stdlib", "eldritch-libsys/stdlib", "eldritch-libtime/stdlib", + "eldritch-libevents/stdlib", ] [dev-dependencies] diff --git a/implants/lib/eldritchv2/eldritchv2/src/events_test.rs b/implants/lib/eldritchv2/eldritchv2/src/events_test.rs new file mode 100644 index 000000000..f1ff30cca --- /dev/null +++ b/implants/lib/eldritchv2/eldritchv2/src/events_test.rs @@ -0,0 +1,58 @@ +use crate::Interpreter; +#[cfg(feature = "stdlib")] +use crate::agent::fake::AgentFake; +use eldritch_core::Value; +use std::sync::Arc; + +fn create_interp() -> Interpreter { + #[cfg(feature = "stdlib")] + { + use eldritch_libassets::std::EmptyAssets; + + let agent_mock = Arc::new(AgentFake); + let task_id = 123; + let backend = Arc::new(EmptyAssets {}); + Interpreter::new().with_default_libs().with_task_context( + agent_mock, + task_id, + vec![], + backend, + ) + } + #[cfg(not(feature = "stdlib"))] + { + Interpreter::new().with_default_libs() + } +} + +#[test] +fn test_events_constants() { + let mut interp = create_interp(); + + let val = interp.interpret("events.ON_CALLBACK_START").unwrap(); + assert_eq!(val, Value::String("on_callback_start".to_string())); + + let val = interp.interpret("events.ON_CALLBACK_END").unwrap(); + assert_eq!(val, Value::String("on_callback_end".to_string())); +} + +#[test] +fn test_events_register() { + let mut interp = create_interp(); + + let code = r#" +called = [False] +def my_hook(): + called[0] = True + +events.register(events.ON_CALLBACK_START, my_hook) +"#; + interp.interpret(code).unwrap(); + + // Manually trigger the event + use crate::events::std::trigger_event; + trigger_event(&mut interp.inner, "on_callback_start", vec![]); + + let val = interp.interpret("called[0]").unwrap(); + assert_eq!(val, Value::Bool(true)); +} diff --git a/implants/lib/eldritchv2/eldritchv2/src/lib.rs b/implants/lib/eldritchv2/eldritchv2/src/lib.rs index 890f89912..d21aeb1b3 100644 --- a/implants/lib/eldritchv2/eldritchv2/src/lib.rs +++ b/implants/lib/eldritchv2/eldritchv2/src/lib.rs @@ -18,12 +18,13 @@ pub use eldritch_libregex as regex; pub use eldritch_libreport as report; pub use eldritch_libsys as sys; pub use eldritch_libtime as time; +pub use eldritch_libevents as events; pub use eldritch_repl as repl; // Re-export core types pub use eldritch_core::{ BufferPrinter, Environment, ForeignValue, Interpreter as CoreInterpreter, Printer, Span, - StdoutPrinter, TokenKind, Value, conversion, + StdoutPrinter, TokenKind, Value, conversion, Parser }; pub use eldritch_macros as macros; @@ -65,6 +66,8 @@ use crate::report::std::StdReportLibrary; use crate::sys::std::StdSysLibrary; #[cfg(feature = "stdlib")] use crate::time::std::StdTimeLibrary; +#[cfg(feature = "stdlib")] +use crate::events::std::StdEventsLibrary; #[cfg(feature = "fake_agent")] use crate::agent::fake::AgentLibraryFake; @@ -90,6 +93,8 @@ use crate::report::fake::ReportLibraryFake; use crate::sys::fake::SysLibraryFake; #[cfg(feature = "fake_time")] use crate::time::fake::TimeLibraryFake; +#[cfg(feature = "fake_bindings")] +use crate::events::fake::EventsLibraryFake; pub struct Interpreter { inner: CoreInterpreter, @@ -126,6 +131,7 @@ impl Interpreter { self.inner.register_lib(StdRegexLibrary); self.inner.register_lib(StdSysLibrary); self.inner.register_lib(StdTimeLibrary); + self.inner.register_lib(StdEventsLibrary::new()); } #[cfg(feature = "fake_crypto")] @@ -258,3 +264,6 @@ mod bindings_test; #[cfg(all(test, feature = "fake_bindings"))] mod input_params_test; + +#[cfg(all(test, feature = "stdlib"))] +mod events_test; diff --git a/implants/lib/eldritchv2/stdlib/eldritch-libevents/Cargo.toml b/implants/lib/eldritchv2/stdlib/eldritch-libevents/Cargo.toml new file mode 100644 index 000000000..96dceed34 --- /dev/null +++ b/implants/lib/eldritchv2/stdlib/eldritch-libevents/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "eldritch-libevents" +version = "0.1.0" +edition = "2024" + +[dependencies] +eldritch-core = { workspace = true } +eldritch-macros = { workspace = true } +eldritch-agent = { workspace = true } +anyhow = { version = "1.0" } +rust-embed = { version = "8.0" } +spin = { workspace = true } + +[features] +default = ["std", "stdlib"] +std = ["eldritch-core/std"] +fake_bindings = [] +stdlib = [] diff --git a/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/fake.rs b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/fake.rs new file mode 100644 index 000000000..484659bc3 --- /dev/null +++ b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/fake.rs @@ -0,0 +1,19 @@ +use super::EventsLibrary; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use eldritch_core::Value; +use eldritch_macros::eldritch_library_impl; + +#[eldritch_library_impl(EventsLibrary)] +#[derive(Debug, Default)] +pub struct EventsLibraryFake; + +impl EventsLibrary for EventsLibraryFake { + fn list(&self) -> Result, String> { + Ok(alloc::vec!["fake_event".to_string()]) + } + + fn register(&self, _event: Value, _f: Value) -> Result<(), String> { + Ok(()) + } +} diff --git a/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/lib.rs b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/lib.rs new file mode 100644 index 000000000..260435fba --- /dev/null +++ b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/lib.rs @@ -0,0 +1,52 @@ +extern crate alloc; +use alloc::string::String; +use alloc::vec::Vec; +use eldritch_core::Value; +use eldritch_macros::{eldritch_library, eldritch_method}; + +#[cfg(feature = "fake_bindings")] +pub mod fake; +#[cfg(feature = "stdlib")] +pub mod std; + +pub const ON_CALLBACK_START: &str = "on_callback_start"; +pub const ON_CALLBACK_END: &str = "on_callback_end"; +pub const ON_TASK_START: &str = "on_task_start"; +pub const ON_TASK_END: &str = "on_task_end"; + +#[eldritch_library("events")] +/// The `events` library provides a mechanism for registering callbacks that are executed when specific agent events occur. +/// +/// This allows you to: +/// - Hook into the agent's lifecycle (e.g., before or after a callback). +/// - Monitor task execution. +/// - Implement custom logic in response to agent activities. +pub trait EventsLibrary { + #[eldritch_method] + /// Returns a list of all available events. + /// + /// **Returns** + /// - `List`: A list of event names that can be registered. + fn list(&self) -> Result, String>; + + #[eldritch_method] + /// Registers a callback function for a specific event. + /// + /// **Parameters** + /// - `event` (`str`): The name of the event (e.g., `events.ON_CALLBACK_START`). + /// - `f` (`function`): The callback function to execute. + /// + /// **Returns** + /// - `None` + fn register(&self, event: Value, f: Value) -> Result<(), String>; + + fn _eldritch_get_attr(&self, name: &str) -> Option { + match name { + "ON_CALLBACK_START" => Some(Value::String(ON_CALLBACK_START.to_string())), + "ON_CALLBACK_END" => Some(Value::String(ON_CALLBACK_END.to_string())), + "ON_TASK_START" => Some(Value::String(ON_TASK_START.to_string())), + "ON_TASK_END" => Some(Value::String(ON_TASK_END.to_string())), + _ => None, + } + } +} diff --git a/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/std/mod.rs b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/std/mod.rs new file mode 100644 index 000000000..7a6b51d3d --- /dev/null +++ b/implants/lib/eldritchv2/stdlib/eldritch-libevents/src/std/mod.rs @@ -0,0 +1,64 @@ +use super::{EventsLibrary, ON_CALLBACK_END, ON_CALLBACK_START, ON_TASK_END, ON_TASK_START}; +use alloc::collections::BTreeMap; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use eldritch_core::Value; +use eldritch_macros::eldritch_library_impl; +use spin::RwLock; + +// The global registry for event callbacks +pub static EVENT_REGISTRY: RwLock>> = RwLock::new(BTreeMap::new()); + +#[eldritch_library_impl(EventsLibrary)] +#[derive(Debug, Default)] +pub struct StdEventsLibrary {} + +impl StdEventsLibrary { + pub fn new() -> Self { + Self::default() + } +} + +impl EventsLibrary for StdEventsLibrary { + fn list(&self) -> Result, String> { + Ok(alloc::vec![ + ON_CALLBACK_START.to_string(), + ON_CALLBACK_END.to_string(), + ON_TASK_START.to_string(), + ON_TASK_END.to_string(), + ]) + } + + fn register(&self, event: Value, f: Value) -> Result<(), String> { + let event_name = match event { + Value::String(s) => s, + _ => return Err("Event name must be a string".to_string()), + }; + + match f { + Value::Function(_) | Value::NativeFunction(_, _) | Value::NativeFunctionWithKwargs(_, _) | Value::BoundMethod(_, _) => { + let mut registry = EVENT_REGISTRY.write(); + registry.entry(event_name).or_default().push(f); + Ok(()) + } + _ => Err("Callback must be a function".to_string()), + } + } +} + +/// Triggers an event and executes all registered callbacks. +pub fn trigger_event(interp: &mut eldritch_core::Interpreter, event: &str, args: Vec) { + let callbacks = { + let registry = EVENT_REGISTRY.read(); + registry.get(event).cloned() + }; + + if let Some(callbacks) = callbacks { + for callback in callbacks { + if let Err(e) = interp.call_value(&callback, &args, &BTreeMap::new()) { + let printer = interp.env.read().printer.clone(); + printer.print_err(&eldritch_core::Span::new(0, 0, 0), &alloc::format!("Event callback error for event '{}': {}", event, e)); + } + } + } +}