@cryptotaxi247 / CoPilot
Code Commits Issues Pulls Sessions CI/CD
main
main
py 231 lines 8.24 KB
Raw
1 from enum import Enum
2 from typing import Dict
3 from typing import Optional
4
5 import httpx
6 from loguru import logger
7
8 from app.integrations.copilot_mcp.schema.copilot_mcp import MCPQueryRequest
9 from app.integrations.copilot_mcp.schema.copilot_mcp import MCPQueryResponse
10 from app.integrations.copilot_mcp.schema.copilot_mcp import MCPServerType
11
12
13 class MCPServiceType(str, Enum):
14 """Enumeration of MCP service deployment types"""
15
16 LOCAL = "local"
17 CLOUD = "cloud"
18
19
20 class MCPServerConfig:
21 """Configuration for MCP server endpoints"""
22
23 def __init__(self, service_type: MCPServiceType, endpoint: str):
24 self.service_type = service_type
25 self.endpoint = endpoint
26
27
28 class MCPService:
29 """Service for handling MCP queries with modular server routing"""
30
31 # Base URLs for different service types
32 _BASE_URLS = {
33 MCPServiceType.LOCAL: "http://copilot-mcp/mcp",
34 MCPServiceType.CLOUD: "https://mcp.socfortress.co/query",
35 }
36
37 # Define the mapping of MCP server types to their configurations
38 _SERVER_CONFIGS: Dict[MCPServerType, MCPServerConfig] = {
39 # Local services
40 MCPServerType.WAZUH_INDEXER: MCPServerConfig(MCPServiceType.LOCAL, "opensearch-query"),
41 MCPServerType.WAZUH_MANAGER: MCPServerConfig(MCPServiceType.LOCAL, "wazuh-query"),
42 MCPServerType.COPILOT: MCPServerConfig(MCPServiceType.LOCAL, "mysql-query"),
43 MCPServerType.VELOCIRAPTOR: MCPServerConfig(MCPServiceType.LOCAL, "velociraptor-query"),
44 # Cloud services
45 MCPServerType.THREAT_INTEL: MCPServerConfig(MCPServiceType.CLOUD, "threat_intel"),
46 MCPServerType.CYBER_NEWS: MCPServerConfig(MCPServiceType.CLOUD, "cyber_news"),
47 MCPServerType.KNOWLEDGEBASE: MCPServerConfig(MCPServiceType.CLOUD, "knowledgebase"),
48 MCPServerType.ATTACK_SURFACE: MCPServerConfig(MCPServiceType.CLOUD, "attack_surface"),
49 }
50
51 @classmethod
52 def get_server_config(cls, mcp_server: MCPServerType) -> MCPServerConfig:
53 """
54 Get the configuration for a specific MCP server type.
55
56 Args:
57 mcp_server: The MCP server type
58
59 Returns:
60 MCPServerConfig: The configuration for the server
61
62 Raises:
63 ValueError: If the server type is not supported
64 """
65 config = cls._SERVER_CONFIGS.get(mcp_server)
66 if not config:
67 raise ValueError(f"Unsupported MCP server type: {mcp_server}")
68 return config
69
70 @classmethod
71 def build_full_url(cls, mcp_server: MCPServerType) -> str:
72 """
73 Build the full URL for a specific MCP server type.
74
75 Args:
76 mcp_server: The MCP server type
77
78 Returns:
79 str: The full URL for the server endpoint
80 """
81 config = cls.get_server_config(mcp_server)
82 base_url = cls._BASE_URLS[config.service_type]
83 return f"{base_url}/{config.endpoint}"
84
85 @classmethod
86 def is_cloud_service(cls, mcp_server: MCPServerType) -> bool:
87 """
88 Check if the MCP server is a cloud service.
89
90 Args:
91 mcp_server: The MCP server type
92
93 Returns:
94 bool: True if it's a cloud service, False if local
95 """
96 config = cls.get_server_config(mcp_server)
97 return config.service_type == MCPServiceType.CLOUD
98
99 @classmethod
100 async def execute_query(cls, data: MCPQueryRequest, license_key: Optional[str] = None) -> MCPQueryResponse:
101 """
102 Execute a query on the appropriate MCP server based on the request.
103
104 Args:
105 data: The MCP query request containing the server type and query
106 license_key: Optional license key for cloud services authentication
107
108 Returns:
109 MCPQueryResponse: The response from the MCP server
110
111 Raises:
112 httpx.HTTPError: If the HTTP request fails
113 ValueError: If the server type is not supported
114 """
115 try:
116 # Get the full URL for the specified server
117 full_url = cls.build_full_url(data.mcp_server)
118 is_cloud = cls.is_cloud_service(data.mcp_server)
119
120 logger.info(f"Sending MCP query to {data.mcp_server.value} ({'cloud' if is_cloud else 'local'}) at {full_url}")
121 logger.debug(f"Query data: {data.model_dump()}")
122
123 # Set different timeout for cloud vs local services
124 timeout = 300 if is_cloud else 300
125
126 # Prepare headers
127 headers = {"Content-Type": "application/json"}
128
129 # Add license key as x-api-key header for cloud services
130 if is_cloud and license_key:
131 headers["x-api-key"] = license_key
132 logger.debug("Added x-api-key header for cloud service")
133 elif is_cloud and not license_key:
134 logger.warning(f"No license key provided for cloud service {data.mcp_server.value}")
135
136 async with httpx.AsyncClient() as client:
137 response = await client.post(
138 full_url,
139 json=data.model_dump(),
140 headers=headers,
141 timeout=timeout,
142 )
143
144 # Raise an exception for HTTP error status codes
145 response.raise_for_status()
146
147 logger.info(f"Successfully received response from {data.mcp_server.value}: {response.json()}")
148 return MCPQueryResponse(**response.json())
149
150 except ValueError as e:
151 logger.error(f"Invalid server type: {str(e)}")
152 return MCPQueryResponse(message=f"Error: {str(e)}", success=False, result=None, structured_result=None, execution_time=0.0)
153
154 except httpx.HTTPError as e:
155 logger.error(f"HTTP error when querying {data.mcp_server.value}: {str(e)}")
156 return MCPQueryResponse(
157 message=f"HTTP error when querying {data.mcp_server.value}: {str(e)}",
158 success=False,
159 result=None,
160 structured_result=None,
161 execution_time=0.0,
162 )
163
164 except Exception as e:
165 logger.error(f"Unexpected error when querying {data.mcp_server.value}: {str(e)}")
166 return MCPQueryResponse(
167 message=f"Unexpected error: {str(e)}",
168 success=False,
169 result=None,
170 structured_result=None,
171 execution_time=0.0,
172 )
173
174 @classmethod
175 def add_local_service(cls, server_type: MCPServerType, endpoint: str) -> None:
176 """
177 Add a new local MCP service.
178
179 Args:
180 server_type: The MCP server type enum
181 endpoint: The endpoint path for the local service
182 """
183 cls._SERVER_CONFIGS[server_type] = MCPServerConfig(MCPServiceType.LOCAL, endpoint)
184
185 @classmethod
186 def add_cloud_service(cls, server_type: MCPServerType, endpoint: str) -> None:
187 """
188 Add a new cloud MCP service.
189
190 Args:
191 server_type: The MCP server type enum
192 endpoint: The endpoint path for the cloud service
193 """
194 cls._SERVER_CONFIGS[server_type] = MCPServerConfig(MCPServiceType.CLOUD, endpoint)
195
196 @classmethod
197 def get_service_info(cls) -> Dict[str, Dict[str, str]]:
198 """
199 Get information about all configured services.
200
201 Returns:
202 Dict containing service information grouped by type
203 """
204 local_services = {}
205 cloud_services = {}
206
207 for server_type, config in cls._SERVER_CONFIGS.items():
208 service_info = {"endpoint": config.endpoint, "full_url": cls.build_full_url(server_type)}
209
210 if config.service_type == MCPServiceType.LOCAL:
211 local_services[server_type.value] = service_info
212 else:
213 cloud_services[server_type.value] = service_info
214
215 return {"local": local_services, "cloud": cloud_services}
216
217
218 # Convenience function to maintain backward compatibility
219 async def post_to_copilot_mcp(data: MCPQueryRequest) -> MCPQueryResponse:
220 """
221 Send a POST request to the appropriate copilot-mcp endpoint based on server type.
222
223 This function maintains backward compatibility while using the new modular service.
224
225 Args:
226 data: The MCP query request
227
228 Returns:
229 MCPQueryResponse: The response from the MCP server
230 """
231 return await MCPService.execute_query(data)