Merge pull request #28 from NVIDIA-NeMo/fix/secret-log-redaction

Author: alessiodevoto

src/nooa/mcp/oauth.py

@@ -550,19 +550,14 @@ async def exchange_code_for_token(self, code: str) -> OAuthToken:
                     client_secret=self.config.client_secret,
                 )
             except httpx.HTTPStatusError as e:
-                error_detail = e.response.text if e.response else str(e)
                 error_json = None
                 try:
                     if e.response:
                         error_json = e.response.json()
                 except Exception:
                     pass
 
-                logger.error(
-                    f"Token exchange failed: {e.response.status_code}, Request data: {data}, Response: {error_detail}",
-                )
-                if error_json:
-                    logger.error(f"Error details: {error_json}")
+                logger.error("Token exchange failed (HTTP %s)", e.response.status_code)
 
                 # Provide helpful error message for common issues
                 error_msg = f"Failed to exchange authorization code for token (HTTP {e.response.status_code})"
@@ -572,9 +567,6 @@ async def exchange_code_for_token(self, code: str) -> OAuthToken:
                         "Authorization codes are typically valid for only a few minutes. "
                         "Please try the OAuth flow again to get a fresh code."
                     )
-                else:
-                    error_msg += f": {error_detail}"
-
                 raise RuntimeError(error_msg) from e
 
     async def client_credentials_token(self) -> OAuthToken:
@@ -643,7 +635,7 @@ async def complete_flow(self, open_browser: bool = True) -> OAuthToken:
             logger.info("OAuth flow completed successfully")
             return token
         except Exception as e:
-            logger.error(f"OAuth flow failed: {e}")
+            logger.error("OAuth flow failed (%s)", type(e).__name__)
             raise
 
 

src/nooa/tracing/_secret_scrubber.py

@@ -24,6 +24,37 @@
 
 REDACTED = "[REDACTED]"
 
+_SENSITIVE_KEYS = frozenset(
+    {
+        "authorization",
+        "proxy_authorization",
+        "api_key",
+        "x_api_key",
+        "access_token",
+        "refresh_token",
+        "auth_token",
+        "session_token",
+        "id_token",
+        "client_secret",
+        "client_assertion",
+        "password",
+        "passwd",
+        "private_key",
+        "secret_key",
+        "code_verifier",
+        "cookie",
+        "set_cookie",
+    }
+)
+
+
+def _is_sensitive_key(key: Any) -> bool:
+    normalized = re.sub(r"[^a-z0-9]+", "_", str(key).lower()).strip("_")
+    return normalized in _SENSITIVE_KEYS or any(
+        normalized.endswith(f"_{sensitive}") for sensitive in _SENSITIVE_KEYS
+    )
+
+
 # ---------------------------------------------------------------------------
 # Regex-based secret patterns — high precision, low false positives
 # ---------------------------------------------------------------------------
@@ -62,8 +93,9 @@
         "generic_api_key",
         re.compile(
             r"(?:api[_-]?key|api[_-]?token|auth[_-]?token|access[_-]?token|"
+            r"client[_-]?secret|refresh[_-]?token|code[_-]?verifier|"
             r"secret[_-]?key|private[_-]?key|password|passwd)"
-            r"[\"\']?\s*[=:]+\s*[\"\'\s]*(?P<secret>[A-Za-z0-9_.\-/+=]{20,})",
+            r"[\"\']?\s*[=:]+\s*[\"\'\s]*(?P<secret>[A-Za-z0-9_.\-/+=]{1,})",
             re.IGNORECASE,
         ),
     ),
@@ -153,7 +185,7 @@ def scrub_string(text: str) -> tuple[str, int]:
     span processing). Returns the original string unchanged with a count of 0
     if no secrets are found (fast path).
     """
-    if not text or len(text) < 10:
+    if not text:
         return text, 0
 
     count = 0
@@ -176,12 +208,25 @@ def _replace(m: re.Match, _name: str = name) -> str:
 def scrub_value(value: Any) -> tuple[Any, int]:
     """Scrub secrets from a span attribute value.
 
-    Handles strings and sequences of strings. Other types pass through.
+    Handles strings, mappings, and sequences. Values under credential-bearing
+    keys are always replaced, including short or provider-specific secrets.
     Returns ``(scrubbed_value, redaction_count)`` where the count is the
     number of secrets redacted within this value.
     """
     if isinstance(value, str):
         return scrub_string(value)
+    if isinstance(value, dict):
+        scrubbed: dict[Any, Any] = {}
+        count = 0
+        for key, item in value.items():
+            if _is_sensitive_key(key):
+                scrubbed[key] = REDACTED
+                stats.record("sensitive_key")
+                count += 1
+            else:
+                scrubbed[key], n = scrub_value(item)
+                count += n
+        return scrubbed, count
     if isinstance(value, (list, tuple)):
         new_items = []
         count = 0
@@ -223,7 +268,11 @@ def on_end(self, span: ReadableSpan) -> None:
                 redacted_count = 0
 
                 for key, value in span.attributes.items():
-                    new_value, n = scrub_value(value)
+                    if _is_sensitive_key(key):
+                        new_value, n = REDACTED, 1
+                        stats.record("sensitive_key")
+                    else:
+                        new_value, n = scrub_value(value)
                     scrubbed[key] = new_value
                     redacted_count += n
 

src/nooa/unifiedllm/http_logging.py

@@ -5,6 +5,8 @@
 from datetime import UTC, datetime
 from pathlib import Path
 
+from nooa.tracing._secret_scrubber import REDACTED, _is_sensitive_key, scrub_value
+
 
 def enable_http_request_logging(
     output_dir: str | Path = ".",
@@ -89,12 +91,24 @@ def enable_http_request_logging(
 
     def _redact_headers(headers: dict) -> dict:
         """Redact sensitive headers like Authorization."""
-        redacted = dict(headers)
-        for key in list(redacted.keys()):
-            if key.lower() in ["authorization", "api-key", "x-api-key"]:
+        redacted = {}
+        for key, value in headers.items():
+            if _is_sensitive_key(key):
                 redacted[key] = "***REDACTED***"
+            else:
+                redacted[key], _ = scrub_value(value)
         return redacted
 
+    def _redact_body(body):
+        """Recursively scrub secrets before a parsed HTTP body is logged."""
+        scrubbed, _ = scrub_value(body)
+        if isinstance(scrubbed, dict) and scrubbed.get("grant_type") == "authorization_code":
+            # "code" is too generic for global key matching, but is a credential
+            # in an OAuth authorization-code exchange.
+            if "code" in scrubbed:
+                scrubbed["code"] = REDACTED
+        return scrubbed
+
     def _write_jsonl_entry(entry: dict):
         """Append a JSON entry to the JSONL error file."""
         if jsonl_file is None:
@@ -113,7 +127,7 @@ def _log_request(request, counter):
                 filename = output_path / f"request_{counter}_{model_name}.json"
 
                 with open(filename, "w") as f:
-                    json.dump(body_dict, f, indent=2)
+                    json.dump(_redact_body(body_dict), f, indent=2)
 
                 if verbose:
                     print(f"\n💾 Saved HTTP request to: {filename}")
@@ -144,7 +158,7 @@ def _log_response_sync(response, request, counter, model_name):
         try:
             response_dict = {
                 "status_code": response.status_code,
-                "headers": dict(response.headers),
+                "headers": _redact_headers(dict(response.headers)),
             }
 
             try:
@@ -169,6 +183,8 @@ def _log_response_sync(response, request, counter, model_name):
             except Exception as read_error:
                 response_dict["body"] = f"[Error reading response: {read_error}]"
 
+            response_dict["body"] = _redact_body(response_dict["body"])
+
             _save_response_file(response_dict, counter, model_name)
         except Exception as e:
             if verbose:
@@ -180,7 +196,7 @@ async def _log_response_async(response, request, counter, model_name):
         try:
             response_dict = {
                 "status_code": response.status_code,
-                "headers": dict(response.headers),
+                "headers": _redact_headers(dict(response.headers)),
             }
 
             try:
@@ -206,6 +222,8 @@ async def _log_response_async(response, request, counter, model_name):
             except Exception as read_error:
                 response_dict["body"] = f"[Error reading response: {read_error}]"
 
+            response_dict["body"] = _redact_body(response_dict["body"])
+
             _save_response_file(response_dict, counter, model_name)
         except Exception as e:
             if verbose:
@@ -236,7 +254,7 @@ def logging_send(self, request):
                         "url": str(request.url),
                         "method": request.method,
                         "headers": _redact_headers(dict(request.headers)),
-                        "body": body_dict,
+                        "body": _redact_body(body_dict),
                     }
             except Exception as e:
                 if verbose:
@@ -265,8 +283,8 @@ def logging_send(self, request):
 
                     response_data = {
                         "status_code": response.status_code,
-                        "headers": dict(response.headers),
-                        "body": body_content,
+                        "headers": _redact_headers(dict(response.headers)),
+                        "body": _redact_body(body_content),
                     }
 
                     # Write JSONL entry
@@ -324,7 +342,7 @@ async def async_logging_send(self, request):
                         "url": str(request.url),
                         "method": request.method,
                         "headers": _redact_headers(dict(request.headers)),
-                        "body": body_dict,
+                        "body": _redact_body(body_dict),
                     }
             except Exception as e:
                 if verbose:
@@ -353,8 +371,8 @@ async def async_logging_send(self, request):
 
                     response_data = {
                         "status_code": response.status_code,
-                        "headers": dict(response.headers),
-                        "body": body_content,
+                        "headers": _redact_headers(dict(response.headers)),
+                        "body": _redact_body(body_content),
                     }
 
                     # Write JSONL entry

tests/tracing/test_secret_scrubber.py

@@ -139,9 +139,14 @@ def test_empty_string(self):
         assert scrub_string("") == ("", 0)
 
     def test_short_string(self):
-        """Strings below the minimum length are skipped (fast path)."""
+        """A clean short string passes through unchanged."""
         assert scrub_string("hello") == ("hello", 0)
 
+    def test_short_generic_secret(self):
+        result, count = scrub_string("client_secret=s3cr3t")
+        assert "s3cr3t" not in result
+        assert count == 1
+
     def test_preserves_surrounding_text(self):
         """Surrounding non-secret text is preserved around a redaction."""
         text = "before AKIAIOSFODNN7EXAMPLE after"
@@ -203,6 +208,13 @@ def test_none_passthrough(self):
         """None passes through unchanged with a zero count."""
         assert scrub_value(None) == (None, 0)
 
+    def test_nested_sensitive_keys(self):
+        result, count = scrub_value(
+            {"safe": {"client_secret": "short", "refresh_token": "provider-specific"}}
+        )
+        assert result == {"safe": {"client_secret": REDACTED, "refresh_token": REDACTED}}
+        assert count == 2
+
 
 class TestScrubStats:
     def test_record_and_snapshot(self):