@cryptotaxi247 / CoPilot / commits / 90ca7d87

Threat intel ai (#485)

* feat: add Threat Intel server type and example questions; enhance license error messaging * precommit fixes

taylor_socfortress committed Aug 5, 2025 at 14:13 UTC 90ca7d87472cf88fbbcf3550f1d3ef74e244ce88
5 files changed +177 -26
backend/app/integrations/copilot_mcp/routes/copilot_mcp.py
+28 -5
@@ -1,11 +1,14 @@
1 from typing import Optional
2
3 from fastapi import APIRouter
4 +from fastapi import Depends
5 from fastapi import Query
6 from fastapi import Security
7 from loguru import logger
8 +from sqlalchemy.ext.asyncio import AsyncSession
9
10 from app.auth.routes.auth import AuthHandler
11 +from app.db.db_session import get_db
12 from app.integrations.copilot_mcp.schema.copilot_mcp import AvailableMCPServersResponse
13 from app.integrations.copilot_mcp.schema.copilot_mcp import ExampleQuestionsResponse
14 from app.integrations.copilot_mcp.schema.copilot_mcp import MCPQueryRequest
@@ -15,6 +18,8 @@ from app.integrations.copilot_mcp.services.copilot_mcp import MCPService
18 from app.integrations.copilot_mcp.services.example_questions import (
19 ExampleQuestionsService,
20 )
21 +from app.middleware.license import get_license
22 +from app.middleware.license import is_feature_enabled
23
24 copilot_mcp_router = APIRouter()
25 auth_handler = AuthHandler()
@@ -192,16 +197,34 @@ async def get_question_categories(mcp_server: MCPServerType = Query(..., descrip
197
198 @copilot_mcp_router.post(
199 "/query",
195 - description="Get all disabled rules",
200 + description="Process a query to the appropriate MCP server",
201 dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])],
202 )
198 -async def query_mcp(request: MCPQueryRequest) -> MCPQueryResponse:
203 +async def query_mcp(request: MCPQueryRequest, session: AsyncSession = Depends(get_db)) -> MCPQueryResponse:
204 """
205 Process a query to the MCP agent and return structured response.
206 + Routes the query to the appropriate MCP server based on the selected server type.
207 """
202 - logger.info("Processing MCP query.")
203 -
208 logger.info(f"Processing MCP query for server: {request.mcp_server.value}")
209
210 + license_key = None
211 +
212 + # Check if it's a cloud service and get license key if needed
213 + if MCPService.is_cloud_service(request.mcp_server):
214 + try:
215 + await is_feature_enabled("SOCFORTRESS AI", session=session)
216 + license_info = await get_license(session)
217 + license_key = license_info.license_key
218 + logger.info(f"Retrieved license key for cloud service {request.mcp_server.value}")
219 + except Exception as e:
220 + logger.error(f"Failed to get license key for cloud service: {str(e)}")
221 + return MCPQueryResponse(
222 + message=f"License validation failed: {str(e)}",
223 + success=False,
224 + result=None,
225 + structured_result=None,
226 + execution_time=0.0,
227 + )
228 +
229 # Use the modular service to execute the query
207 - return await MCPService.execute_query(request)
230 + return await MCPService.execute_query(request, license_key=license_key)
backend/app/integrations/copilot_mcp/schema/copilot_mcp.py
+1
@@ -17,6 +17,7 @@ class MCPServerType(str, Enum):
17 WAZUH_MANAGER = "wazuh-manager"
18 COPILOT = "copilot"
19 VELOCIRAPTOR = "velociraptor"
20 + THREAT_INTEL = "threat-intel"
21
22
23 class MCPServerConfig(BaseModel):
backend/app/integrations/copilot_mcp/services/copilot_mcp.py
+117 -20
@@ -1,4 +1,6 @@
1 +from enum import Enum
2 from typing import Dict
3 +from typing import Optional
4
5 import httpx
6 from loguru import logger
@@ -8,38 +10,59 @@ 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
14 - # Define the mapping of MCP server types to their endpoint paths
15 - _SERVER_ENDPOINTS: Dict[MCPServerType, str] = {
16 - MCPServerType.WAZUH_INDEXER: "opensearch-query",
17 - MCPServerType.WAZUH_MANAGER: "wazuh-query",
18 - MCPServerType.COPILOT: "mysql-query",
19 - MCPServerType.VELOCIRAPTOR: "velociraptor-query",
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
22 - # Base URL for the copilot-mcp service
23 - _BASE_URL = "http://copilot-mcp/mcp"
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 + }
47
48 @classmethod
26 - def get_endpoint_for_server(cls, mcp_server: MCPServerType) -> str:
49 + def get_server_config(cls, mcp_server: MCPServerType) -> MCPServerConfig:
50 """
28 - Get the endpoint path for a specific MCP server type.
51 + Get the configuration for a specific MCP server type.
52
53 Args:
54 mcp_server: The MCP server type
55
56 Returns:
34 - str: The endpoint path for the server
57 + MCPServerConfig: The configuration for the server
58
59 Raises:
60 ValueError: If the server type is not supported
61 """
39 - endpoint = cls._SERVER_ENDPOINTS.get(mcp_server)
40 - if not endpoint:
62 + config = cls._SERVER_CONFIGS.get(mcp_server)
63 + if not config:
64 raise ValueError(f"Unsupported MCP server type: {mcp_server}")
42 - return endpoint
65 + return config
66
67 @classmethod
68 def build_full_url(cls, mcp_server: MCPServerType) -> str:
@@ -52,16 +75,32 @@ class MCPService:
75 Returns:
76 str: The full URL for the server endpoint
77 """
55 - endpoint = cls.get_endpoint_for_server(mcp_server)
56 - return f"{cls._BASE_URL}/{endpoint}"
78 + config = cls.get_server_config(mcp_server)
79 + base_url = cls._BASE_URLS[config.service_type]
80 + return f"{base_url}/{config.endpoint}"
81 +
82 + @classmethod
83 + def is_cloud_service(cls, mcp_server: MCPServerType) -> bool:
84 + """
85 + Check if the MCP server is a cloud service.
86 +
87 + Args:
88 + mcp_server: The MCP server type
89 +
90 + Returns:
91 + bool: True if it's a cloud service, False if local
92 + """
93 + config = cls.get_server_config(mcp_server)
94 + return config.service_type == MCPServiceType.CLOUD
95
96 @classmethod
59 - async def execute_query(cls, data: MCPQueryRequest) -> MCPQueryResponse:
97 + async def execute_query(cls, data: MCPQueryRequest, license_key: Optional[str] = None) -> MCPQueryResponse:
98 """
99 Execute a query on the appropriate MCP server based on the request.
100
101 Args:
102 data: The MCP query request containing the server type and query
103 + license_key: Optional license key for cloud services authentication
104
105 Returns:
106 MCPQueryResponse: The response from the MCP server
@@ -73,21 +112,36 @@ class MCPService:
112 try:
113 # Get the full URL for the specified server
114 full_url = cls.build_full_url(data.mcp_server)
115 + is_cloud = cls.is_cloud_service(data.mcp_server)
116
77 - logger.info(f"Sending MCP query to {data.mcp_server.value} at {full_url}")
117 + logger.info(f"Sending MCP query to {data.mcp_server.value} ({'cloud' if is_cloud else 'local'}) at {full_url}")
118 logger.debug(f"Query data: {data.dict()}")
119
120 + # Set different timeout for cloud vs local services
121 + timeout = 180 if is_cloud else 120
122 +
123 + # Prepare headers
124 + headers = {"Content-Type": "application/json"}
125 +
126 + # Add license key as x-api-key header for cloud services
127 + if is_cloud and license_key:
128 + headers["x-api-key"] = license_key
129 + logger.debug("Added x-api-key header for cloud service")
130 + elif is_cloud and not license_key:
131 + logger.warning(f"No license key provided for cloud service {data.mcp_server.value}")
132 +
133 async with httpx.AsyncClient() as client:
134 response = await client.post(
135 full_url,
136 json=data.dict(),
84 - timeout=120,
137 + headers=headers,
138 + timeout=timeout,
139 )
140
141 # Raise an exception for HTTP error status codes
142 response.raise_for_status()
143
90 - logger.info(f"Successfully received response from {data.mcp_server.value}")
144 + logger.info(f"Successfully received response from {data.mcp_server.value}: {response.json()}")
145 return MCPQueryResponse(**response.json())
146
147 except ValueError as e:
@@ -114,6 +168,49 @@ class MCPService:
168 execution_time=0.0,
169 )
170
171 + @classmethod
172 + def add_local_service(cls, server_type: MCPServerType, endpoint: str) -> None:
173 + """
174 + Add a new local MCP service.
175 +
176 + Args:
177 + server_type: The MCP server type enum
178 + endpoint: The endpoint path for the local service
179 + """
180 + cls._SERVER_CONFIGS[server_type] = MCPServerConfig(MCPServiceType.LOCAL, endpoint)
181 +
182 + @classmethod
183 + def add_cloud_service(cls, server_type: MCPServerType, endpoint: str) -> None:
184 + """
185 + Add a new cloud MCP service.
186 +
187 + Args:
188 + server_type: The MCP server type enum
189 + endpoint: The endpoint path for the cloud service
190 + """
191 + cls._SERVER_CONFIGS[server_type] = MCPServerConfig(MCPServiceType.CLOUD, endpoint)
192 +
193 + @classmethod
194 + def get_service_info(cls) -> Dict[str, Dict[str, str]]:
195 + """
196 + Get information about all configured services.
197 +
198 + Returns:
199 + Dict containing service information grouped by type
200 + """
201 + local_services = {}
202 + cloud_services = {}
203 +
204 + for server_type, config in cls._SERVER_CONFIGS.items():
205 + service_info = {"endpoint": config.endpoint, "full_url": cls.build_full_url(server_type)}
206 +
207 + if config.service_type == MCPServiceType.LOCAL:
208 + local_services[server_type.value] = service_info
209 + else:
210 + cloud_services[server_type.value] = service_info
211 +
212 + return {"local": local_services, "cloud": cloud_services}
213 +
214
215 # Convenience function to maintain backward compatibility
216 async def post_to_copilot_mcp(data: MCPQueryRequest) -> MCPQueryResponse:
backend/app/integrations/copilot_mcp/services/example_questions.py
+27
@@ -139,6 +139,23 @@ class ExampleQuestionsService:
139 category="artifacts",
140 ),
141 ],
142 + MCPServerType.THREAT_INTEL: [
143 + ExampleQuestion(
144 + question="What is the IP reputation for 8.8.8.8?",
145 + description="Retrieve the IP reputation for a specific IP address",
146 + category="threat_intel",
147 + ),
148 + ExampleQuestion(
149 + question="What is the domain analysis for example.com?",
150 + description="Get threat intelligence data for a specific domain",
151 + category="threat_intel",
152 + ),
153 + ExampleQuestion(
154 + question="What is the file hash analysis for 1234567890abcdef1234567890abcdef?",
155 + description="Analyze a file hash for malware or other threats",
156 + category="threat_intel",
157 + ),
158 + ],
159 }
160
161 # Define server information with descriptions and capabilities
@@ -191,6 +208,16 @@ class ExampleQuestionsService:
208 "Hardware device tracking",
209 ],
210 ),
211 + MCPServerType.THREAT_INTEL: MCPServerInfo(
212 + name="Threat Intel",
213 + value=MCPServerType.THREAT_INTEL.value,
214 + description="Access threat intelligence data, such as IP reputation scores and malware indicators",
215 + capabilities=[
216 + "Threat intelligence lookups",
217 + "IP and domain reputation scoring",
218 + "File hash and URL analysis",
219 + ],
220 + ),
221 }
222
223 @classmethod
backend/app/middleware/license.py
+4 -1
@@ -413,7 +413,10 @@ async def is_feature_enabled(feature_name: str, session: AsyncSession, message:
413 if message:
414 raise HTTPException(status_code=400, detail=message)
415
416 - raise HTTPException(status_code=400, detail="Feature not enabled. You must purchase a license to use this feature.")
416 + raise HTTPException(
417 + status_code=400,
418 + detail=f"Feature is not enabled. You must purchase the {feature_name} license to use this feature.",
419 + )
420
421
422 async def send_get_request(endpoint: str) -> Dict[str, Any]: