Add secrets detection tool pre-invoke hook

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "secrets_detection"
-version = "0.3.1"
+version = "0.3.2"
 edition.workspace = true
 authors.workspace = true
 license.workspace = true

plugins/rust/python-package/secrets_detection/Cargo.toml

@@ -1,9 +1,10 @@
 description: "Detect likely credentials and secrets in prompt input, tool output, and resource content"
 author: "ContextForge Contributors"
-version: "0.3.1"
+version: "0.3.2"
 kind: "cpex_secrets_detection.secrets_detection.SecretsDetectionPlugin"
 available_hooks:
   - "prompt_pre_fetch"
+  - "tool_pre_invoke"
   - "tool_post_invoke"
   - "resource_post_fetch"
 default_configs:

plugins/rust/python-package/secrets_detection/cpex_secrets_detection/plugin-manifest.yaml

@@ -20,6 +20,9 @@ def __init__(self, config) -> None:
     async def prompt_pre_fetch(self, payload, context):
         return self._core.prompt_pre_fetch(payload, context)
 
+    async def tool_pre_invoke(self, payload, context):
+        return self._core.tool_pre_invoke(payload, context)
+
     async def tool_post_invoke(self, payload, context):
         return self._core.tool_post_invoke(payload, context)
 

plugins/rust/python-package/secrets_detection/cpex_secrets_detection/secrets_detection.py

@@ -12,6 +12,7 @@ __all__ = [
 class SecretsDetectionPluginCore:
     def __new__(cls, config: typing.Any) -> SecretsDetectionPluginCore: ...
     def prompt_pre_fetch(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
+    def tool_pre_invoke(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
     def tool_post_invoke(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
     def resource_post_fetch(self, payload: typing.Any, _context: typing.Any) -> typing.Any: ...
 

plugins/rust/python-package/secrets_detection/cpex_secrets_detection/secrets_detection_rust/__init__.pyi

@@ -8,7 +8,7 @@ use pyo3_stub_gen::derive::*;
 
 use crate::config::SecretsDetectionConfig;
 use crate::object_model::copy_object_with_updates;
-use crate::scanner::scan_container;
+use crate::scanner::{scan_container, scan_container_findings};
 
 #[gen_stub_pyclass]
 #[pyclass]
@@ -32,54 +32,28 @@ impl SecretsDetectionPluginCore {
         payload: &Bound<'_, PyAny>,
         _context: &Bound<'_, PyAny>,
     ) -> PyResult<Py<PyAny>> {
-        let args = payload.getattr("args")?;
-        let (count, redacted_args, findings) = scan_container(py, &args, &self.config)?;
-        if self.should_block(count) {
-            let modified_payload = if self.config.redact && count > 0 {
-                copy_with_update(py, payload, [("args", redacted_args.clone().unbind())])?
-            } else {
-                payload.clone().unbind()
-            };
-            return blocked_result(
-                py,
-                "PromptPrehookResult",
-                "Potential secrets detected in prompt arguments",
-                count,
-                findings.as_any(),
-                modified_payload,
-            );
-        }
-
-        if self.config.redact && count > 0 {
-            let modified_payload =
-                copy_with_update(py, payload, [("args", redacted_args.unbind())])?;
-            return build_framework_object(
-                py,
-                "PromptPrehookResult",
-                [
-                    ("modified_payload", modified_payload),
-                    (
-                        "metadata",
-                        redaction_metadata(py, count)?.into_any().unbind(),
-                    ),
-                ],
-            );
-        }
-
-        if count > 0 {
-            return build_framework_object(
-                py,
-                "PromptPrehookResult",
-                [(
-                    "metadata",
-                    findings_metadata(py, count, findings.as_any())?
-                        .into_any()
-                        .unbind(),
-                )],
-            );
-        }
+        self.scan_payload_attr(
+            py,
+            payload,
+            "args",
+            "PromptPrehookResult",
+            "Potential secrets detected in prompt arguments",
+        )
+    }
 
-        default_result(py, "PromptPrehookResult")
+    pub fn tool_pre_invoke(
+        &self,
+        py: Python<'_>,
+        payload: &Bound<'_, PyAny>,
+        _context: &Bound<'_, PyAny>,
+    ) -> PyResult<Py<PyAny>> {
+        self.scan_payload_attr(
+            py,
+            payload,
+            "args",
+            "ToolPreInvokeResult",
+            "Potential secrets detected in tool arguments",
+        )
     }
 
     pub fn tool_post_invoke(
@@ -88,30 +62,51 @@ impl SecretsDetectionPluginCore {
         payload: &Bound<'_, PyAny>,
         _context: &Bound<'_, PyAny>,
     ) -> PyResult<Py<PyAny>> {
-        let value = payload.getattr("result")?;
-        let (count, redacted_result, findings) = scan_container(py, &value, &self.config)?;
+        self.scan_payload_attr(
+            py,
+            payload,
+            "result",
+            "ToolPostInvokeResult",
+            "Potential secrets detected in tool result",
+        )
+    }
+
+    pub fn resource_post_fetch(
+        &self,
+        py: Python<'_>,
+        payload: &Bound<'_, PyAny>,
+        _context: &Bound<'_, PyAny>,
+    ) -> PyResult<Py<PyAny>> {
+        let content = payload.getattr("content")?;
+        let Ok(text) = content.getattr("text") else {
+            return default_result(py, "ResourcePostFetchResult");
+        };
+        let (count, redacted_text, findings) = scan_container(py, &text, &self.config)?;
         if self.should_block(count) {
             let modified_payload = if self.config.redact && count > 0 {
-                copy_with_update(py, payload, [("result", redacted_result.clone().unbind())])?
+                let modified_content =
+                    copy_with_update(py, &content, [("text", redacted_text.clone().unbind())])?;
+                copy_with_update(py, payload, [("content", modified_content)])?
             } else {
                 payload.clone().unbind()
             };
             return blocked_result(
                 py,
-                "ToolPostInvokeResult",
-                "Potential secrets detected in tool result",
+                "ResourcePostFetchResult",
+                "Potential secrets detected in resource content",
                 count,
                 findings.as_any(),
                 modified_payload,
             );
         }
 
         if self.config.redact && count > 0 {
-            let modified_payload =
-                copy_with_update(py, payload, [("result", redacted_result.unbind())])?;
+            let modified_content =
+                copy_with_update(py, &content, [("text", redacted_text.unbind())])?;
+            let modified_payload = copy_with_update(py, payload, [("content", modified_content)])?;
             return build_framework_object(
                 py,
-                "ToolPostInvokeResult",
+                "ResourcePostFetchResult",
                 [
                     ("modified_payload", modified_payload),
                     (
@@ -125,7 +120,7 @@ impl SecretsDetectionPluginCore {
         if count > 0 {
             return build_framework_object(
                 py,
-                "ToolPostInvokeResult",
+                "ResourcePostFetchResult",
                 [(
                     "metadata",
                     findings_metadata(py, count, findings.as_any())?
@@ -135,45 +130,58 @@ impl SecretsDetectionPluginCore {
             );
         }
 
-        default_result(py, "ToolPostInvokeResult")
+        default_result(py, "ResourcePostFetchResult")
     }
+}
 
-    pub fn resource_post_fetch(
+impl SecretsDetectionPluginCore {
+    fn should_block(&self, count: usize) -> bool {
+        self.config.block_on_detection && count >= self.config.min_findings_to_block
+    }
+
+    fn scan_payload_attr(
         &self,
         py: Python<'_>,
         payload: &Bound<'_, PyAny>,
-        _context: &Bound<'_, PyAny>,
+        attr: &str,
+        result_class: &str,
+        block_description: &str,
     ) -> PyResult<Py<PyAny>> {
-        let content = payload.getattr("content")?;
-        let Ok(text) = content.getattr("text") else {
-            return default_result(py, "ResourcePostFetchResult");
+        let value = payload.getattr(attr)?;
+        let (mut count, mut findings) = scan_container_findings(py, &value, &self.config)?;
+        let redacted_value = if self.config.redact && count > 0 {
+            let (redacted_count, redacted, redacted_findings) =
+                scan_container(py, &value, &self.config)?;
+            count = redacted_count;
+            findings = redacted_findings;
+            Some(redacted)
+        } else {
+            None
         };
-        let (count, redacted_text, findings) = scan_container(py, &text, &self.config)?;
+
         if self.should_block(count) {
             let modified_payload = if self.config.redact && count > 0 {
-                let modified_content =
-                    copy_with_update(py, &content, [("text", redacted_text.clone().unbind())])?;
-                copy_with_update(py, payload, [("content", modified_content)])?
+                let redacted = redacted_value.expect("redacted value exists");
+                copy_with_update(py, payload, [(attr, redacted.unbind())])?
             } else {
                 payload.clone().unbind()
             };
             return blocked_result(
                 py,
-                "ResourcePostFetchResult",
-                "Potential secrets detected in resource content",
+                result_class,
+                block_description,
                 count,
                 findings.as_any(),
                 modified_payload,
             );
         }
 
         if self.config.redact && count > 0 {
-            let modified_content =
-                copy_with_update(py, &content, [("text", redacted_text.unbind())])?;
-            let modified_payload = copy_with_update(py, payload, [("content", modified_content)])?;
+            let redacted = redacted_value.expect("redacted value exists");
+            let modified_payload = copy_with_update(py, payload, [(attr, redacted.unbind())])?;
             return build_framework_object(
                 py,
-                "ResourcePostFetchResult",
+                result_class,
                 [
                     ("modified_payload", modified_payload),
                     (
@@ -187,7 +195,7 @@ impl SecretsDetectionPluginCore {
         if count > 0 {
             return build_framework_object(
                 py,
-                "ResourcePostFetchResult",
+                result_class,
                 [(
                     "metadata",
                     findings_metadata(py, count, findings.as_any())?
@@ -197,13 +205,7 @@ impl SecretsDetectionPluginCore {
             );
         }
 
-        default_result(py, "ResourcePostFetchResult")
-    }
-}
-
-impl SecretsDetectionPluginCore {
-    fn should_block(&self, count: usize) -> bool {
-        self.config.block_on_detection && count >= self.config.min_findings_to_block
+        default_result(py, result_class)
     }
 }
 

plugins/rust/python-package/secrets_detection/src/plugin.rs

@@ -11,7 +11,7 @@ mod cycle_rewrite;
 mod python_scan;
 mod text_scan;
 
-pub use python_scan::scan_container;
+pub use python_scan::{scan_container, scan_container_findings};
 pub use text_scan::detect_and_redact;
 
 #[cfg(test)]