106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Security regression tests for MCP configuration and subprocess isolation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import ai_config
|
|
import mcp_bridge
|
|
import mcp_server
|
|
|
|
|
|
class McpSecurityTest(unittest.TestCase):
|
|
def test_config_tool_never_reveals_or_updates_secrets(self) -> None:
|
|
keys = [
|
|
"AI_API_KEY",
|
|
"GROK_API_KEY",
|
|
"SERVICE_PASSWORD",
|
|
"AI_MODEL",
|
|
]
|
|
patches = [
|
|
mock.patch.object(ai_config, "CONFIGURABLE_KEYS", keys),
|
|
mock.patch.object(ai_config, "AI_API_KEY", "customer-secret-value"),
|
|
mock.patch.object(ai_config, "GROK_API_KEY", "coding-secret-value"),
|
|
mock.patch.object(
|
|
ai_config,
|
|
"SERVICE_PASSWORD",
|
|
"password-value",
|
|
create=True,
|
|
),
|
|
mock.patch.object(ai_config, "AI_MODEL", "old-model"),
|
|
mock.patch.object(ai_config, "load_settings"),
|
|
mock.patch.object(ai_config, "apply_settings"),
|
|
mock.patch.object(ai_config, "build_system_prompt", return_value="prompt"),
|
|
]
|
|
for patcher in patches:
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
visible = mcp_server.get_ai_config()
|
|
self.assertNotIn("customer-secret-value", str(visible))
|
|
self.assertNotIn("coding-secret-value", str(visible))
|
|
self.assertNotIn("password-value", str(visible))
|
|
|
|
result = mcp_server.update_ai_config(
|
|
{
|
|
"AI_API_KEY": "new-customer-secret",
|
|
"GROK_API_KEY": "new-coding-secret",
|
|
"SERVICE_PASSWORD": "new-password",
|
|
"AI_MODEL": "new-model",
|
|
}
|
|
)
|
|
self.assertEqual({"AI_MODEL": "new-model"}, result["changed"])
|
|
self.assertCountEqual(
|
|
["AI_API_KEY", "GROK_API_KEY", "SERVICE_PASSWORD"],
|
|
result["ignored"],
|
|
)
|
|
ai_config.apply_settings.assert_called_once_with(
|
|
{"AI_MODEL": "new-model"},
|
|
persist=True,
|
|
)
|
|
|
|
def test_external_stdio_mcp_gets_minimal_environment(self) -> None:
|
|
with mock.patch.dict(
|
|
os.environ,
|
|
{
|
|
"PATH": "C:\\tools",
|
|
"SYSTEMROOT": "C:\\Windows",
|
|
"WECOM_GROK_API_KEY": "must-not-leak",
|
|
"OPENAI_API_KEY": "must-not-leak-either",
|
|
"PRIVATE_TOKEN": "private",
|
|
},
|
|
clear=True,
|
|
):
|
|
environment = mcp_bridge._minimal_subprocess_env(
|
|
{"SERVER_TOKEN": "explicitly-authorized"}
|
|
)
|
|
|
|
self.assertEqual("C:\\tools", environment["PATH"])
|
|
self.assertEqual("explicitly-authorized", environment["SERVER_TOKEN"])
|
|
self.assertNotIn("WECOM_GROK_API_KEY", environment)
|
|
self.assertNotIn("OPENAI_API_KEY", environment)
|
|
self.assertNotIn("PRIVATE_TOKEN", environment)
|
|
|
|
def test_draft_reply_without_history_does_not_reuse_remote_session(self) -> None:
|
|
with mock.patch.object(
|
|
mcp_server,
|
|
"call_ai_text",
|
|
return_value="草稿",
|
|
) as call:
|
|
result = mcp_server.draft_reply(
|
|
"新问题",
|
|
session_id="customer-1",
|
|
use_history=False,
|
|
)
|
|
|
|
self.assertTrue(result["ok"])
|
|
self.assertIsNone(call.call_args.kwargs["history"])
|
|
self.assertIsNone(call.call_args.kwargs["session_id"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|