refactor: tools

Author: dhsifss

mcp_server/README.md

@@ -21,32 +21,27 @@ This tool used to say hello to someone
 
 ```python
 from pydantic import BaseModel, Field
-from tools import register_tool, Tool
+from tools import Tool, ABCTool, tools
 
+# register to global tools
+@tools.register
 # Hello describe function paramters
-class Hello(BaseModel):
+class Hello(BaseModel, ABCTool):
+    # tools paramters
     name: str = Field(description="username to say hello")
 
-# hello is tool logic
-async def hello(arguments: dict) -> str:
-    """
-    Say hello to someone
-    """
-    return f"Hello {arguments['name']}"
-
-# register tool to global variable
-register_tool(
-    Tool(
-        name="hello",
-        description="say hello to someone",
-        inputSchema=Hello.model_json_schema()
-    ),
-    hello
-)
-```
+    # run is tool logic, must use classmethod
+    @classmethod
+    async def run(arguments: dict) -> str:
+        return f"Hello {arguments['name']}"
 
-2. import this tool in `tools/__init__.py`
+    # tool description, must use classmethod
+    @classmethod
+    def tool(self) -> Tool:
+        return Tool(
+            name="hello",
+            description="say hello to someone",
+            inputSchema=self.model_json_schema()
+        )
 
-```python
-from . import hello
-```
+```

mcp_server/config/config.py

@@ -1,13 +1,21 @@
 import os
+import logging
 
 class Config:
     SAFELINE_ADDRESS: str
     SAFELINE_API_TOKEN: str
     SECRET: str
     LISTEN_PORT: int
     LISTEN_ADDRESS: str
+    DEBUG: bool
 
     def __init__(self):
+        set_log_level()
+
+        if os.getenv("MCP_SERVER_DEBUG"):
+            self.DEBUG = True
+        else:
+            self.DEBUG = False
         self.SAFELINE_ADDRESS = os.getenv("SAFELINE_ADDRESS")
         if self.SAFELINE_ADDRESS:
             self.SAFELINE_ADDRESS = self.SAFELINE_ADDRESS.removesuffix("/")
@@ -26,4 +34,23 @@ def __init__(self):
 
     @staticmethod
     def from_env():
-        return Config()
+        return Config()
+
+
+def set_log_level():
+    level = logging.WARN
+    log_level = os.getenv("MCO_SERVER_LOG_LEVEL")
+    if log_level:
+        match log_level.lower():
+            case "debug":
+                level = logging.DEBUG
+            case "info":
+                level = logging.INFO
+            case "warn":
+                level = logging.WARN
+            case "error":
+                level = logging.ERROR
+            case "critical":
+                level = logging.CRITICAL
+
+    logging.basicConfig(level=level)

mcp_server/middleware.py

@@ -0,0 +1,21 @@
+from starlette.requests import HTTPConnection
+from starlette.responses import PlainTextResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+from config import GLOBAL_CONFIG
+
+
+class AuthenticationMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+    ) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        conn = HTTPConnection(scope)
+        if GLOBAL_CONFIG.SECRET and GLOBAL_CONFIG.SECRET != "" and conn.headers.get("Secret") != GLOBAL_CONFIG.SECRET:
+            response = PlainTextResponse("Unauthorized", status_code=401)
+            await response(scope, receive, send)
+            return
+
+        await self.app(scope, receive, send)

mcp_server/server.py

@@ -6,16 +6,18 @@
 from starlette.requests import Request
 import uvicorn
 from starlette.responses import PlainTextResponse
-import tools
+from tools import tools
 from config import GLOBAL_CONFIG
+from middleware import AuthenticationMiddleware
+from starlette.middleware import Middleware
 
 # Create an MCP server
 mcp_server = Server("SafeLine WAF mcp server")
 sse = mcp.server.sse.SseServerTransport("/messages/")
 
 @mcp_server.list_tools()
 async def list_tools() -> list[Tool]:
-    return tools.ALL_TOOLS
+    return tools.all()
 
 @mcp_server.call_tool()
 async def call_tool(name: str, arguments: dict) -> list[TextContent]:
@@ -29,9 +31,6 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
 
 
 async def handle_sse(request: Request) -> None:
-    if GLOBAL_CONFIG.SECRET and GLOBAL_CONFIG.SECRET != "" and request.headers.get("Secret") != GLOBAL_CONFIG.SECRET:
-        return PlainTextResponse("Unauthorized", status_code=401)
-
     async with sse.connect_sse(
         request.scope, request.receive, request._send
     ) as [read_stream, write_stream]:
@@ -40,9 +39,9 @@ async def handle_sse(request: Request) -> None:
         )
 
 def main():
-    starlette_app = Starlette(debug=True,routes=[
+    starlette_app = Starlette(debug=GLOBAL_CONFIG.DEBUG,routes=[
         Route("/sse", endpoint=handle_sse),
-        Mount("/messages/", app=sse.handle_post_message),
+        Mount("/messages/", app=sse.handle_post_message, middleware=[Middleware(AuthenticationMiddleware)]),
     ])
 
     uvicorn.run(starlette_app, host=GLOBAL_CONFIG.LISTEN_ADDRESS, port=GLOBAL_CONFIG.LISTEN_PORT)

mcp_server/tools/__init__.py

@@ -1,17 +1,49 @@
 from mcp.types import Tool
-from typing import Callable
+from abc import ABC, abstractmethod
+import os
+import importlib
+import logging
+class ABCTool(ABC):
+    @classmethod
+    @abstractmethod
+    async def run(self, arguments:dict) -> str:
+        pass
 
-ALL_TOOLS = []
-TOOL_FUNC_MAP = {}
+    @classmethod
+    @abstractmethod
+    def tool(self) -> Tool:
+        pass
 
-def register_tool(tool: Tool, func: Callable):
-    ALL_TOOLS.append(tool)
-    TOOL_FUNC_MAP[tool.name] = func
+class ToolRegister:
+    _dict: dict[str, ABCTool] = {}
+    
+    @classmethod
+    def register(self, tool: ABCTool) -> ABCTool:
+        tool_name = tool.tool().name
+        logging.info(f"Registering tool: {tool_name}")
+        if tool_name in self._dict:
+            raise ValueError(f"Tool {tool_name} already registered")
+
+        self._dict[tool_name] = tool
+        return tool
 
-async def run(name:str, arguments:dict) -> str:
-    if name not in TOOL_FUNC_MAP:
-        return f"Unknown tool: {name}"
+    def all(self) -> list[Tool]:
+        return [tool.tool() for tool in self._dict.values()]
 
-    return await TOOL_FUNC_MAP[name](arguments)
+    async def run(self, name: str, arguments: dict) -> str:
+        if name not in self._dict:
+            raise ValueError(f"Unknown tool: {name}")
+
+        return await self._dict[name].run(arguments)
+
+def import_all_tools():
+    for module in os.listdir(os.path.dirname(__file__)):
+        if module == "__init__.py" or len(module) < 3 or not module.endswith(".py"):
+            continue
+
+        module_name = module[:-3]
+        importlib.import_module(f".{module_name}", package=__name__)
+
+tools = ToolRegister()
 
-from . import create_black_custom_rule, create_http_application
+import_all_tools()

mcp_server/tools/create_black_custom_rule.py

@@ -1,54 +1,44 @@
 from pydantic import BaseModel, Field
 from utils.request import post_slce_api
-from tools import register_tool, Tool
+from tools import Tool, ABCTool, tools
 from urllib.parse import urlparse
-
-class CreateHttpApplication(BaseModel):
+@tools.register
+class CreateHttpApplication(BaseModel, ABCTool):
     domain: str = Field(default="",description="application domain, if empty, match all domain")
     port: int = Field(description="application listen port, must between 1 and 65535")
     upstream: str = Field(description="application proxy address, must be a valid url")
 
-
-async def create_http_application(arguments:dict) -> str:
-    """
-    Create a new HTTP application.
-
-    Args:
-        domain: application domain
-        port: application listen port
-        upstream: application proxy address
-    """
-
-    port = arguments["port"]
-    upstream = arguments["upstream"]
-    domain = arguments["domain"]
-
-    if port is None or port < 1 or port > 65535:
-        return "invalid port"
-
-    parsed_upstream = urlparse(upstream)
-    if parsed_upstream.scheme != "https" and parsed_upstream.scheme != "http":
-        return "invalid upstream scheme"
-    if parsed_upstream.hostname == "":
-        return "invalid upstream host"
-
-    return await post_slce_api("/api/open/site",{
-            "server_names": [domain],
-            "ports": [ str(port) ],
-            "upstreams": [ upstream ],
-            "type": 0,
-            "static_default": 1,
-            "health_check": True,
-            "load_balance": {
-                "balance_type": 1
-            }
-        })
-
-register_tool(
-    Tool(
-        name="create_http_application",
-        description="在雷池 WAF 上创建一个站点应用",
-        inputSchema=CreateHttpApplication.model_json_schema()
-    ),
-    create_http_application
-)
+    @classmethod
+    async def run(self, arguments:dict) -> str:
+        port = arguments["port"]
+        upstream = arguments["upstream"]
+        domain = arguments["domain"]
+
+        if port is None or port < 1 or port > 65535:
+            return "invalid port"
+
+        parsed_upstream = urlparse(upstream)
+        if parsed_upstream.scheme != "https" and parsed_upstream.scheme != "http":
+            return "invalid upstream scheme"
+        if parsed_upstream.hostname == "":
+            return "invalid upstream host"
+
+        return await post_slce_api("/api/open/site",{
+                "server_names": [domain],
+                "ports": [ str(port) ],
+                "upstreams": [ upstream ],
+                "type": 0,
+                "static_default": 1,
+                "health_check": True,
+                "load_balance": {
+                    "balance_type": 1
+                }
+            })
+
+    @classmethod
+    def tool(self) -> Tool:
+        return Tool(
+            name="create_http_application",
+            description="在雷池 WAF 上创建一个站点应用",
+            inputSchema=self.model_json_schema()
+        )