| 1 | from enum import Enum |
| 2 | from typing import Any |
| 3 | from typing import Dict |
| 4 | from typing import List |
| 5 | from typing import Optional |
| 6 | |
| 7 | from fastapi import HTTPException |
| 8 | from pydantic import BaseModel |
| 9 | from pydantic import Field |
| 10 | from pydantic import field_validator |
| 11 | |
| 12 | |
| 13 | class MCPServerType(str, Enum): |
| 14 | """Enumeration of available MCP servers""" |
| 15 | |
| 16 | WAZUH_INDEXER = "wazuh-indexer" |
| 17 | WAZUH_MANAGER = "wazuh-manager" |
| 18 | COPILOT = "copilot" |
| 19 | VELOCIRAPTOR = "velociraptor" |
| 20 | THREAT_INTEL = "threat-intel" |
| 21 | CYBER_NEWS = "cyber-news" |
| 22 | KNOWLEDGEBASE = "knowledgebase" |
| 23 | ATTACK_SURFACE = "attack-surface" |
| 24 | |
| 25 | |
| 26 | class MCPServerConfig(BaseModel): |
| 27 | """Configuration for MCP server connection""" |
| 28 | |
| 29 | name: str = Field(..., description="Name of the MCP server") |
| 30 | transport: str = Field(default="sse", description="Transport type (sse, stdio)") |
| 31 | url: Optional[str] = Field( |
| 32 | default=None, |
| 33 | description="URL for the MCP server endpoint (for sse transport)", |
| 34 | ) |
| 35 | command: Optional[str] = Field( |
| 36 | default=None, |
| 37 | description="Command to run (for stdio transport)", |
| 38 | ) |
| 39 | args: Optional[List[str]] = Field( |
| 40 | default=None, |
| 41 | description="Arguments for the command (for stdio transport)", |
| 42 | ) |
| 43 | headers: Optional[Dict[str, str]] = Field( |
| 44 | default=None, |
| 45 | description="Headers for authentication", |
| 46 | ) |
| 47 | env: Optional[Dict[str, str]] = Field( |
| 48 | default=None, |
| 49 | description="Environment variables for the process", |
| 50 | ) |
| 51 | |
| 52 | |
| 53 | class MCPQuery(BaseModel): |
| 54 | """Query request for MCP server""" |
| 55 | |
| 56 | input: str = Field(..., description="The query/input to send to the MCP server") |
| 57 | server_name: str = Field( |
| 58 | default="opensearch-mcp-server", |
| 59 | description="Name of the MCP server to use", |
| 60 | ) |
| 61 | verbose: bool = Field(default=True, description="Enable verbose output") |
| 62 | |
| 63 | |
| 64 | class MCPQueryRequest(BaseModel): |
| 65 | """Complete MCP query request""" |
| 66 | |
| 67 | input: str = Field(..., description="The query/input to send to the MCP server") |
| 68 | mcp_server: MCPServerType = Field(..., description="MCP server to use for the query") |
| 69 | verbose: Optional[bool] = Field(default=True, description="Enable verbose output") |
| 70 | |
| 71 | @field_validator("mcp_server", mode="before") |
| 72 | @classmethod |
| 73 | def validate_mcp_server(cls, v): |
| 74 | """Validate that the MCP server type is one of the allowed values""" |
| 75 | if isinstance(v, str): |
| 76 | # Check if the string value is valid |
| 77 | valid_values = [server.value for server in MCPServerType] |
| 78 | if v not in valid_values: |
| 79 | raise HTTPException(status_code=400, detail=f"Invalid MCP server type: '{v}'. Must be one of: {', '.join(valid_values)}") |
| 80 | return v |
| 81 | |
| 82 | |
| 83 | class VulnerabilityInfo(BaseModel): |
| 84 | """Structured vulnerability information""" |
| 85 | |
| 86 | agent_name: Optional[str] = Field(None, description="Name of the agent") |
| 87 | vulnerability_id: Optional[str] = Field( |
| 88 | None, |
| 89 | description="Vulnerability identifier", |
| 90 | ) |
| 91 | score: Optional[float] = Field(None, description="Vulnerability score") |
| 92 | severity: Optional[str] = Field(None, description="Vulnerability severity level") |
| 93 | description: Optional[str] = Field(None, description="Vulnerability description") |
| 94 | cve_id: Optional[str] = Field(None, description="CVE identifier if available") |
| 95 | |
| 96 | |
| 97 | class ClusterHealthInfo(BaseModel): |
| 98 | """Structured cluster health information""" |
| 99 | |
| 100 | status: Optional[str] = Field( |
| 101 | None, |
| 102 | description="Cluster status (green, yellow, red)", |
| 103 | ) |
| 104 | number_of_nodes: Optional[int] = Field(None, description="Total number of nodes") |
| 105 | active_primary_shards: Optional[int] = Field( |
| 106 | None, |
| 107 | description="Number of active primary shards", |
| 108 | ) |
| 109 | active_shards: Optional[int] = Field( |
| 110 | None, |
| 111 | description="Total number of active shards", |
| 112 | ) |
| 113 | relocating_shards: Optional[int] = Field( |
| 114 | None, |
| 115 | description="Number of relocating shards", |
| 116 | ) |
| 117 | initializing_shards: Optional[int] = Field( |
| 118 | None, |
| 119 | description="Number of initializing shards", |
| 120 | ) |
| 121 | unassigned_shards: Optional[int] = Field( |
| 122 | None, |
| 123 | description="Number of unassigned shards", |
| 124 | ) |
| 125 | |
| 126 | |
| 127 | class StructuredAgentResponse(BaseModel): |
| 128 | """Structured response from the MCP agent with consistent format""" |
| 129 | |
| 130 | response: str = Field( |
| 131 | ..., |
| 132 | description="Human-readable response in markdown format - the main answer to the user's query", |
| 133 | ) |
| 134 | thinking_process: Optional[str] = Field( |
| 135 | None, |
| 136 | description="Agent's step-by-step reasoning and exploration process", |
| 137 | ) |
| 138 | |
| 139 | |
| 140 | class MCPQueryResponse(BaseModel): |
| 141 | """Response from MCP query""" |
| 142 | |
| 143 | message: str |
| 144 | success: bool |
| 145 | result: Optional[Any] = Field(None, description="The result from the MCP agent") |
| 146 | structured_result: Optional[StructuredAgentResponse] = Field( |
| 147 | None, |
| 148 | description="Structured response from agent", |
| 149 | ) |
| 150 | execution_time: Optional[float] = Field( |
| 151 | None, |
| 152 | description="Time taken to execute the query", |
| 153 | ) |
| 154 | |
| 155 | |
| 156 | class MCPServerInfo(BaseModel): |
| 157 | """Information about an available MCP server""" |
| 158 | |
| 159 | name: str = Field(..., description="The server name/identifier") |
| 160 | value: str = Field(..., description="The server enum value") |
| 161 | description: str = Field(..., description="Description of what this server does") |
| 162 | capabilities: List[str] = Field(default=[], description="List of server capabilities") |
| 163 | |
| 164 | |
| 165 | class AvailableMCPServersResponse(BaseModel): |
| 166 | """Response containing available MCP servers""" |
| 167 | |
| 168 | servers: List[MCPServerInfo] = Field(..., description="List of available MCP servers") |
| 169 | total_servers: int = Field(..., description="Total number of available servers") |
| 170 | message: str = Field(..., description="Response message") |
| 171 | success: bool = Field(default=True, description="Whether the request was successful") |
| 172 | |
| 173 | |
| 174 | class ExampleQuestion(BaseModel): |
| 175 | """Single example question with metadata""" |
| 176 | |
| 177 | question: str = Field(..., description="The example question text") |
| 178 | description: Optional[str] = Field(None, description="Brief description of what this question does") |
| 179 | category: Optional[str] = Field(None, description="Category of the question (e.g., 'alerts', 'agents', 'health')") |
| 180 | |
| 181 | |
| 182 | class ExampleQuestionsResponse(BaseModel): |
| 183 | """Response containing example questions for a specific MCP server""" |
| 184 | |
| 185 | mcp_server: MCPServerType = Field(..., description="The MCP server these questions are for") |
| 186 | questions: List[ExampleQuestion] = Field(..., description="List of example questions") |
| 187 | total_questions: int = Field(..., description="Total number of example questions") |
| 188 | message: str = Field(..., description="Response message") |
| 189 | success: bool = Field(default=True, description="Whether the request was successful") |