diff --git a/mcp_server/tools/create_http_application.py b/mcp_server/tools/create_http_application.py index 2fe6e55..bf38cf0 100644 --- a/mcp_server/tools/create_http_application.py +++ b/mcp_server/tools/create_http_application.py @@ -2,6 +2,7 @@ from pydantic import BaseModel, Field from utils.request import post_slce_api from tools import Tool, ABCTool, tools from urllib.parse import urlparse + @tools.register class CreateHttpApplication(BaseModel, ABCTool): domain: str = Field(default="",description="application domain, if empty, match all domain") diff --git a/mcp_server/tools/get_attack_events.py b/mcp_server/tools/get_attack_events.py index e5044bb..d506a66 100644 --- a/mcp_server/tools/get_attack_events.py +++ b/mcp_server/tools/get_attack_events.py @@ -1,9 +1,9 @@ from pydantic import BaseModel, Field from utils.request import get_slce_api from tools import Tool, ABCTool, tools -from urllib.parse import urlparse + @tools.register -class CreateHttpApplication(BaseModel, ABCTool): +class GetAttackEvents(BaseModel, ABCTool): ip: str = Field(default="", description="the attacker's client IP address") size: int = Field(default=10, min=1, max=100, description="the number of results to return") start: str = Field(default="", description="start time, millisecond timestamp") @@ -12,13 +12,9 @@ class CreateHttpApplication(BaseModel, ABCTool): @classmethod async def run(self, arguments:dict) -> str: try: - req = CreateHttpApplication.model_validate(arguments) - parsed_upstream = urlparse(req.upstream) - if parsed_upstream.scheme not in ["https", "http"]: - return "invalid upstream scheme" - - if parsed_upstream.hostname == "": - return "invalid upstream host" + req = GetAttackEvents.model_validate(arguments) + if req.size < 1 or req.size > 100: + return "size must be between 1 and 100" except Exception as e: return str(e)