refactor: tools

Author: dhsifss

mcp_server/README.md

@@ -33,7 +33,8 @@ class Hello(BaseModel, ABCTool):
     # run is tool logic, must use classmethod
     @classmethod
     async def run(arguments: dict) -> str:
-        return f"Hello {arguments['name']}"
+        req = Hello.model_validate(arguments)
+        return f"Hello {req.name}"
 
     # tool description, must use classmethod
     @classmethod

mcp_server/tools/create_http_application.py

@@ -5,28 +5,26 @@
 @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")
+    port: int = Field(min=1, max=65535,description="application listen port, must between 1 and 65535")
     upstream: str = Field(description="application proxy address, must be a valid url")
 
     @classmethod
     async def run(self, arguments:dict) -> str:
-        port = arguments["port"]
-        upstream = arguments["upstream"]
-        domain = arguments["domain"]
+        try:
+            req = CreateHttpApplication.model_validate(arguments)
+            parsed_upstream = urlparse(req.upstream)
+            if parsed_upstream.scheme != "https" and parsed_upstream.scheme != "http":
+                return "invalid upstream scheme"
 
-        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"
+            if parsed_upstream.hostname == "":
+                return "invalid upstream host"
+        except Exception as e:
+            return str(e)
 
         return await post_slce_api("/api/open/site",{
-                "server_names": [domain],
-                "ports": [ str(port) ],
-                "upstreams": [ upstream ],
+                "server_names": [req.domain],
+                "ports": [ str(req.port) ],
+                "upstreams": [ req.upstream ],
                 "type": 0,
                 "static_default": 1,
                 "health_check": True,

mcp_server/tools/create_ip_custom_rule.py

@@ -1,30 +1,34 @@
 from pydantic import BaseModel, Field
 from utils.request import post_slce_api
 from tools import Tool, ABCTool, tools
+import ipaddress
 
 @tools.register
 class CreateIPCustomRule(BaseModel, ABCTool):
     ip: str = Field(description="request ip to allow or block")
-    action: int = Field(description="1: block, 0: allow")
+    action: int = Field(min=0, max=1,description="1: block, 0: allow")
 
     @classmethod
     async def run(self, arguments:dict) -> str:
-        print(arguments)
-        ip = arguments["ip"]
-        action = arguments["action"]
+        try:
+            req = CreateIPCustomRule.model_validate(arguments)
+            ipaddress.ip_address(req.ip)
+        except Exception as e:
+            return str(e)
+
         name = ""
-        match action:
+        match req.action:
             case 0:
                 name += "allow "
             case 1:
                 name += "block "
             case _:
                 return "invalid action"
 
-        if not ip or ip == "":
+        if not req.ip or req.ip == "":
             return "ip is required"
 
-        name += f"ip: {ip}"
+        name += f"ip: {req.ip}"
 
         return await post_slce_api("/api/open/policy",{
             "name": name,
@@ -34,12 +38,12 @@ async def run(self, arguments:dict) -> str:
                     {
                         "k": "src_ip",
                         "op": "eq",
-                        "v": [ip],
+                        "v": [req.ip],
                         "sub_k": ""
                     },
                 ]
             ],
-            "action": action
+            "action": req.action
         })
 
     @classmethod

mcp_server/tools/create_path_custom_rule.py

@@ -5,25 +5,28 @@
 @tools.register
 class CreatePathCustomRule(BaseModel, ABCTool):
     path: str = Field(description="request path to block or allow")
-    action: int = Field(description="1: block, 0: allow")
+    action: int = Field(min=0, max=1,description="1: block, 0: allow")
 
     @classmethod
     async def run(self, arguments:dict) -> str:
-        path = arguments["path"]
-        action = arguments["action"]
+        try:
+            req = CreatePathCustomRule.model_validate(arguments)
+        except Exception as e:
+            return str(e)
+
         name = ""
-        match action:
+        match req.action:
             case 0:
                 name += "allow "
             case 1:
                 name += "block "
             case _:
                 return "invalid action"
 
-        if not path or path == "":
+        if not req.path or req.path == "":
             return "path is required"
 
-        name += f"path: {path}"
+        name += f"path: {req.path}"
 
         return await post_slce_api("/api/open/policy",{
             "name": name,
@@ -33,12 +36,12 @@ async def run(self, arguments:dict) -> str:
                     {
                         "k": "uri_no_query",
                         "op": "eq",
-                        "v": [path],
+                        "v": [req.path],
                         "sub_k": ""
                     },
                 ]
             ],
-            "action": action
+            "action": req.action
         })
 
     @classmethod