Skip to content

Commit 13e4882

Browse files
committed
fix(weather): stop logging API keys and simplify validation
Remove the middleware that logged Authorization headers on every request, exposing API keys in pod logs. Defense-in-depth changes: - Change root log level from DEBUG to INFO - Add SecretRedactionFilter (Bearer tokens + literal configured key) - Add has_valid_api_key check with clear user-facing error message - Simplified from prior iteration per review feedback (esnible, mrsabath): dropped sk-* regex (literal key covers it), inlined is_local_llm, removed log_warnings(), handled dict args in filter 18 new tests cover redaction and API key validation edge cases. Closes #119 Assisted-By: Claude (Anthropic AI) <noreply@anthropic.com> Signed-off-by: Paolo Dettori <dettori@us.ibm.com>
1 parent 1091806 commit 13e4882

3 files changed

Lines changed: 241 additions & 12 deletions

File tree

a2a/weather_service/src/weather_service/agent.py

Lines changed: 41 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import logging
22
import os
3+
import re
34
from textwrap import dedent
45

56
import uvicorn
@@ -14,13 +15,43 @@
1415
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
1516
from a2a.types import AgentCapabilities, AgentCard, AgentSkill, TaskState, TextPart
1617
from a2a.utils import new_agent_text_message, new_task
18+
from weather_service.configuration import Configuration
1719
from weather_service.graph import get_graph, get_mcpclient
1820
from weather_service.observability import (
1921
create_tracing_middleware,
2022
get_root_span,
2123
set_span_output,
2224
)
2325

26+
27+
class SecretRedactionFilter(logging.Filter):
28+
"""Redacts Bearer tokens and the configured API key from log messages."""
29+
30+
_BEARER_RE = re.compile(r"(Bearer\s+)\S+", re.IGNORECASE)
31+
32+
def __init__(self):
33+
super().__init__()
34+
key = os.environ.get("LLM_API_KEY", "").strip()
35+
self._key_re = re.compile(re.escape(key)) if len(key) > 8 else None
36+
37+
def _redact(self, text: str) -> str:
38+
text = self._BEARER_RE.sub(r"\1[REDACTED]", text)
39+
if self._key_re:
40+
text = self._key_re.sub("[REDACTED]", text)
41+
return text
42+
43+
def filter(self, record: logging.LogRecord) -> bool:
44+
if isinstance(record.msg, str):
45+
record.msg = self._redact(record.msg)
46+
if isinstance(record.args, dict):
47+
record.args = {k: self._redact(v) if isinstance(v, str) else v for k, v in record.args.items()}
48+
elif isinstance(record.args, tuple):
49+
record.args = tuple(self._redact(a) if isinstance(a, str) else a for a in record.args)
50+
return True
51+
52+
53+
logging.basicConfig(level=logging.INFO)
54+
logging.getLogger().addFilter(SecretRedactionFilter())
2455
logger = logging.getLogger(__name__)
2556

2657

@@ -110,6 +141,16 @@ async def execute(self, context: RequestContext, event_queue: EventQueue):
110141
task_updater = TaskUpdater(event_queue, task.id, task.context_id)
111142
event_emitter = A2AEvent(task_updater)
112143

144+
# Check API key before attempting LLM calls
145+
config = Configuration()
146+
if not config.has_valid_api_key:
147+
await event_emitter.emit_event(
148+
"Error: No LLM API key configured. Set the LLM_API_KEY "
149+
"environment variable.",
150+
failed=True,
151+
)
152+
return
153+
113154
# Get user input for the agent
114155
user_input = context.get_user_input()
115156

@@ -225,16 +266,4 @@ def run():
225266
# Add tracing middleware - creates root span with MLflow/GenAI attributes
226267
app.add_middleware(BaseHTTPMiddleware, dispatch=create_tracing_middleware())
227268

228-
class LogAuthorizationMiddleware(BaseHTTPMiddleware):
229-
async def dispatch(self, request, call_next):
230-
auth_header = request.headers.get("authorization", "No Authorization header")
231-
logger.info(
232-
f"🔐 Incoming request to {request.url.path} with Authorization: {auth_header[:80] + '...' if len(auth_header) > 80 else auth_header}"
233-
)
234-
response = await call_next(request)
235-
return response
236-
237-
# Add logging middleware
238-
app.add_middleware(LogAuthorizationMiddleware)
239-
240269
uvicorn.run(app, host="0.0.0.0", port=8000)
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,23 @@
1+
from urllib.parse import urlparse
2+
13
from pydantic_settings import BaseSettings
24

5+
_PLACEHOLDER_KEYS = {"dummy", "changeme", "your-api-key-here", ""}
6+
37

48
class Configuration(BaseSettings):
59
llm_model: str = "llama3.1"
610
llm_api_base: str = "http://localhost:11434/v1"
711
llm_api_key: str = "dummy"
12+
13+
@property
14+
def has_valid_api_key(self) -> bool:
15+
"""Placeholder keys are only invalid for remote APIs.
16+
17+
Local LLMs (Ollama, vLLM) accept any key, so we skip validation
18+
when the API base points to localhost.
19+
"""
20+
host = urlparse(self.llm_api_base).hostname or ""
21+
if host in {"localhost", "127.0.0.1", "0.0.0.0"}:
22+
return True
23+
return self.llm_api_key.strip() not in _PLACEHOLDER_KEYS
Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
"""Tests for secret redaction and API key validation in the weather service.
2+
3+
Loads agent.py and configuration.py in isolation (same approach as
4+
test_weather_service.py) to avoid pulling in heavy deps like opentelemetry.
5+
"""
6+
7+
import importlib.util
8+
import logging
9+
import pathlib
10+
import sys
11+
from types import ModuleType
12+
from unittest.mock import MagicMock
13+
14+
# --- Isolation setup (must happen before any weather_service imports) ---
15+
_fake_ws = ModuleType("weather_service")
16+
_fake_ws.__path__ = [] # type: ignore[attr-defined]
17+
sys.modules.setdefault("weather_service", _fake_ws)
18+
sys.modules.setdefault("weather_service.observability", MagicMock())
19+
20+
_BASE = pathlib.Path(__file__).parent.parent.parent / "a2a" / "weather_service" / "src" / "weather_service"
21+
22+
23+
def _load_module(name: str, path: pathlib.Path):
24+
spec = importlib.util.spec_from_file_location(name, path)
25+
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
26+
sys.modules[name] = mod
27+
spec.loader.exec_module(mod) # type: ignore[union-attr]
28+
return mod
29+
30+
31+
_config_mod = _load_module("weather_service.configuration", _BASE / "configuration.py")
32+
33+
# Mock modules that agent.py imports but we don't need
34+
for mod_name in [
35+
"uvicorn",
36+
"langchain_core",
37+
"langchain_core.messages",
38+
"starlette",
39+
"starlette.middleware",
40+
"starlette.middleware.base",
41+
"starlette.routing",
42+
"a2a",
43+
"a2a.server",
44+
"a2a.server.agent_execution",
45+
"a2a.server.apps",
46+
"a2a.server.events",
47+
"a2a.server.events.event_queue",
48+
"a2a.server.request_handlers",
49+
"a2a.server.tasks",
50+
"a2a.types",
51+
"a2a.utils",
52+
"weather_service.graph",
53+
]:
54+
sys.modules.setdefault(mod_name, MagicMock())
55+
56+
_agent_mod = _load_module("weather_service.agent", _BASE / "agent.py")
57+
58+
Configuration = _config_mod.Configuration
59+
SecretRedactionFilter = _agent_mod.SecretRedactionFilter
60+
61+
62+
# --- Tests ---
63+
64+
65+
class TestSecretRedactionFilter:
66+
"""Test the logging filter that redacts Bearer tokens and API keys."""
67+
68+
def setup_method(self):
69+
self.filt = SecretRedactionFilter()
70+
71+
def _make_record(self, msg: str, args=None) -> logging.LogRecord:
72+
return logging.LogRecord(
73+
name="test", level=logging.INFO, pathname="", lineno=0,
74+
msg=msg, args=args, exc_info=None,
75+
)
76+
77+
def test_redacts_bearer_token(self):
78+
record = self._make_record("Authorization: Bearer sk-abc123xyz789secret")
79+
self.filt.filter(record)
80+
assert "sk-abc123xyz789secret" not in record.msg
81+
assert "[REDACTED]" in record.msg
82+
83+
def test_bearer_case_insensitive(self):
84+
record = self._make_record("header: bearer my-secret-token-value")
85+
self.filt.filter(record)
86+
assert "my-secret-token-value" not in record.msg
87+
88+
def test_preserves_non_secret_messages(self):
89+
record = self._make_record("Processing weather request for New York")
90+
self.filt.filter(record)
91+
assert record.msg == "Processing weather request for New York"
92+
93+
def test_redacts_bearer_in_tuple_args(self):
94+
record = self._make_record("Header: %s", ("Bearer secret123",))
95+
self.filt.filter(record)
96+
assert "secret123" not in record.args[0]
97+
98+
def test_redacts_bearer_in_dict_args(self):
99+
record = self._make_record("%(auth)s")
100+
record.args = {"auth": "Bearer secret123"}
101+
self.filt.filter(record)
102+
assert "secret123" not in record.args["auth"]
103+
104+
def test_always_returns_true(self):
105+
record = self._make_record("Bearer secret123")
106+
assert self.filt.filter(record) is True
107+
108+
def test_redacts_literal_configured_key(self, monkeypatch):
109+
rhoai_key = "a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6"
110+
monkeypatch.setenv("LLM_API_KEY", rhoai_key)
111+
filt = SecretRedactionFilter()
112+
record = self._make_record(f"Sending request with api-key={rhoai_key}")
113+
filt.filter(record)
114+
assert rhoai_key not in record.msg
115+
assert "[REDACTED]" in record.msg
116+
117+
def test_literal_key_redaction_in_args(self, monkeypatch):
118+
rhoai_key = "a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6"
119+
monkeypatch.setenv("LLM_API_KEY", rhoai_key)
120+
filt = SecretRedactionFilter()
121+
record = self._make_record("key=%s", (rhoai_key,))
122+
filt.filter(record)
123+
assert rhoai_key not in record.args[0]
124+
125+
def test_short_key_not_redacted(self, monkeypatch):
126+
monkeypatch.setenv("LLM_API_KEY", "dummy")
127+
filt = SecretRedactionFilter()
128+
record = self._make_record("Using dummy config for testing dummy values")
129+
filt.filter(record)
130+
assert "dummy" in record.msg
131+
132+
def test_no_crash_when_key_unset(self, monkeypatch):
133+
monkeypatch.delenv("LLM_API_KEY", raising=False)
134+
filt = SecretRedactionFilter()
135+
record = self._make_record("Normal log message")
136+
filt.filter(record)
137+
assert record.msg == "Normal log message"
138+
139+
140+
class TestConfigurationApiKeyValidation:
141+
"""Test API key validation logic."""
142+
143+
def test_dummy_key_with_remote_api_is_invalid(self, monkeypatch):
144+
monkeypatch.setenv("LLM_API_BASE", "https://api.openai.com/v1")
145+
monkeypatch.setenv("LLM_API_KEY", "dummy")
146+
assert Configuration().has_valid_api_key is False
147+
148+
def test_dummy_key_with_localhost_is_valid(self):
149+
config = Configuration() # defaults: localhost + dummy
150+
assert config.has_valid_api_key is True
151+
152+
def test_empty_key_with_remote_api_is_invalid(self, monkeypatch):
153+
monkeypatch.setenv("LLM_API_BASE", "https://api.openai.com/v1")
154+
monkeypatch.setenv("LLM_API_KEY", "")
155+
assert Configuration().has_valid_api_key is False
156+
157+
def test_placeholder_keys_with_remote_are_invalid(self, monkeypatch):
158+
monkeypatch.setenv("LLM_API_BASE", "https://api.openai.com/v1")
159+
for key in ["changeme", "your-api-key-here"]:
160+
monkeypatch.setenv("LLM_API_KEY", key)
161+
assert Configuration().has_valid_api_key is False, f"'{key}' should be invalid"
162+
163+
def test_placeholder_keys_with_local_llm_are_valid(self, monkeypatch):
164+
monkeypatch.setenv("LLM_API_BASE", "http://localhost:11434/v1")
165+
for key in ["dummy", "changeme", ""]:
166+
monkeypatch.setenv("LLM_API_KEY", key)
167+
assert Configuration().has_valid_api_key is True
168+
169+
def test_real_key_is_valid(self, monkeypatch):
170+
monkeypatch.setenv("LLM_API_BASE", "https://api.openai.com/v1")
171+
monkeypatch.setenv("LLM_API_KEY", "sk-proj-realkey123")
172+
assert Configuration().has_valid_api_key is True
173+
174+
def test_rhoai_maas_key_is_valid(self, monkeypatch):
175+
monkeypatch.setenv(
176+
"LLM_API_BASE",
177+
"https://model--maas-apicast-production.apps.prod.rhoai.rh-aiservices-bu.com:443/v1",
178+
)
179+
monkeypatch.setenv("LLM_API_KEY", "a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6")
180+
assert Configuration().has_valid_api_key is True
181+
182+
def test_127_is_local(self, monkeypatch):
183+
monkeypatch.setenv("LLM_API_BASE", "http://127.0.0.1:8080/v1")
184+
assert Configuration().has_valid_api_key is True

0 commit comments

Comments
 (0)