update pyfunc action interface

mcp-guardian-core/src/guard_profile/profiles/block-get-current-time.json

@@ -16,6 +16,7 @@
                 "type": "PyFunc",
                 "code_lines": [
                     "import sys",
+                    "import mcp_guardian",
                     "",
                     "# msg_direction is 'inbound' or 'outbound'",
                     "msg_direction = sys.argv[0]",
@@ -25,10 +26,9 @@
                     "raw_msg = sys.argv[2]",
                     "",
                     "if 'get_current_time' in raw_msg:",
-                    "    action = 'drop'",
+                    "    mcp_guardian.drop()",
                     "else:",
-                    "    action = 'send'",
-                    "    outbound_msg = raw_msg"
+                    "    mcp_guardian.send(raw_msg)"
                 ]
             }
         },

mcp-guardian-core/src/message_interceptor/py_func.rs → mcp-guardian-core/src/message_interceptor/py_func/mod.rs

@@ -1,3 +1,5 @@
+pub mod sys_mcp_guardian;
+
 use anyhow::{anyhow, bail, Result};
 use async_trait::async_trait;
 use rustpython_vm::{builtins::PyStr, Interpreter, Settings};
@@ -39,31 +41,33 @@ impl MessageInterceptor for PyFuncInterceptor {
         settings.argv = args;
 
         let action = Interpreter::without_stdlib(settings).enter(|vm| {
+            sys_mcp_guardian::add_to_vm(vm)?;
+
             let scope = vm.new_scope_with_builtins();
 
             vm.run_code_string(scope.clone(), &code, "<embedded>".to_owned())
                 .map_err(|_| anyhow!("Error running Python code."))?;
 
             let action = scope
                 .globals
-                .get_item("action", vm)
+                .get_item("__action", vm)
                 .map_err(|_| anyhow!("Failed to get 'action' from Python scope."))?;
             let Some(action) = action.downcast_ref::<PyStr>() else {
                 bail!("Failed to downcast 'action' to PyStr.");
             };
             let action = match action.as_str() {
                 "send" => {
-                    let outbound_msg = scope
+                    let msg = scope
                         .globals
-                        .get_item("outbound_msg", vm)
-                        .map_err(|_| anyhow!("Failed to get 'outbound_msg' from Python scope."))?;
-                    let Some(outbound_msg) = outbound_msg.downcast_ref::<PyStr>() else {
-                        bail!("Failed to downcast 'outbound_msg' to PyStr.");
+                        .get_item("__msg", vm)
+                        .map_err(|_| anyhow!("Failed to get '__msg' from Python scope."))?;
+                    let Some(msg) = msg.downcast_ref::<PyStr>() else {
+                        bail!("Failed to downcast '__msg' to PyStr.");
                     };
 
                     let message = Message {
                         type_: message.type_,
-                        raw_msg: serde_json::from_str(outbound_msg.as_str())?,
+                        raw_msg: serde_json::from_str(msg.as_str())?,
                     };
 
                     Send(message)

mcp-guardian-core/src/message_interceptor/py_func/sys_mcp_guardian.rs

@@ -0,0 +1,43 @@
+use anyhow::{anyhow, Result};
+use rustpython_vm::{builtins::PyStr, PyObjectRef, PyPayload, PyResult, VirtualMachine};
+
+pub fn add_to_vm(vm: &VirtualMachine) -> Result<()> {
+    let mcp_guardian = vm.new_module("mcp_guardian", vm.ctx.new_dict(), None);
+
+    mcp_guardian
+        .set_attr("send", vm.new_function("send", send), vm)
+        .map_err(|_| anyhow!("Failed to set 'send' in mcp_guardian module."))?;
+
+    mcp_guardian
+        .set_attr("drop", vm.new_function("drop", drop), vm)
+        .map_err(|_| anyhow!("Failed to set 'drop' in mcp_guardian module."))?;
+
+    mcp_guardian
+        .set_attr("return", vm.new_function("return", return_), vm)
+        .map_err(|_| anyhow!("Failed to set 'return' in mcp_guardian module."))?;
+
+    Ok(())
+}
+
+fn send(_self: PyObjectRef, msg: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+    vm.current_globals()
+        .set_item("__action", PyStr::from("send").into_pyobject(vm), vm)?;
+    vm.current_globals().set_item("__msg", msg, vm)?;
+
+    Ok(())
+}
+
+fn drop(_self: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+    vm.current_globals()
+        .set_item("__action", PyStr::from("drop").into_pyobject(vm), vm)?;
+
+    Ok(())
+}
+
+fn return_(_self: PyObjectRef, msg: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
+    vm.current_globals()
+        .set_item("__action", PyStr::from("return").into_pyobject(vm), vm)?;
+    vm.current_globals().set_item("__msg", msg, vm)?;
+
+    Ok(())
+}