diff --git a/implants/imix/src/shell/manager.rs b/implants/imix/src/shell/manager.rs index 5e001f6b7..733793873 100644 --- a/implants/imix/src/shell/manager.rs +++ b/implants/imix/src/shell/manager.rs @@ -7,7 +7,7 @@ use tokio::sync::mpsc; use eldritch::agent::agent::Agent; use eldritch::assets::std::EmptyAssets; -use eldritch::{Interpreter, Printer, Span, Value}; +use eldritch::{Interpreter, Printer, Span, Value, format_tprint, pretty_format}; use eldritch_agent::Context; use pb::c2::{ ReportOutputRequest, ReportShellTaskOutputMessage, ShellTask, ShellTaskContext, @@ -327,7 +327,8 @@ impl ShellManager { Ok(Ok(value)) => { if !matches!(value, Value::None) { let ctx = context.lock().unwrap().clone(); - dispatch_output(agent, shell_id, &ctx, format!("{:?}\n", value), false); + let formatted = format_value_smart(&value); + dispatch_output(agent, shell_id, &ctx, formatted, false); } } Ok(Err(e)) => { @@ -372,6 +373,41 @@ impl ShellManager { } } } + +/// Formats a Value for REPL output using smart formatting: +/// - Dictionaries are pretty-printed (pprint style) +/// - Lists of dictionaries are table-printed (tprint style) +/// - Everything else uses the default Debug format +fn format_value_smart(value: &Value) -> String { + match value { + Value::Dictionary(_) => { + let mut buf = String::new(); + pretty_format(value, 0, 2, &mut buf); + buf.push('\n'); + buf + } + Value::List(l) => { + let items = l.read(); + let is_list_of_dicts = + !items.is_empty() && items.iter().all(|v| matches!(v, Value::Dictionary(_))); + drop(items); + if is_list_of_dicts { + match format_tprint(value) { + Ok(Some(table)) => table, + Ok(None) => format!("{:?}\n", value), + Err(_) => format!("{:?}\n", value), + } + } else { + let mut buf = String::new(); + pretty_format(value, 0, 2, &mut buf); + buf.push('\n'); + buf + } + } + _ => format!("{:?}\n", value), + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/mod.rs b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/mod.rs index a6c2e283b..956e54442 100644 --- a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/mod.rs +++ b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/mod.rs @@ -17,7 +17,7 @@ mod int; mod len; mod libs; mod ord; -mod pprint; +pub mod pprint; mod print; mod range; mod str; @@ -36,7 +36,7 @@ mod min; mod repr; mod reversed; mod set; -mod tprint; +pub mod tprint; mod tuple; mod zip; diff --git a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/pprint.rs b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/pprint.rs index 6808397c2..f85dcc231 100644 --- a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/pprint.rs +++ b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/pprint.rs @@ -32,7 +32,10 @@ pub fn builtin_pprint(env: &Arc>, args: &[Value]) -> Result< Ok(Value::None) } -fn pretty_format(val: &Value, current_indent: usize, indent_width: usize, buf: &mut String) { +/// Formats a value using pretty-printing into a string buffer. +/// This is the core formatting logic used by `pprint()` and can be +/// called directly to format values without printing. +pub fn pretty_format(val: &Value, current_indent: usize, indent_width: usize, buf: &mut String) { match val { Value::List(l) => { let list = l.read(); diff --git a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/tprint.rs b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/tprint.rs index 03fe7c4d1..ea4761c97 100644 --- a/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/tprint.rs +++ b/implants/lib/eldritch/eldritch-core/src/interpreter/builtins/tprint.rs @@ -13,14 +13,26 @@ pub fn builtin_tprint(env: &Arc>, args: &[Value]) -> Result< return Err("tprint() takes at least 1 argument".to_string()); } - let list_val = &args[0]; + let output = format_tprint(&args[0])?; + + if let Some(text) = output { + env.read().printer.print_out(&Span::new(0, 0, 0), &text); + } + + Ok(Value::None) +} + +/// Formats a list of dictionaries as a markdown table string. +/// Returns `Ok(None)` if the list or all dictionaries are empty. +/// Returns an error if the value is not a list of dictionaries. +pub fn format_tprint(list_val: &Value) -> Result, String> { let items_snapshot: Vec = match list_val { Value::List(l) => l.read().clone(), _ => return Err("tprint() argument must be a list of dictionaries".to_string()), }; if items_snapshot.is_empty() { - return Ok(Value::None); + return Ok(None); } // Collect all unique keys (columns) @@ -47,7 +59,7 @@ pub fn builtin_tprint(env: &Arc>, args: &[Value]) -> Result< } if columns.is_empty() { - return Ok(Value::None); + return Ok(None); } let columns_vec: Vec = columns.into_iter().collect(); @@ -98,7 +110,5 @@ pub fn builtin_tprint(env: &Arc>, args: &[Value]) -> Result< output.push('\n'); } - env.read().printer.print_out(&Span::new(0, 0, 0), &output); - - Ok(Value::None) + Ok(Some(output)) } diff --git a/implants/lib/eldritch/eldritch-core/src/interpreter/mod.rs b/implants/lib/eldritch/eldritch-core/src/interpreter/mod.rs index b0ed5b1e2..e8a1f6df2 100644 --- a/implants/lib/eldritch/eldritch-core/src/interpreter/mod.rs +++ b/implants/lib/eldritch/eldritch-core/src/interpreter/mod.rs @@ -1,4 +1,4 @@ -mod builtins; +pub mod builtins; mod core; pub mod error; mod eval; diff --git a/implants/lib/eldritch/eldritch-core/src/lib.rs b/implants/lib/eldritch/eldritch-core/src/lib.rs index 29065c9f7..3fdaa79e1 100644 --- a/implants/lib/eldritch/eldritch-core/src/lib.rs +++ b/implants/lib/eldritch/eldritch-core/src/lib.rs @@ -17,6 +17,8 @@ pub use analysis::find_node_at_offset; pub use ast::{ Argument, Environment, ExprKind, FStringSegment, ForeignValue, Param, Stmt, StmtKind, Value, }; +pub use interpreter::builtins::pprint::pretty_format; +pub use interpreter::builtins::tprint::format_tprint; pub use interpreter::{BufferPrinter, Interpreter, NoopPrinter, Printer, StdoutPrinter}; pub use lexer::Lexer; pub use token::{Span, TokenKind}; diff --git a/implants/lib/eldritch/eldritch/src/lib.rs b/implants/lib/eldritch/eldritch/src/lib.rs index f0a3ddd6b..5d4e0b3c2 100644 --- a/implants/lib/eldritch/eldritch/src/lib.rs +++ b/implants/lib/eldritch/eldritch/src/lib.rs @@ -25,7 +25,7 @@ pub use eldritch_repl as repl; // Re-export core types pub use eldritch_core::{ BufferPrinter, Environment, ForeignValue, Interpreter as CoreInterpreter, NoopPrinter, Printer, - Span, StdoutPrinter, TokenKind, Value, conversion, + Span, StdoutPrinter, TokenKind, Value, conversion, format_tprint, pretty_format, }; pub use eldritch_macros as macros;