diff --git a/src/fgai/graylog_mcp.py b/src/fgai/graylog_mcp.py index 5b61423..fab925c 100644 --- a/src/fgai/graylog_mcp.py +++ b/src/fgai/graylog_mcp.py @@ -9,13 +9,22 @@ from urllib import error, request class GraylogMcpClient: """Small Streamable HTTP MCP client used for Graylog connection checks.""" - def __init__(self, url: str, token: str, *, timeout: int = 15, verify_tls: bool = True) -> None: + def __init__(self, url: str, token: str, *, timeout: int = 15, verify_tls: bool = True, use_proxy: bool = False) -> None: self.url = url.rstrip("/") self.token = token self.timeout = timeout self.verify_tls = verify_tls + self.use_proxy = use_proxy self.session_id: str | None = None + def _open(self, req: request.Request, context: ssl.SSLContext | None): + if self.use_proxy: + return request.urlopen(req, timeout=self.timeout, context=context) + handlers = [request.ProxyHandler({})] + if context is not None: + handlers.append(request.HTTPSHandler(context=context)) + return request.build_opener(*handlers).open(req, timeout=self.timeout) + def _call(self, method: str, params: dict[str, object] | None = None, *, notification: bool = False) -> dict[str, object]: payload: dict[str, object] = {"jsonrpc": "2.0", "method": method} if not notification: @@ -33,7 +42,7 @@ class GraylogMcpClient: req = request.Request(self.url, data=json.dumps(payload).encode("utf-8"), method="POST", headers=headers) try: context = None if self.verify_tls else ssl._create_unverified_context() - with request.urlopen(req, timeout=self.timeout, context=context) as response: + with self._open(req, context) as response: self.session_id = response.headers.get("Mcp-Session-Id", self.session_id) body = response.read().decode("utf-8") except error.HTTPError as exc: diff --git a/tests/test_graylog_mcp.py b/tests/test_graylog_mcp.py index d57c1d5..5c78c99 100644 --- a/tests/test_graylog_mcp.py +++ b/tests/test_graylog_mcp.py @@ -1,7 +1,7 @@ import json import unittest from urllib.error import HTTPError -from unittest.mock import patch +from unittest.mock import Mock, patch from fgai.graylog_mcp import GraylogMcpClient @@ -34,7 +34,7 @@ class GraylogMcpTests(unittest.TestCase): _Response({}), _Response({"result": {"tools": [{"name": "search_messages"}, {"name": "aggregate_messages"}]}}), ]) - with patch("fgai.graylog_mcp.request.urlopen", side_effect=responses): + with patch("fgai.graylog_mcp.GraylogMcpClient._open", side_effect=responses): status = GraylogMcpClient("http://graylog/api/mcp", "raw-token").probe() self.assertEqual(status["status"], "connected") self.assertIn("search_messages", status["tools"]) @@ -48,7 +48,7 @@ class GraylogMcpTests(unittest.TestCase): client = GraylogMcpClient("http://graylog/api/mcp", "raw-token") client.session_id = "existing-session" - with patch("fgai.graylog_mcp.request.urlopen", side_effect=responses): + with patch("fgai.graylog_mcp.GraylogMcpClient._open", side_effect=responses): client.probe() self.assertEqual(client.session_id, "session-1") @@ -64,12 +64,25 @@ class GraylogMcpTests(unittest.TestCase): client = GraylogMcpClient("http://graylog/api/mcp", "raw-token") client.session_id = "old-session" - with patch("fgai.graylog_mcp.request.urlopen", side_effect=responses): + with patch("fgai.graylog_mcp.GraylogMcpClient._open", side_effect=responses): result = client.call_tool("search_messages", {"query": "*"}) self.assertEqual(result["result"]["content"], []) self.assertEqual(client.session_id, "session-2") + def test_default_open_bypasses_environment_proxies(self): + client = GraylogMcpClient("http://graylog/api/mcp", "raw-token") + req = object() + opener = Mock() + opener.open.return_value = "response" + + with patch("fgai.graylog_mcp.request.build_opener", return_value=opener) as build_opener: + result = client._open(req, None) + + self.assertEqual(result, "response") + self.assertEqual(opener.open.call_args.kwargs["timeout"], 15) + self.assertEqual(build_opener.call_args.args[0].proxies, {}) + if __name__ == "__main__": unittest.main()