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

@@ -82,6 +82,66 @@ impl SecretsDetectionPluginCore {
         default_result(py, "PromptPrehookResult")
     }
 
+    pub fn tool_pre_invoke(
+        &self,
+        py: Python<'_>,
+        payload: &Bound<'_, PyAny>,
+        _context: &Bound<'_, PyAny>,
+    ) -> PyResult<Py<PyAny>> {
+        let arguments = payload.getattr("arguments")?;
+        let (count, redacted_arguments, findings) = scan_container(py, &arguments, &self.config)?;
+        if self.should_block(count) {
+            let modified_payload = if self.config.redact && count > 0 {
+                copy_with_update(
+                    py,
+                    payload,
+                    [("arguments", redacted_arguments.clone().unbind())],
+                )?
+            } else {
+                payload.clone().unbind()
+            };
+            return blocked_result(
+                py,
+                "ToolPreInvokeResult",
+                "Potential secrets detected in tool arguments",
+                count,
+                findings.as_any(),
+                modified_payload,
+            );
+        }
+
+        if self.config.redact && count > 0 {
+            let modified_payload =
+                copy_with_update(py, payload, [("arguments", redacted_arguments.unbind())])?;
+            return build_framework_object(
+                py,
+                "ToolPreInvokeResult",
+                [
+                    ("modified_payload", modified_payload),
+                    (
+                        "metadata",
+                        redaction_metadata(py, count)?.into_any().unbind(),
+                    ),
+                ],
+            );
+        }
+
+        if count > 0 {
+            return build_framework_object(
+                py,
+                "ToolPreInvokeResult",
+                [(
+                    "metadata",
+                    findings_metadata(py, count, findings.as_any())?
+                        .into_any()
+                        .unbind(),
+                )],
+            );
+        }
+
+        default_result(py, "ToolPreInvokeResult")
+    }
+
     pub fn tool_post_invoke(
         &self,
         py: Python<'_>,

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

@@ -13,6 +13,7 @@
     ResourceHookType,
     ResourcePostFetchPayload,
     ToolHookType,
+    ToolPreInvokePayload,
     ToolPostInvokePayload,
 )
 
@@ -32,6 +33,7 @@
     "RootModel",
     "SecretsDetectionPlugin",
     "ToolHookType",
+    "ToolPreInvokePayload",
     "ToolPostInvokePayload",
     "make_config",
     "make_context",

plugins/tests/secrets_detection/helpers.py

@@ -28,6 +28,7 @@ async def manager(
         config_path = tmp_path / "secrets_detection.yaml"
         configured_hooks = hooks or [
             PromptHookType.PROMPT_PRE_FETCH.value,
+            ToolHookType.TOOL_PRE_INVOKE.value,
             ToolHookType.TOOL_POST_INVOKE.value,
             ResourceHookType.RESOURCE_POST_FETCH.value,
         ]
@@ -189,6 +190,28 @@ async def test_tool_post_invoke_blocks_without_redaction_via_plugin_manager(
         finally:
             await manager.shutdown()
 
+    async def test_tool_pre_invoke_blocks_without_redaction_via_plugin_manager(
+        self, tmp_path
+    ):
+        manager = await self.manager(
+            tmp_path, {"block_on_detection": True, "redact": False}
+        )
+        try:
+            payload = ToolPreInvokePayload(
+                name="echo",
+                arguments={"message": "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"},
+            )
+            result, _ = await manager.invoke_hook(
+                ToolHookType.TOOL_PRE_INVOKE,
+                payload,
+                global_context=self.global_context(),
+            )
+            assert result.continue_processing is False
+            assert result.violation.code == "SECRETS_DETECTED"
+            assert result.modified_payload == payload
+        finally:
+            await manager.shutdown()
+
     async def test_resource_post_fetch_blocks_without_redaction_via_plugin_manager(
         self, tmp_path
     ):

plugins/tests/secrets_detection/test_hook_dispatch.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+import yaml
+
+
+def test_manifest_declares_tool_pre_invoke_hook():
+    manifest_path = (
+        Path(__file__).resolve().parents[3]
+        / "plugins"
+        / "rust"
+        / "python-package"
+        / "secrets_detection"
+        / "cpex_secrets_detection"
+        / "plugin-manifest.yaml"
+    )
+
+    manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8"))
+
+    assert "tool_pre_invoke" in manifest["available_hooks"]

plugins/tests/secrets_detection/test_manifest.py

@@ -50,6 +50,45 @@ async def test_prompt_pre_fetch_blocks_without_redaction(self):
         assert result.violation.code == "SECRETS_DETECTED"
         assert result.modified_payload == payload
 
+    async def test_tool_pre_invoke_redacts_arguments_without_blocking(self, plugin):
+        payload = ToolPreInvokePayload(
+            name="echo",
+            arguments={"message": "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"},
+        )
+
+        result = await plugin.tool_pre_invoke(payload, make_context())
+
+        assert result.continue_processing is True
+        assert result.violation is None
+        assert result.modified_payload is not None
+        assert result.modified_payload is not payload
+        assert (
+            result.modified_payload.arguments["message"]
+            == "AWS_ACCESS_KEY_ID=[REDACTED]"
+        )
+        assert (
+            payload.arguments["message"]
+            == "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"
+        )
+        assert result.metadata == {"secrets_redacted": True, "count": 1}
+
+    async def test_tool_pre_invoke_blocks_without_redaction(self):
+        plugin = SecretsDetectionPlugin(make_config(block_on_detection=True, redact=False))
+        payload = ToolPreInvokePayload(
+            name="echo",
+            arguments={"message": "AWS_ACCESS_KEY_ID=AKIAFAKE12345EXAMPLE"},
+        )
+
+        result = await plugin.tool_pre_invoke(payload, make_context())
+
+        assert result.continue_processing is False
+        assert result.violation is not None
+        assert result.violation.code == "SECRETS_DETECTED"
+        assert result.violation.description == (
+            "Potential secrets detected in tool arguments"
+        )
+        assert result.modified_payload == payload
+
     async def test_prompt_pre_fetch_blocks_with_redaction_without_leaking_secret(self):
         plugin = SecretsDetectionPlugin(make_config(block_on_detection=True, redact=True))
         payload = PromptPrehookPayload(