main
py 505 lines 21.1 KB
Raw
1 # App specific imports
2 from typing import List
3 from typing import Optional
4
5 from fastapi import APIRouter
6 from fastapi import HTTPException
7 from fastapi import Path
8 from fastapi import Query
9 from fastapi import Security
10 from loguru import logger
11
12 from app.auth.routes.auth import AuthHandler
13 from app.connectors.wazuh_manager.schema.mitre import AtomicRedTeamMarkdownResponse
14 from app.connectors.wazuh_manager.schema.mitre import AtomicTestsListResponse
15 from app.connectors.wazuh_manager.schema.mitre import MitreTechniqueAlertsResponse
16 from app.connectors.wazuh_manager.schema.mitre import MitreTechniquesInAlertsResponse
17 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreGroupsResponse
18 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreMitigationsResponse
19 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreReferencesResponse
20 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreSoftwareResponse
21 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreTacticsResponse
22 from app.connectors.wazuh_manager.schema.mitre import WazuhMitreTechniquesResponse
23 from app.connectors.wazuh_manager.services.mitre import AtomicRedTeamService
24 from app.connectors.wazuh_manager.services.mitre import get_alerts_by_mitre_id
25 from app.connectors.wazuh_manager.services.mitre import get_mitre_groups
26 from app.connectors.wazuh_manager.services.mitre import get_mitre_mitigations
27 from app.connectors.wazuh_manager.services.mitre import get_mitre_references
28 from app.connectors.wazuh_manager.services.mitre import get_mitre_software
29 from app.connectors.wazuh_manager.services.mitre import get_mitre_tactics
30 from app.connectors.wazuh_manager.services.mitre import get_mitre_techniques
31 from app.connectors.wazuh_manager.services.mitre import (
32 search_mitre_techniques_in_alerts,
33 )
34
35 # Initialize router and auth handler
36 wazuh_manager_mitre_router = APIRouter()
37 auth_handler = AuthHandler()
38
39
40 @wazuh_manager_mitre_router.get(
41 "/groups",
42 response_model=WazuhMitreGroupsResponse,
43 description="List MITRE ATT&CK groups",
44 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
45 )
46 async def list_mitre_groups(
47 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
48 offset: Optional[int] = Query(None, description="First item to return"),
49 select: Optional[List[str]] = Query(None, description="List of fields to return"),
50 sort: Optional[str] = Query(None, description="Fields to sort by"),
51 search: Optional[str] = Query(None, description="Text to search in fields"),
52 q: Optional[str] = Query(None, description="Query to filter results"),
53 ):
54 """
55 List MITRE ATT&CK groups with optional filtering parameters.
56
57 Args:
58 limit: Maximum number of items to return
59 offset: First item to return
60 select: List of fields to return
61 sort: Fields to sort by
62 search: Text to search in fields
63 q: Query to filter results
64
65 Returns:
66 WazuhMitreGroupsResponse: A list of MITRE ATT&CK groups matching the criteria.
67 """
68 return await get_mitre_groups(limit=limit, offset=offset, select=select, sort=sort, search=search, q=q)
69
70
71 @wazuh_manager_mitre_router.get(
72 "/mitigations",
73 response_model=WazuhMitreMitigationsResponse,
74 description="List MITRE ATT&CK mitigations",
75 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
76 )
77 async def list_mitre_mitigations(
78 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
79 offset: Optional[int] = Query(None, description="First item to return"),
80 select: Optional[List[str]] = Query(None, description="List of fields to return"),
81 sort: Optional[str] = Query(None, description="Fields to sort by"),
82 search: Optional[str] = Query(None, description="Text to search in fields"),
83 q: Optional[str] = Query(None, description="Query to filter results"),
84 ):
85 """
86 List MITRE ATT&CK mitigations with optional filtering parameters.
87
88 Args:
89 limit: Maximum number of items to return
90 offset: First item to return
91 select: List of fields to return
92 sort: Fields to sort by
93 search: Text to search in fields
94 q: Query to filter results
95
96 Returns:
97 WazuhMitreMitigationsResponse: A list of MITRE ATT&CK mitigations matching the criteria.
98 """
99 return await get_mitre_mitigations(limit=limit, offset=offset, select=select, sort=sort, search=search, q=q)
100
101
102 @wazuh_manager_mitre_router.get(
103 "/references",
104 response_model=WazuhMitreReferencesResponse,
105 description="List MITRE ATT&CK references",
106 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
107 )
108 async def list_mitre_references(
109 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
110 offset: Optional[int] = Query(None, description="First item to return"),
111 sort: Optional[str] = Query(None, description="Fields to sort by"),
112 search: Optional[str] = Query(None, description="Text to search in fields"),
113 q: Optional[str] = Query(None, description="Query to filter results"),
114 ):
115 """
116 List MITRE ATT&CK references with optional filtering parameters.
117
118 Args:
119 limit: Maximum number of items to return
120 offset: First item to return
121 sort: Fields to sort by
122 search: Text to search in fields
123 q: Query to filter results
124
125 Returns:
126 WazuhMitreReferencesResponse: A list of MITRE ATT&CK references matching the criteria.
127 """
128 return await get_mitre_references(limit=limit, offset=offset, sort=sort, search=search, q=q)
129
130
131 @wazuh_manager_mitre_router.get(
132 "/software",
133 response_model=WazuhMitreSoftwareResponse,
134 description="List MITRE ATT&CK software",
135 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
136 )
137 async def list_mitre_software(
138 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
139 offset: Optional[int] = Query(None, description="First item to return"),
140 select: Optional[List[str]] = Query(None, description="List of fields to return"),
141 sort: Optional[str] = Query(None, description="Fields to sort by"),
142 search: Optional[str] = Query(None, description="Text to search in fields"),
143 q: Optional[str] = Query(None, description="Query to filter results"),
144 ):
145 """
146 List MITRE ATT&CK software with optional filtering parameters.
147
148 Args:
149 limit: Maximum number of items to return
150 offset: First item to return
151 select: List of fields to return
152 sort: Fields to sort by
153 search: Text to search in fields
154 q: Query to filter results
155
156 Returns:
157 WazuhMitreSoftwareResponse: A list of MITRE ATT&CK software matching the criteria.
158 """
159 return await get_mitre_software(limit=limit, offset=offset, select=select, sort=sort, search=search, q=q)
160
161
162 @wazuh_manager_mitre_router.get(
163 "/tactics",
164 response_model=WazuhMitreTacticsResponse,
165 description="List MITRE ATT&CK tactics",
166 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
167 )
168 async def list_mitre_tactics(
169 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
170 offset: Optional[int] = Query(None, description="First item to return"),
171 select: Optional[List[str]] = Query(None, description="List of fields to return"),
172 sort: Optional[str] = Query(None, description="Fields to sort by"),
173 search: Optional[str] = Query(None, description="Text to search in fields"),
174 q: Optional[str] = Query(None, description="Query to filter results"),
175 ):
176 """
177 List MITRE ATT&CK tactics with optional filtering parameters.
178
179 Args:
180 limit: Maximum number of items to return
181 offset: First item to return
182 select: List of fields to return
183 sort: Fields to sort by
184 search: Text to search in fields
185 q: Query to filter results
186
187 Returns:
188 WazuhMitreTacticsResponse: A list of MITRE ATT&CK tactics matching the criteria.
189 """
190 return await get_mitre_tactics(limit=limit, offset=offset, select=select, sort=sort, search=search, q=q)
191
192
193 @wazuh_manager_mitre_router.get(
194 "/techniques",
195 response_model=WazuhMitreTechniquesResponse,
196 description="List MITRE ATT&CK techniques",
197 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
198 )
199 async def list_mitre_techniques(
200 limit: Optional[int] = Query(None, description="Maximum number of items to return"),
201 offset: Optional[int] = Query(None, description="First item to return"),
202 select: Optional[List[str]] = Query(None, description="List of fields to return"),
203 sort: Optional[str] = Query(None, description="Fields to sort by"),
204 search: Optional[str] = Query(None, description="Text to search in fields"),
205 q: Optional[str] = Query(None, description="Query to filter results"),
206 ):
207 """
208 List MITRE ATT&CK techniques with optional filtering parameters.
209
210 Args:
211 limit: Maximum number of items to return
212 offset: First item to return
213 select: List of fields to return
214 sort: Fields to sort by
215 search: Text to search in fields
216 q: Query to filter results
217
218 Returns:
219 WazuhMitreTechniquesResponse: A list of MITRE ATT&CK techniques matching the criteria.
220 """
221 return await get_mitre_techniques(limit=limit, offset=offset, select=select, sort=sort, search=search, q=q)
222
223
224 # @wazuh_manager_mitre_router.get(
225 # "/atomic-tests",
226 # response_model=AtomicTestsListResponse,
227 # description="List all available Atomic Red Team tests",
228 # dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
229 # )
230 # async def list_atomic_tests(
231 # size: int = Query(25, description="Maximum number of techniques to return per page"),
232 # page: int = Query(1, description="Page number for pagination", gt=0),
233 # ):
234 # """
235 # List all available Atomic Red Team tests across all techniques.
236
237 # Args:
238 # size: Maximum number of techniques to return per page
239 # page: Page number for pagination
240
241 # Returns:
242 # AtomicTestsListResponse: A paginated list of techniques with Atomic Red Team tests.
243 # """
244 # logger.info(f"Request for list of all Atomic Red Team tests (page {page}, size {size})")
245
246 # try:
247 # # Get the list of all atomic tests
248 # result = await AtomicRedTeamService.list_all_atomic_tests()
249
250 # # Apply pagination to the results
251 # total_techniques = result["total_techniques"]
252 # all_tests = result["tests"]
253
254 # # Calculate total pages
255 # total_pages = (total_techniques + size - 1) // size if total_techniques > 0 else 1
256
257 # # Apply pagination
258 # start_idx = (page - 1) * size
259 # end_idx = start_idx + size
260 # paginated_tests = all_tests[start_idx:end_idx]
261
262 # return AtomicTestsListResponse(
263 # success=True,
264 # message=f"Found {total_techniques} MITRE techniques in {result['total_techniques']} alerts (page {page} of {total_pages},)",
265 # total_techniques=total_techniques,
266 # total_tests=result.get("total_tests"),
267 # tests=paginated_tests,
268 # last_updated=result["last_updated"],
269 # page=page,
270 # page_size=size,
271 # total_pages=total_pages,
272 # )
273 # except Exception as e:
274 # logger.error(f"Error retrieving atomic tests: {str(e)}")
275 # raise HTTPException(status_code=500, detail=f"Error retrieving atomic tests: {str(e)}")
276
277
278 @wazuh_manager_mitre_router.get(
279 "/atomic-tests",
280 response_model=AtomicTestsListResponse,
281 description="List all available Atomic Red Team tests",
282 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
283 )
284 async def list_atomic_tests(
285 size: int = Query(25, description="Maximum number of techniques to return per page"),
286 page: int = Query(1, description="Page number for pagination", gt=0),
287 os_category: Optional[str] = Query(None, description="Filter by operating system category (windows, linux, macos)"),
288 ):
289 """
290 List all available Atomic Red Team tests across all techniques.
291
292 Args:
293 size: Maximum number of techniques to return per page
294 page: Page number for pagination
295 os_category: Optional filter for operating system category
296
297 Returns:
298 AtomicTestsListResponse: A paginated list of techniques with Atomic Red Team tests.
299 """
300 logger.info(f"Request for list of all Atomic Red Team tests (page {page}, size {size}, os_category: {os_category})")
301
302 try:
303 # Get the list of all atomic tests
304 result = await AtomicRedTeamService.list_all_atomic_tests()
305
306 # Filter by OS category if provided
307 all_tests = result["tests"]
308 if os_category:
309 os_category_lower = os_category.lower()
310 # Filter tests that have the specified OS category in their categories list
311 filtered_tests = [test for test in all_tests if os_category_lower in [cat.lower() for cat in test.get("categories", [])]]
312 logger.info(f"Filtered {len(all_tests)} tests down to {len(filtered_tests)} tests for OS category '{os_category}'")
313 else:
314 filtered_tests = all_tests
315
316 # Apply pagination to the filtered results
317 total_techniques = len(filtered_tests)
318
319 # Calculate total pages
320 total_pages = (total_techniques + size - 1) // size if total_techniques > 0 else 1
321
322 # Apply pagination
323 start_idx = (page - 1) * size
324 end_idx = start_idx + size
325 paginated_tests = filtered_tests[start_idx:end_idx]
326
327 # Build the message
328 if os_category:
329 message = f"Found {total_techniques} MITRE techniques for OS '{os_category}' (page {page} of {total_pages})"
330 else:
331 message = f"Found {total_techniques} MITRE techniques (page {page} of {total_pages})"
332
333 return AtomicTestsListResponse(
334 success=True,
335 message=message,
336 total_techniques=total_techniques,
337 total_tests=result.get("total_tests"),
338 tests=paginated_tests,
339 last_updated=result["last_updated"],
340 page=page,
341 page_size=size,
342 total_pages=total_pages,
343 )
344 except Exception as e:
345 logger.error(f"Error retrieving atomic tests: {str(e)}")
346 raise HTTPException(status_code=500, detail=f"Error retrieving atomic tests: {str(e)}")
347
348
349 @wazuh_manager_mitre_router.get(
350 "/techniques/{technique_id}/atomic-tests",
351 response_model=AtomicRedTeamMarkdownResponse,
352 description="Get Atomic Red Team tests for a MITRE ATT&CK technique",
353 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
354 )
355 async def get_technique_atomic_tests(technique_id: str = Path(..., description="MITRE ATT&CK technique ID (e.g., T1003, T1003.004)")):
356 """
357 Get Atomic Red Team tests for a specific MITRE ATT&CK technique.
358
359 Args:
360 technique_id: The MITRE ATT&CK technique ID
361
362 Returns:
363 AtomicRedTeamMarkdownResponse: The Atomic Red Team tests for the technique
364 """
365 logger.info(f"Request for Atomic Red Team tests for technique {technique_id}")
366
367 # Extract the technique ID from the full ID if needed (e.g., "T1003.004" -> "T1003.004")
368 clean_technique_id = technique_id.split("-")[-1] if "-" in technique_id else technique_id
369
370 # Fetch the markdown content
371 markdown_content = await AtomicRedTeamService.get_technique_markdown(clean_technique_id)
372
373 if markdown_content is None:
374 return AtomicRedTeamMarkdownResponse(
375 success=False,
376 message=f"No Atomic Red Team tests found for technique {technique_id}",
377 technique_id=clean_technique_id,
378 )
379
380 return AtomicRedTeamMarkdownResponse(
381 success=True,
382 message=f"Atomic Red Team tests retrieved for technique {technique_id}",
383 technique_id=clean_technique_id,
384 markdown_content=markdown_content,
385 )
386
387
388 @wazuh_manager_mitre_router.get(
389 "/techniques/alerts",
390 response_model=MitreTechniquesInAlertsResponse,
391 description="Search for MITRE ATT&CK techniques in alerts",
392 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
393 )
394 async def list_mitre_techniques_in_alerts(
395 time_range: str = Query("now-24h", description="Time range for the search (e.g., now-24h, now-7d)"),
396 size: int = Query(25, description="Maximum number of techniques to return per page"),
397 page: int = Query(1, description="Page number for pagination", gt=0),
398 rule_level: Optional[int] = Query(None, description="Filter by rule level"),
399 rule_group: Optional[str] = Query(None, description="Filter by rule group"),
400 mitre_field: Optional[str] = Query(None, description="Override the field containing MITRE IDs"),
401 index_pattern: str = Query("wazuh-*", description="Index pattern to search"),
402 ) -> MitreTechniquesInAlertsResponse:
403 """Search for MITRE ATT&CK techniques in Wazuh alerts."""
404 logger.info(f"Searching for MITRE techniques in alerts from {time_range} (page {page}, size {size})")
405
406 # Calculate the offset based on page and size
407 offset = (page - 1) * size
408
409 # Build additional filters based on request parameters
410 additional_filters = []
411
412 if rule_level is not None:
413 additional_filters.append({"match_phrase": {"rule_level": {"query": str(rule_level)}}})
414
415 if rule_group is not None:
416 additional_filters.append({"match_phrase": {"rule_groups": {"query": rule_group}}})
417
418 # Execute the search with the specified parameters
419 results = await search_mitre_techniques_in_alerts(
420 time_range=time_range,
421 size=size,
422 offset=offset,
423 additional_filters=additional_filters,
424 index_pattern=index_pattern,
425 mitre_field=mitre_field,
426 )
427
428 # Get the total number of techniques (from all pages)
429 total_techniques = results.get("total_techniques_count", results["techniques_count"])
430
431 # Calculate total pages based on the total number of techniques
432 total_pages = (total_techniques + size - 1) // size if total_techniques > 0 else 1
433
434 return MitreTechniquesInAlertsResponse(
435 success=True,
436 message=f"Found {total_techniques} MITRE techniques in {results['total_alerts']} alerts (page {page} of {total_pages},)",
437 total_alerts=results["total_alerts"],
438 techniques_count=total_techniques, # Use the total count for all pages
439 techniques=results["techniques"], # Use current page techniques
440 time_range=time_range,
441 field_used=results.get("field_used", "unknown"),
442 page=page,
443 page_size=size,
444 total_pages=total_pages,
445 )
446
447
448 @wazuh_manager_mitre_router.get(
449 "/techniques/{technique_id}/alerts",
450 response_model=MitreTechniqueAlertsResponse,
451 description="Get alert documents for a specific MITRE ATT&CK technique",
452 dependencies=[Security(auth_handler.require_any_scope("admin", "analyst"))],
453 )
454 async def get_mitre_technique_alerts(
455 technique_id: str = Path(..., description="MITRE ATT&CK technique ID (e.g., T1047, 1047)"),
456 time_range: str = Query("now-24h", description="Time range for the search (e.g., now-24h, now-7d)"),
457 size: int = Query(25, description="Maximum number of alerts to return per page"),
458 page: int = Query(1, description="Page number for pagination", gt=0),
459 rule_level: Optional[int] = Query(None, description="Filter by rule level"),
460 rule_group: Optional[str] = Query(None, description="Filter by rule group"),
461 mitre_field: Optional[str] = Query(None, description="Override the field containing MITRE IDs"),
462 index_pattern: str = Query("wazuh-*", description="Index pattern to search"),
463 ) -> MitreTechniqueAlertsResponse:
464 """Get alert documents for a specific MITRE ATT&CK technique."""
465 logger.info(f"Request for alerts related to MITRE technique {technique_id} from {time_range} (page {page}, size {size})")
466
467 # Clean up technique ID if needed
468 clean_technique_id = technique_id.strip()
469
470 # Calculate the offset based on page and size
471 offset = (page - 1) * size
472
473 # Build additional filters based on request parameters
474 additional_filters = []
475
476 if rule_level is not None:
477 additional_filters.append({"match_phrase": {"rule_level": {"query": str(rule_level)}}})
478
479 if rule_group is not None:
480 additional_filters.append({"match_phrase": {"rule_groups": {"query": rule_group}}})
481
482 # Get the alerts
483 results = await get_alerts_by_mitre_id(
484 technique_id=clean_technique_id,
485 time_range=time_range,
486 size=size,
487 offset=offset,
488 additional_filters=additional_filters,
489 index_pattern=index_pattern,
490 mitre_field=mitre_field,
491 )
492
493 return MitreTechniqueAlertsResponse(
494 success=True,
495 message=f"Found {results['total_alerts']} alerts for MITRE technique {clean_technique_id} (page {page} of {(results['total_alerts'] + size - 1) // size},)",
496 technique_id=results["technique_id"],
497 technique_name=results["technique_name"],
498 total_alerts=results["total_alerts"],
499 alerts=results["alerts"],
500 field_used=results.get("field_used", "unknown"),
501 time_range=time_range,
502 page=page,
503 page_size=size,
504 total_pages=(results["total_alerts"] + size - 1) // size,
505 )