From 56f2424f75753ab81316ec04fbd37f736dbf4f4d Mon Sep 17 00:00:00 2001 From: rohanbuilds-ai Date: Thu, 15 Oct 2026 10:03:58 +0530 Subject: [PATCH] fix(mcp): enforce SystemUtils enabled_tools allow-list SystemUtils MCP server accepted an enabled_tools configuration setting defined in DEFAULT_CONFIG, but create_systemutils_server statically decorated all 8 tool functions with @mcp.tool. Consequently, all tools (including administrative and arbitrary script execution tools) were exposed and callable even when administrators or challenge setups restricted enabled_tools. - Conditionally register tools with FastMCP only when their name appears in enabled_tools. - Gracefully fall back to defaults when enabled_tools is None or non-iterable. - Ensure disabled tools are excluded from list_tools() and reject calls with ToolError. - Add comprehensive test suite in tests/unit/mcp/test_systemutils_server.py. --- finbot/mcp/servers/systemutils/server.py | 30 ++- tests/unit/mcp/test_systemutils_server.py | 279 ++++++++++++++++++++++ 2 files changed, 300 insertions(+), 9 deletions(-) create mode 100644 tests/unit/mcp/test_systemutils_server.py diff --git a/finbot/mcp/servers/systemutils/server.py b/finbot/mcp/servers/systemutils/server.py index 7f2ddd52..b5ec18f0 100644 --- a/finbot/mcp/servers/systemutils/server.py +++ b/finbot/mcp/servers/systemutils/server.py @@ -10,7 +10,7 @@ import logging from datetime import UTC, datetime -from typing import Any +from typing import Any, Callable from fastmcp import FastMCP @@ -42,7 +42,19 @@ def create_systemutils_server( config = {**DEFAULT_CONFIG, **(server_config or {})} mcp = FastMCP("SystemUtils") - @mcp.tool + raw_enabled = config.get("enabled_tools") + if raw_enabled is None or not hasattr(raw_enabled, "__iter__"): + raw_enabled = DEFAULT_CONFIG["enabled_tools"] + elif isinstance(raw_enabled, str): + raw_enabled = [raw_enabled] + enabled_set = set(raw_enabled) + + def _tool(fn: Callable[..., Any]) -> Callable[..., Any]: + if fn.__name__ in enabled_set: + return mcp.tool(fn) + return fn + + @_tool def run_diagnostics(command: str) -> dict[str, Any]: """Run a system diagnostic command for health checks. @@ -65,7 +77,7 @@ def run_diagnostics(command: str) -> dict[str, Any]: "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def manage_storage(action: str, target: str) -> dict[str, Any]: """Manage file storage operations. @@ -88,7 +100,7 @@ def manage_storage(action: str, target: str) -> dict[str, Any]: "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def rotate_logs(service: str, options: str = "") -> dict[str, Any]: """Rotate service logs for maintenance. @@ -112,7 +124,7 @@ def rotate_logs(service: str, options: str = "") -> dict[str, Any]: "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def database_maintenance(query: str) -> dict[str, Any]: """Run a database maintenance query. @@ -134,7 +146,7 @@ def database_maintenance(query: str) -> dict[str, Any]: "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def network_request( url: str, method: str = "GET", headers: str = "", body: str = "" ) -> dict[str, Any]: @@ -163,7 +175,7 @@ def network_request( "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def read_config(filepath: str) -> dict[str, Any]: """Read a system or application configuration file. @@ -186,7 +198,7 @@ def read_config(filepath: str) -> dict[str, Any]: "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def manage_users(action: str, username: str, role: str = "viewer") -> dict[str, Any]: """Manage system user accounts. @@ -211,7 +223,7 @@ def manage_users(action: str, username: str, role: str = "viewer") -> dict[str, "timestamp": datetime.now(UTC).isoformat().replace("+00:00", "Z"), } - @mcp.tool + @_tool def execute_script(script_content: str, interpreter: str = "bash") -> dict[str, Any]: """Execute a maintenance script on the system. diff --git a/tests/unit/mcp/test_systemutils_server.py b/tests/unit/mcp/test_systemutils_server.py new file mode 100644 index 00000000..66f81d98 --- /dev/null +++ b/tests/unit/mcp/test_systemutils_server.py @@ -0,0 +1,279 @@ +"""Unit tests for SystemUtils MCP Server. + +Validates that SystemUtils MCP server correctly enforces the configured +'enabled_tools' allow-list, exposing only permitted tools and rejecting +unregistered tools, while preserving mock execution behaviors and defaults. +""" + +from unittest.mock import MagicMock + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError + +from finbot.core.auth.session import SessionContext +from finbot.mcp.factory import _apply_tool_overrides +from finbot.mcp.servers.systemutils.server import ( + DEFAULT_CONFIG, + create_systemutils_server, +) + +ALL_SYSTEMUTILS_TOOLS = [ + "run_diagnostics", + "manage_storage", + "rotate_logs", + "database_maintenance", + "network_request", + "read_config", + "manage_users", + "execute_script", +] + + +@pytest.fixture +def session_context(): + """Create a mock SessionContext with a valid namespace.""" + ctx = MagicMock(spec=SessionContext) + ctx.namespace = "test-systemutils-namespace" + ctx.user_id = "test-user-001" + ctx.session_id = "test-session-001" + return ctx + + +class TestSystemUtilsServerToolAllowlist: + """Tests verifying enforcement of enabled_tools allow-list.""" + + @pytest.mark.asyncio + async def test_default_config_enables_all_eight_tools(self, session_context): + """When server_config is None, all 8 default tools must be exposed.""" + server = create_systemutils_server(session_context, server_config=None) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert len(tool_names) == 8 + assert set(tool_names) == set(DEFAULT_CONFIG["enabled_tools"]) + assert set(tool_names) == set(ALL_SYSTEMUTILS_TOOLS) + + @pytest.mark.asyncio + async def test_empty_server_config_enables_all_eight_tools(self, session_context): + """When server_config is an empty dict, all 8 default tools must be exposed.""" + server = create_systemutils_server(session_context, server_config={}) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert len(tool_names) == 8 + assert set(tool_names) == set(DEFAULT_CONFIG["enabled_tools"]) + + @pytest.mark.asyncio + async def test_server_config_none_enabled_tools_falls_back_to_defaults(self, session_context): + """When enabled_tools is explicitly None, fallback to all default tools.""" + server = create_systemutils_server(session_context, server_config={"enabled_tools": None}) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert len(tool_names) == 8 + assert set(tool_names) == set(ALL_SYSTEMUTILS_TOOLS) + + @pytest.mark.asyncio + async def test_server_config_non_iterable_falls_back_to_defaults(self, session_context): + """When enabled_tools is a non-iterable type, safely fallback to defaults.""" + for invalid_val in [12345, True]: + server = create_systemutils_server( + session_context, server_config={"enabled_tools": invalid_val} + ) + async with Client(server) as client: + tools = await client.list_tools() + assert len(tools) == 8 + + @pytest.mark.asyncio + async def test_enabled_tools_subset_exposes_only_selected_tools(self, session_context): + """Only explicitly permitted tools must be exposed in list_tools().""" + selected = ["run_diagnostics", "read_config"] + server = create_systemutils_server( + session_context, server_config={"enabled_tools": selected} + ) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert len(tool_names) == 2 + assert set(tool_names) == set(selected) + assert "execute_script" not in tool_names + assert "manage_storage" not in tool_names + + @pytest.mark.asyncio + async def test_enabled_tools_single_tool(self, session_context): + """A single configured tool exposes only that specific tool.""" + server = create_systemutils_server( + session_context, server_config={"enabled_tools": ["execute_script"]} + ) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert tool_names == ["execute_script"] + + @pytest.mark.asyncio + async def test_enabled_tools_empty_list_disables_all_tools(self, session_context): + """When enabled_tools is empty list, no tools are exposed.""" + server = create_systemutils_server(session_context, server_config={"enabled_tools": []}) + async with Client(server) as client: + tools = await client.list_tools() + assert len(tools) == 0 + + @pytest.mark.asyncio + async def test_enabled_tools_as_string_coerced_to_single_item(self, session_context): + """If enabled_tools is passed as a bare string, treat as single-item list.""" + server = create_systemutils_server( + session_context, server_config={"enabled_tools": "network_request"} + ) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert tool_names == ["network_request"] + + @pytest.mark.asyncio + async def test_enabled_tools_as_tuple_or_set(self, session_context): + """Tuples and sets are valid iterables for enabled_tools.""" + server = create_systemutils_server( + session_context, + server_config={"enabled_tools": ("manage_storage", "rotate_logs")}, + ) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = {t.name for t in tools} + + assert tool_names == {"manage_storage", "rotate_logs"} + + @pytest.mark.asyncio + async def test_unknown_tools_in_allowlist_safely_ignored(self, session_context): + """Unrecognized tool names in enabled_tools do not cause crashes or ghost tools.""" + server = create_systemutils_server( + session_context, + server_config={"enabled_tools": ["run_diagnostics", "nonexistent_destructive_tool"]}, + ) + async with Client(server) as client: + tools = await client.list_tools() + tool_names = [tool.name for tool in tools] + + assert tool_names == ["run_diagnostics"] + + +class TestSystemUtilsToolExecutionAndSecurity: + """Tests verifying execution of enabled tools and rejection of disabled tools.""" + + @pytest.mark.asyncio + async def test_calling_enabled_tool_succeeds(self, session_context): + """An enabled tool can be invoked and returns mock execution output.""" + server = create_systemutils_server( + session_context, + server_config={"enabled_tools": ["run_diagnostics"]}, + ) + async with Client(server) as client: + res = await client.call_tool("run_diagnostics", {"command": "disk_usage"}) + assert res.data["command"] == "disk_usage" + assert res.data["status"] == "completed" + assert res.data["exit_code"] == 0 + assert "disk_usage" in res.data["output"] + + @pytest.mark.asyncio + async def test_calling_disabled_tool_raises_tool_error(self, session_context): + """Calling a disabled tool raises fastmcp.exceptions.ToolError.""" + server = create_systemutils_server( + session_context, + server_config={"enabled_tools": ["run_diagnostics"]}, + ) + async with Client(server) as client: + with pytest.raises(ToolError, match="Unknown tool"): + await client.call_tool( + "execute_script", + {"script_content": "rm -rf /", "interpreter": "bash"}, + ) + + @pytest.mark.asyncio + async def test_all_individual_tools_execute_correctly_when_enabled(self, session_context): + """Each of the 8 tools functions properly when enabled in the server.""" + server = create_systemutils_server( + session_context, + server_config={ + "enabled_tools": ALL_SYSTEMUTILS_TOOLS, + "mock_hostname": "test-finbot-node", + "mock_os": "TestOS 1.0", + }, + ) + async with Client(server) as client: + # 1. run_diagnostics + diag = await client.call_tool("run_diagnostics", {"command": "process_list"}) + assert diag.data["command"] == "process_list" + assert diag.data["hostname"] == "test-finbot-node" + assert diag.data["os"] == "TestOS 1.0" + assert diag.data["status"] == "completed" + + # 2. manage_storage + storage = await client.call_tool( + "manage_storage", {"action": "check", "target": "/var/log"} + ) + assert storage.data["action"] == "check" + assert storage.data["target"] == "/var/log" + assert storage.data["status"] == "completed" + + # 3. rotate_logs + logs = await client.call_tool("rotate_logs", {"service": "payments"}) + assert logs.data["service"] == "payments" + assert logs.data["status"] == "completed" + + # 4. database_maintenance + db = await client.call_tool("database_maintenance", {"query": "VACUUM ANALYZE"}) + assert db.data["query"] == "VACUUM ANALYZE" + assert db.data["status"] == "completed" + + # 5. network_request + net = await client.call_tool( + "network_request", + {"url": "https://api.internal/health", "method": "GET"}, + ) + assert net.data["url"] == "https://api.internal/health" + assert net.data["status_code"] == 200 + + # 6. read_config + cfg = await client.call_tool("read_config", {"filepath": "/etc/finbot/app.conf"}) + assert cfg.data["filepath"] == "/etc/finbot/app.conf" + assert "DATABASE_URL" in cfg.data["content"] + + # 7. manage_users + usr = await client.call_tool( + "manage_users", {"action": "list", "username": "admin", "role": "admin"} + ) + assert usr.data["action"] == "list" + assert usr.data["username"] == "admin" + + # 8. execute_script + script = await client.call_tool( + "execute_script", + {"script_content": "echo 'safe test'", "interpreter": "bash"}, + ) + assert script.data["interpreter"] == "bash" + assert script.data["exit_code"] == 0 + + @pytest.mark.asyncio + async def test_tool_overrides_applied_only_to_enabled_tools(self, session_context): + """Tool poisoning overrides via factory apply cleanly to enabled tools.""" + server = create_systemutils_server( + session_context, + server_config={"enabled_tools": ["manage_storage"]}, + ) + overrides = { + "manage_storage": {"description": "Poisoned storage description"}, + "execute_script": {"description": "Poisoned script description"}, + } + await _apply_tool_overrides(server, overrides) + + async with Client(server) as client: + tools = await client.list_tools() + assert len(tools) == 1 + assert tools[0].name == "manage_storage" + assert tools[0].description == "Poisoned storage description"