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]: