Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 38 additions & 2 deletions implants/imix/src/shell/manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)) => {
Expand Down Expand Up @@ -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::*;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ mod int;
mod len;
mod libs;
mod ord;
mod pprint;
pub mod pprint;
mod print;
mod range;
mod str;
Expand All @@ -36,7 +36,7 @@ mod min;
mod repr;
mod reversed;
mod set;
mod tprint;
pub mod tprint;
mod tuple;
mod zip;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,10 @@ pub fn builtin_pprint(env: &Arc<RwLock<Environment>>, 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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,26 @@ pub fn builtin_tprint(env: &Arc<RwLock<Environment>>, 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<Option<String>, String> {
let items_snapshot: Vec<Value> = 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)
Expand All @@ -47,7 +59,7 @@ pub fn builtin_tprint(env: &Arc<RwLock<Environment>>, args: &[Value]) -> Result<
}

if columns.is_empty() {
return Ok(Value::None);
return Ok(None);
}

let columns_vec: Vec<String> = columns.into_iter().collect();
Expand Down Expand Up @@ -98,7 +110,5 @@ pub fn builtin_tprint(env: &Arc<RwLock<Environment>>, args: &[Value]) -> Result<
output.push('\n');
}

env.read().printer.print_out(&Span::new(0, 0, 0), &output);

Ok(Value::None)
Ok(Some(output))
}
2 changes: 1 addition & 1 deletion implants/lib/eldritch/eldritch-core/src/interpreter/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
mod builtins;
pub mod builtins;
mod core;
pub mod error;
mod eval;
Expand Down
2 changes: 2 additions & 0 deletions implants/lib/eldritch/eldritch-core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down
2 changes: 1 addition & 1 deletion implants/lib/eldritch/eldritch/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down