security(audit): redact nested arguments and result previews (#397, #525)

coworker/audit.py

@@ -10,6 +10,10 @@
 
 from .connectors import connector_for_tool
 
+# Matched as substrings of a lowercased key name, at EVERY level of a structure (#397):
+# an HTTP-shaped or MCP tool takes its credential in a nested `headers` / `auth` / `config`
+# object, and `_truncate` keeps the first 500 characters, so an unredacted bearer token
+# lands in the log whole.
 _SECRET_KEYS = (
     "token",
     "secret",
@@ -18,9 +22,23 @@
     "access_token",
     "bot_token",
     "app_token",
+    "authorization",
+    "cookie",
+    "credential",
+    "private_key",
     "raw",
 )
 _BODY_KEYS = ("body", "content", "html")
+# A tool RESULT carries the same content the argument policy redacts, under other names:
+# a shell command's stdout, an email body, one message's text (#525). The audit row is for
+# triage — who ran what, against which resource — never for replaying the content.
+_RESULT_BODY_KEYS = _BODY_KEYS + ("output", "stdout", "stderr", "text", "snippet")
+# The engine's own preview length, so a rebuilt preview is the same size as the one it
+# replaces.
+_PREVIEW_LIMIT = 300
+# How deep the walk goes before it stops describing a structure. Past this the keys are no
+# longer being checked, so the value is dropped rather than copied through.
+_MAX_DEPTH = 6
 
 
 class AuditStore:
@@ -83,6 +101,7 @@ def append(self, event: dict[str, Any]) -> None:
         resource = _resource(
             tool, event.get("arguments") or {}, event.get("result") or {}
         )
+        preview = _result_preview(tool, event)
         with self._lock:
             self._conn.execute(
                 """
@@ -100,7 +119,7 @@ def append(self, event: dict[str, Any]) -> None:
                     event.get("status") or "",
                     event.get("approval") or "",
                     json.dumps(args, default=str),
-                    _truncate(str(event.get("result_preview") or "")),
+                    preview,
                     _truncate(str(event.get("reason") or "")),
                     _truncate(str(resource or "")),
                     str(event.get("call_id") or ""),
@@ -190,34 +209,79 @@ def close(self) -> None:
 
 
 def _sanitize_args(tool: str, args: dict[str, Any]) -> dict[str, Any]:
+    """Tool arguments, with every secret-like and body-like value replaced by a marker."""
     if not isinstance(args, dict):
         return {}
-    out: dict[str, Any] = {}
-    for key, value in args.items():
-        lk = str(key).lower()
-        if any(s in lk for s in _SECRET_KEYS):
-            out[key] = "[redacted]"
-        elif tool == "browser_type" and lk == "text":
-            out[key] = "[redacted input]"
-        elif any(b == lk or lk.endswith("_" + b) for b in _BODY_KEYS):
-            out[key] = "[redacted body]"
-        else:
-            out[key] = _summarize(value)
-    return out
-
-
-def _summarize(value: Any) -> Any:
+    return _sanitize_mapping(tool, args, _BODY_KEYS, 0)
+
+
+def _sanitize_result(tool: str, result: Any) -> Any:
+    """A tool result under the same policy, plus the result-side content keys."""
+    return _sanitize_value(tool, None, result, _RESULT_BODY_KEYS, 0)
+
+
+def _redaction_marker(
+    tool: str, lower_key: str, body_keys: tuple[str, ...]
+) -> Optional[str]:
+    """The marker this key's value must be replaced by, or None to keep the value."""
+    if any(s in lower_key for s in _SECRET_KEYS):
+        return "[redacted]"
+    if tool == "browser_type" and lower_key == "text":
+        return "[redacted input]"
+    if any(b == lower_key or lower_key.endswith("_" + b) for b in body_keys):
+        return "[redacted body]"
+    return None
+
+
+def _sanitize_mapping(
+    tool: str, mapping: dict[Any, Any], body_keys: tuple[str, ...], depth: int
+) -> dict[str, Any]:
+    return {
+        str(key): _sanitize_value(tool, key, value, body_keys, depth)
+        for key, value in list(mapping.items())[:20]
+    }
+
+
+def _sanitize_value(
+    tool: str, key: Any, value: Any, body_keys: tuple[str, ...], depth: int
+) -> Any:
+    if key is not None:
+        marker = _redaction_marker(tool, str(key).lower(), body_keys)
+        if marker is not None:
+            return marker
     if isinstance(value, str):
         return _truncate(value)
     if isinstance(value, (int, float, bool)) or value is None:
         return value
+    if depth >= _MAX_DEPTH:
+        # Stringifying it here would copy through the very keys we stopped checking.
+        return "[nested]"
     if isinstance(value, list):
-        return [_summarize(v) for v in value[:10]]
+        # An item carries no key of its own; a dict item is checked when we recurse
+        # into it, and a list under a secret-like key never reaches here at all.
+        return [
+            _sanitize_value(tool, None, v, body_keys, depth + 1) for v in value[:10]
+        ]
     if isinstance(value, dict):
-        return {str(k): _summarize(v) for k, v in list(value.items())[:20]}
+        return _sanitize_mapping(tool, value, body_keys, depth + 1)
     return _truncate(str(value))
 
 
+def _result_preview(tool: str, event: dict[str, Any]) -> str:
+    """The stored preview of a tool result, redacted at the structured stage.
+
+    The caller's `result_preview` is already flattened to a string, so nothing can be
+    redacted in it by key any more — when the raw `result` rides along, the preview is
+    rebuilt from the sanitized structure instead (#525).
+    """
+    result = event.get("result")
+    if result is None:
+        return _truncate(str(event.get("result_preview") or ""))
+    sanitized = _sanitize_result(tool, result)
+    text = sanitized if isinstance(sanitized, str) else json.dumps(sanitized, default=str)
+    return _truncate(text, limit=_PREVIEW_LIMIT)
+
+
 def _resource(tool: str, args: dict[str, Any], result: Any) -> str:
     for key in (
         "url",

tests/test_audit_redaction.py

@@ -0,0 +1,122 @@
+"""Audit-log redaction — nested arguments (#397) and tool-result previews (#525).
+
+The audit row is the record of who ran what against which resource. It is not a place
+for the credential the call carried, nor for the content the call returned.
+"""
+
+from __future__ import annotations
+
+import json
+
+from coworker.audit import AuditStore
+
+
+def _row(tmp_path, event: dict) -> dict:
+    store = AuditStore(tmp_path / "audit.db")
+    try:
+        store.append({"session_id": "s1", "tool": "http_request", **event})
+        return store.list(limit=1)[0]
+    finally:
+        store.close()
+
+
+def test_nested_credentials_are_redacted(tmp_path):
+    row = _row(
+        tmp_path,
+        {
+            "arguments": {
+                "url": "https://example.com",
+                "headers": {"Authorization": "Bearer sk-live-NESTED"},
+                "config": {"api_key": "sk-live-NESTED2"},
+            }
+        },
+    )
+    assert row["args"]["headers"]["Authorization"] == "[redacted]"
+    assert row["args"]["config"]["api_key"] == "[redacted]"
+    assert row["args"]["url"] == "https://example.com"  # the resource still reads
+    assert "sk-live" not in json.dumps(row["args"])
+
+
+def test_credential_keys_beyond_the_token_family(tmp_path):
+    row = _row(
+        tmp_path,
+        {
+            "arguments": {
+                "authorization": "Bearer sk-live-AUTHZ",
+                "cookie": "session=sk-live-COOKIE",
+                "credential": "sk-live-CRED",
+                "private_key": "-----BEGIN...",
+            }
+        },
+    )
+    assert set(row["args"].values()) == {"[redacted]"}
+
+
+def test_nested_bodies_are_redacted(tmp_path):
+    row = _row(
+        tmp_path,
+        {"arguments": {"draft": {"to": "a@example.com", "body": "private text"}}},
+    )
+    assert row["args"]["draft"]["body"] == "[redacted body]"
+    assert row["args"]["draft"]["to"] == "a@example.com"
+
+
+def test_a_credential_below_the_walk_limit_is_dropped_not_copied(tmp_path):
+    deep: dict = {"api_key": "sk-live-DEEP"}
+    for _ in range(10):
+        deep = {"wrap": deep}
+    row = _row(tmp_path, {"arguments": deep})
+    assert "sk-live" not in json.dumps(row["args"])
+
+
+def test_browser_typing_is_still_redacted_input(tmp_path):
+    row = _row(
+        tmp_path, {"tool": "browser_type", "arguments": {"text": "hunter2"}}
+    )
+    assert row["args"]["text"] == "[redacted input]"
+
+
+def test_result_preview_drops_an_email_body_and_keeps_the_envelope(tmp_path):
+    row = _row(
+        tmp_path,
+        {
+            "tool": "email_read",
+            "arguments": {"uid": "42"},
+            "stage": "finished",
+            "result": {
+                "ok": True,
+                "subject": "Q3 numbers",
+                "body": "the confidential text of the message",
+            },
+            "result_preview": "unsanitized preview from the caller",
+        },
+    )
+    assert "confidential" not in row["result_preview"]
+    assert "[redacted body]" in row["result_preview"]
+    assert "Q3 numbers" in row["result_preview"]  # triage still works
+
+
+def test_result_preview_drops_shell_output_and_keeps_the_command(tmp_path):
+    row = _row(
+        tmp_path,
+        {
+            "tool": "run_shell",
+            "arguments": {"command": "printenv"},
+            "stage": "finished",
+            "result": {
+                "command": "printenv",
+                "exit_code": 0,
+                "output": "AWS_SECRET_ACCESS_KEY=sk-live-ENV",
+            },
+        },
+    )
+    assert "sk-live-ENV" not in row["result_preview"]
+    assert "printenv" in row["result_preview"]
+
+
+def test_preview_without_a_raw_result_is_still_stored(tmp_path):
+    row = _row(
+        tmp_path,
+        {"stage": "finished", "result_preview": "{\"ok\": true}"},
+    )
+    assert row["result_preview"] == '{"ok": true}'