use crate::{Error, Result};
use rhai::{Engine, Scope};
pub struct ScriptEngine {
engine: Engine,
}
impl ScriptEngine {
pub fn new() -> Self {
let mut engine = Engine::new();
crate::harden(&mut engine, 100_000);
crate::functions::register_functions(&mut engine);
crate::types::register_types(&mut engine);
Self { engine }
}
pub fn validate(&self, script: &str, content: &str) -> Result<bool> {
let mut scope = Scope::new();
scope.push("content", content.to_string());
self.engine
.eval_with_scope::<bool>(&mut scope, script)
.map_err(|e| Error::ExecutionFailed(e.to_string()))
}
pub fn transform(&self, script: &str, input: rhai::Map) -> Result<String> {
let mut scope = Scope::new();
scope.push("input", input);
self.engine
.eval_with_scope::<String>(&mut scope, script)
.map_err(|e| Error::ExecutionFailed(e.to_string()))
}
pub fn execute(&self, script: &str, scope: &mut Scope) -> Result<rhai::Dynamic> {
self.engine
.eval_with_scope(scope, script)
.map_err(|e| Error::ExecutionFailed(e.to_string()))
}
pub fn check_gate_rule(
&self,
script: &str,
tool: &str,
target: Option<&str>,
taint_level: &str,
) -> Result<bool> {
let mut context = rhai::Map::new();
context.insert("tool".into(), rhai::Dynamic::from(tool.to_string()));
context.insert(
"target".into(),
rhai::Dynamic::from(target.unwrap_or("").to_string()),
);
context.insert(
"taint_level".into(),
rhai::Dynamic::from(taint_level.to_string()),
);
let mut scope = Scope::new();
scope.push("context", context);
self.engine
.eval_with_scope::<bool>(&mut scope, script)
.map_err(|e| Error::ExecutionFailed(e.to_string()))
}
}
impl Default for ScriptEngine {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_engine_creation() {
let engine = ScriptEngine::new();
assert!(engine.engine.max_operations() > 0);
}
#[test]
fn test_simple_validation() {
let engine = ScriptEngine::new();
let script = r#"
content.contains("test")
"#;
let result = engine.validate(script, "this is a test");
assert!(result.is_ok());
assert!(result.unwrap());
}
#[test]
fn test_string_operations() {
let engine = ScriptEngine::new();
let script = r#"content.starts_with("Hello")"#;
assert!(engine.validate(script, "Hello, world!").unwrap());
let script = r#"content.ends_with("!")"#;
assert!(engine.validate(script, "Hello, world!").unwrap());
let script = r#"content.trim() == "hello""#;
assert!(engine.validate(script, " hello ").unwrap());
}
#[test]
fn test_json_validation() {
let engine = ScriptEngine::new();
let script = r#"is_json(content)"#;
assert!(engine.validate(script, r#"{"key": "value"}"#).unwrap());
assert!(!engine.validate(script, "not json").unwrap());
}
#[test]
fn test_mermaid_validation() {
let engine = ScriptEngine::new();
let script = r#"is_mermaid(content)"#;
assert!(engine.validate(script, "graph TD\n A --> B").unwrap());
assert!(
engine
.validate(script, "sequenceDiagram\n Alice->>Bob: Hello")
.unwrap()
);
assert!(!engine.validate(script, "just text").unwrap());
}
#[test]
fn test_token_counting() {
let engine = ScriptEngine::new();
let script = r#"count_tokens(content) > 10"#;
let long_text = "a".repeat(50);
assert!(engine.validate(script, &long_text).unwrap());
assert!(!engine.validate(script, "short").unwrap());
}
#[test]
fn test_split_join() {
let engine = ScriptEngine::new();
let script = r#"
let parts = content.split(",");
parts.len() == 3 && join(parts, "|") != ""
"#;
assert!(engine.validate(script, "a,b,c").unwrap());
}
#[test]
fn test_sandbox_limits() {
let engine = ScriptEngine::new();
let script = r#"
let x = 0;
loop {
x = x + 1;
if x > 200000 { break; }
}
true
"#;
let result = engine.validate(script, "test");
assert!(result.is_err());
}
#[test]
fn test_transform_success() {
let engine = ScriptEngine::new();
let mut input = rhai::Map::new();
input.insert("name".into(), rhai::Dynamic::from("world".to_string()));
let script = r#"("hello " + input["name"])"#;
let result = engine.transform(script, input);
assert_eq!(result.unwrap(), "hello world");
}
#[test]
fn test_transform_script_error_returns_execution_failed() {
let engine = ScriptEngine::new();
let input = rhai::Map::new();
let script = "this is not valid rhai {{{";
let result = engine.transform(script, input);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.starts_with("Script execution failed:")
);
}
#[test]
fn test_transform_wrong_return_type_returns_execution_failed() {
let engine = ScriptEngine::new();
let input = rhai::Map::new();
let script = "42";
let result = engine.transform(script, input);
assert!(result.is_err());
}
#[test]
fn test_execute_returns_dynamic_value() {
let engine = ScriptEngine::new();
let mut scope = Scope::new();
scope.push("x", 10_i64);
let result = engine.execute("x * 2", &mut scope);
let value = result.unwrap();
assert_eq!(value.as_int().unwrap(), 20);
}
#[test]
fn test_execute_script_error_returns_execution_failed() {
let engine = ScriptEngine::new();
let mut scope = Scope::new();
let result = engine.execute("undefined_function_call()", &mut scope);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.starts_with("Script execution failed:")
);
}
#[test]
fn test_print_and_debug_statements_invoke_noop_callbacks() {
let engine = ScriptEngine::new();
let mut scope = Scope::new();
let result = engine.execute(r#"print("hello"); debug("world"); true"#, &mut scope);
assert!(result.unwrap().as_bool().unwrap());
}
#[test]
fn test_default_creates_working_engine() {
let engine = ScriptEngine::default();
let result = engine.validate("content.len() > 0", "hi");
assert!(result.unwrap());
}
#[test]
fn test_gate_rule_allows_matching_tool() {
let engine = ScriptEngine::new();
let script = r#"
context["tool"] == "send_email"
&& context["target"].ends_with("@mycompany.com")
&& (context["taint_level"] == "public" || context["taint_level"] == "internal")
"#;
let result = engine
.check_gate_rule(
script,
"send_email",
Some("alice@mycompany.com"),
"internal",
)
.unwrap();
assert!(result);
}
#[test]
fn test_gate_rule_blocks_external_email() {
let engine = ScriptEngine::new();
let script = r#"
context["tool"] == "send_email"
&& context["target"].ends_with("@mycompany.com")
"#;
let result = engine
.check_gate_rule(script, "send_email", Some("bob@external.com"), "internal")
.unwrap();
assert!(!result);
}
#[test]
fn test_gate_rule_blocks_wrong_tool() {
let engine = ScriptEngine::new();
let script = r#"context["tool"] == "send_email""#;
let result = engine
.check_gate_rule(script, "post_to_slack", None, "public")
.unwrap();
assert!(!result);
}
#[test]
fn test_gate_rule_no_target_uses_empty_string() {
let engine = ScriptEngine::new();
let script = r#"context["target"] == """#;
let result = engine
.check_gate_rule(script, "shell", None, "public")
.unwrap();
assert!(result);
}
#[test]
fn test_gate_rule_script_error() {
let engine = ScriptEngine::new();
let result = engine.check_gate_rule("invalid {{ syntax", "shell", None, "public");
assert!(result.is_err());
}
}