From fce9dc1b2940ee6699d2acf0c9e56456d9899581 Mon Sep 17 00:00:00 2001
From: tanaka takashi <isazaka@gmail.com>
Date: Sat, 27 Jun 2026 01:09:00 +0900
Subject: MCP server機能をセキュア化
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/ollama_web/llm.py | 34 ++++++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 8 deletions(-)

(limited to 'src/ollama_web/llm.py')

diff --git a/src/ollama_web/llm.py b/src/ollama_web/llm.py
index 0ab00a3..11e93d2 100644
--- a/src/ollama_web/llm.py
+++ b/src/ollama_web/llm.py
@@ -11,6 +11,7 @@ import ollama
 from ollama import Message
 
 from .config import settings
+from .mcp import is_dangerous_mcp_tool, redact_secrets
 from .tools.registry import ToolRegistry, default_registry
 
 # Maximum successive tool-call rounds before we force a final answer.
@@ -28,6 +29,25 @@ MAX_TOOL_CONTEXT_CHARS = 24000
 ToolCallable = Callable[..., Any]
 
 
+def _safe_tool_arguments(arguments: Any) -> Any:
+    return redact_secrets(arguments)
+
+
+def _safe_tool_result(result: str) -> str:
+    redacted = redact_secrets(result)
+    return str(redacted)
+
+
+def _execute_tool_with_policy(
+    reg: ToolRegistry,
+    name: str,
+    args: Any,
+) -> str:
+    if is_dangerous_mcp_tool(name):
+        return json.dumps({"error": "MCP tool requires explicit approval"}, ensure_ascii=False)
+    return reg.execute(name, args if isinstance(args, dict) else json.dumps(args))
+
+
 def get_client(host: str | None = None) -> ollama.Client:
     """Return an ``ollama.Client`` configured for the configured host."""
     import httpx
@@ -245,9 +265,9 @@ def chat_with_tools(
         for call in calls:
             name = call.get("name", "")
             args = call.get("arguments", {})
-            yield {"type": "tool_start", "name": name, "arguments": args}
-            result = reg.execute(name, args if isinstance(args, dict) else json.dumps(args))
-            result = _trim_tool_result(result, settings.max_tool_result_chars)
+            yield {"type": "tool_start", "name": name, "arguments": _safe_tool_arguments(args)}
+            result = _execute_tool_with_policy(reg, name, args)
+            result = _safe_tool_result(_trim_tool_result(result, settings.max_tool_result_chars))
             yield {"type": "tool_end", "name": name, "result": result}
             tool_results.append((name, result))
 
@@ -423,18 +443,16 @@ def stream_chat_with_tools(
         for call in tool_calls:
             name = call.get("name", "")
             args = call.get("arguments", {})
-            yield {"type": "tool_start", "name": name, "arguments": args}
+            yield {"type": "tool_start", "name": name, "arguments": _safe_tool_arguments(args)}
             t1 = time.time()
             try:
-                result = reg.execute(
-                    name, args if isinstance(args, dict) else json.dumps(args)
-                )
+                result = _execute_tool_with_policy(reg, name, args)
             except Exception as exc:  # noqa: BLE001
                 logger.warning("tool %s execution failed in %.2fs: %s", name, time.time() - t1, exc)
                 result = f"Tool execution failed: {exc}"
             # Always truncate per-tool result immediately so it does not bloat
             # the context even before we reach the cumulative limit.
-            result = _trim_tool_result(result, settings.max_tool_result_chars)
+            result = _safe_tool_result(_trim_tool_result(result, settings.max_tool_result_chars))
             logger.info(
                 "tool %s executed len=%d elapsed=%.2fs",
                 name,
-- 
cgit v1.2.3