main
py 607 lines 22 KB
Raw
1 import asyncio
2 from collections import defaultdict
3 from datetime import datetime
4 from datetime import timedelta
5 from typing import Dict
6 from typing import List
7 from typing import Optional
8 from typing import Tuple
9 from typing import Type
10
11 from elasticsearch7 import AsyncElasticsearch
12 from elasticsearch7.exceptions import NotFoundError
13 from elasticsearch7.exceptions import RequestError
14 from fastapi import HTTPException
15 from loguru import logger
16
17 # from app.connectors.wazuh_indexer.schema.alerts import Alert
18 from app.connectors.wazuh_indexer.schema.alerts import AlertNotFound
19 from app.connectors.wazuh_indexer.schema.alerts import AlertsByHost
20 from app.connectors.wazuh_indexer.schema.alerts import AlertsByHostResponse
21 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRule
22 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHost
23 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRulePerHostResponse
24 from app.connectors.wazuh_indexer.schema.alerts import AlertsByRuleResponse
25 from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchBody
26 from app.connectors.wazuh_indexer.schema.alerts import AlertsSearchResponse
27 from app.connectors.wazuh_indexer.schema.alerts import CollectAlertsResponse
28 from app.connectors.wazuh_indexer.schema.alerts import GraylogAlertsSearchBody
29 from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchBody
30 from app.connectors.wazuh_indexer.schema.alerts import HostAlertsSearchResponse
31 from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchBody
32 from app.connectors.wazuh_indexer.schema.alerts import IndexAlertsSearchResponse
33 from app.connectors.wazuh_indexer.schema.alerts import SkippableWazuhIndexerClientErrors
34 from app.connectors.wazuh_indexer.utils.universal import AlertsQueryBuilder
35 from app.connectors.wazuh_indexer.utils.universal import collect_indices
36 from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
37 from app.connectors.wazuh_indexer.utils.universal import (
38 create_wazuh_indexer_client_async,
39 )
40
41
42 async def collect_and_aggregate_alerts(
43 field_names: List[str],
44 search_body: AlertsSearchBody,
45 ) -> Dict[str, int]:
46 """
47 Collects and aggregates alerts based on the specified field names and search body.
48
49 Args:
50 field_names (List[str]): The list of field names to use for aggregation.
51 search_body (AlertsSearchBody): The search body to filter alerts.
52
53 Returns:
54 Dict[str, int]: A dictionary containing the aggregated alerts, where the keys are composite keys
55 based on the specified field names, and the values are the count of alerts for each composite key.
56 """
57 indices = await collect_indices()
58 aggregated_alerts_dict = {}
59
60 for index_name in indices.indices_list:
61 try:
62 alerts_response = await collect_alerts_generic(index_name, body=search_body)
63 if alerts_response.success:
64 for alert in alerts_response.alerts:
65 composite_key = tuple(alert["_source"][field] for field in field_names)
66 aggregated_alerts_dict[composite_key] = aggregated_alerts_dict.get(composite_key, 0) + 1
67 except HTTPException as e:
68 detail_str = str(e.detail) # Convert to string to make sure it's comparable
69 if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors):
70 logger.warning(
71 f"Skipping index {index_name} due to specific error: {e.detail}",
72 )
73 continue # Skip this index and continue with the next one
74 else:
75 logger.warning(
76 f"An error occurred while processing index {index_name}: {e.detail}",
77 )
78 raise HTTPException(
79 status_code=500,
80 detail=f"An error occurred while processing index {index_name}: {e.detail}",
81 )
82
83 return aggregated_alerts_dict
84
85
86 async def collect_alerts_generic(
87 index_name: str,
88 body: AlertsSearchBody,
89 is_host_specific: bool = False,
90 ) -> CollectAlertsResponse:
91 """
92 Collects alerts from the specified index based on the provided search criteria.
93
94 Args:
95 index_name (str): The name of the index to search for alerts.
96 body (AlertsSearchBody): The search criteria for filtering alerts.
97 is_host_specific (bool, optional): Flag indicating whether the search should be limited to a specific host.
98 Defaults to False.
99
100 Returns:
101 CollectAlertsResponse: The response containing the collected alerts.
102
103 Raises:
104 HTTPException: If an error occurs while collecting alerts.
105
106 """
107 es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
108 query_builder = AlertsQueryBuilder()
109
110 try:
111 query_builder.add_time_range(
112 timerange=body.timerange,
113 timestamp_field=body.timestamp_field,
114 )
115 query_builder.add_matches(matches=[(body.alert_field, body.alert_value)])
116 query_builder.add_sort(body.timestamp_field)
117
118 if is_host_specific:
119 query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)])
120
121 query = query_builder.build()
122
123 alerts = es_client.search(index=index_name, body=query, size=body.size)
124 except RequestError as e:
125 logger.warning(f"An error occurred while collecting alerts: {e}")
126 if "No mapping found for [timestamp_utc] in order to sort on" in str(e):
127 logger.warning("Retrying with timestamp field set to 'timestamp'")
128 body.timestamp_field = "timestamp"
129 try:
130 return await collect_alerts_generic(index_name, body, is_host_specific)
131 except RequestError as e:
132 if "No mapping found for [timestamp] in order to sort on" in str(e):
133 logger.warning("Retrying with timestamp field set to '@timestamp'")
134 body.timestamp_field = "@timestamp"
135 return await collect_alerts_generic(index_name, body, is_host_specific)
136 else:
137 logger.warning(f"An error occurred while collecting alerts: {e}")
138 raise HTTPException(
139 status_code=500,
140 detail=f"An error occurred while collecting alerts: {e}",
141 )
142
143 logger.info(f"Alerts collected: {alerts}")
144 alerts_list = [alert for alert in alerts["hits"]["hits"]]
145 logger.info(f"Alerts collected: {alerts_list}")
146 return CollectAlertsResponse(
147 alerts=alerts_list,
148 success=True,
149 message="Alerts collected successfully",
150 )
151
152
153 async def get_alerts_generic(
154 search_body: Type[AlertsSearchBody],
155 is_host_specific: bool = False,
156 index_name: Optional[str] = None,
157 ):
158 """
159 Retrieves alerts from the Wazuh Indexer based on the provided search criteria.
160
161 Args:
162 search_body (Type[AlertsSearchBody]): The search criteria for the alerts.
163 is_host_specific (bool, optional): Specifies whether the search is host-specific. Defaults to False.
164 index_name (str, optional): The name of the index to search in. If not provided, all indices will be searched.
165
166 Returns:
167 dict: A dictionary containing the alerts summary, success status, and a message.
168 """
169 logger.info(
170 f"Collecting Wazuh Indexer alerts for host {search_body.agent_name if is_host_specific else ''}",
171 )
172 alerts_summary = []
173 indices = await collect_indices()
174 index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices
175
176 for index_name in index_list:
177 try:
178 alerts = await collect_alerts_generic(
179 index_name,
180 body=search_body,
181 is_host_specific=is_host_specific,
182 )
183 if alerts.success and len(alerts.alerts) > 0:
184 alerts_summary.append(
185 {
186 "index_name": index_name,
187 "total_alerts": len(alerts.alerts),
188 "alerts": alerts.alerts,
189 },
190 )
191 except HTTPException as e:
192 detail_str = str(e.detail) # Convert to string to make sure it's comparable
193 if any(err.value in detail_str for err in SkippableWazuhIndexerClientErrors):
194 logger.warning(
195 f"Skipping index {index_name} due to specific error: {e.detail}",
196 )
197 continue # Skip this index and continue with the next one
198 else:
199 logger.warning(
200 f"An error occurred while processing index {index_name}: {e.detail}",
201 )
202 raise HTTPException(
203 status_code=500,
204 detail=f"An error occurred while processing index {index_name}: {e.detail}",
205 )
206
207 if len(alerts_summary) == 0:
208 message = "No alerts found"
209 else:
210 message = f"Succesfully collected top {search_body.size} alerts for each index"
211
212 return {
213 "alerts_summary": alerts_summary,
214 "success": len(alerts_summary) > 0,
215 "message": message,
216 }
217
218
219 async def get_alerts(search_body: AlertsSearchBody) -> AlertsSearchResponse:
220 """
221 Retrieves alerts based on the provided search criteria.
222
223 Args:
224 search_body (AlertsSearchBody): The search criteria for retrieving alerts.
225
226 Returns:
227 AlertsSearchResponse: The response containing the retrieved alerts.
228 """
229 result = await get_alerts_generic(search_body)
230 return AlertsSearchResponse(**result)
231
232
233 async def get_host_alerts(
234 search_body: HostAlertsSearchBody,
235 ) -> HostAlertsSearchResponse:
236 """
237 Retrieves alerts specific to a host.
238
239 Args:
240 search_body (HostAlertsSearchBody): The search criteria for retrieving host alerts.
241
242 Returns:
243 HostAlertsSearchResponse: The response containing the host alerts.
244 """
245 result = await get_alerts_generic(search_body, is_host_specific=True)
246 return HostAlertsSearchResponse(**result)
247
248
249 async def get_index_alerts(
250 search_body: IndexAlertsSearchBody,
251 ) -> IndexAlertsSearchResponse:
252 """
253 Retrieves alerts from the specified index based on the search criteria.
254
255 Args:
256 search_body (IndexAlertsSearchBody): The search criteria for retrieving alerts.
257
258 Returns:
259 IndexAlertsSearchResponse: The response containing the retrieved alerts.
260 """
261 result = await get_alerts_generic(search_body, index_name=search_body.index_name)
262 return IndexAlertsSearchResponse(**result)
263
264
265 async def get_alerts_by_host(search_body: AlertsSearchBody) -> AlertsByHostResponse:
266 """
267 Retrieves alerts grouped by host.
268
269 Args:
270 search_body (AlertsSearchBody): The search criteria for retrieving alerts.
271
272 Returns:
273 AlertsByHostResponse: The response containing alerts grouped by host.
274
275 """
276 aggregated_by_host = await collect_and_aggregate_alerts(["agent_name"], search_body)
277 alerts_by_host_list: List[AlertsByHost] = [
278 AlertsByHost(
279 agent_name=host[0],
280 number_of_alerts=count,
281 ) # host[0] because host is now a tuple
282 for host, count in aggregated_by_host.items()
283 ]
284 return AlertsByHostResponse(
285 alerts_by_host=alerts_by_host_list,
286 success=bool(alerts_by_host_list),
287 message="Successfully collected alerts by host",
288 )
289
290
291 async def get_alerts_by_rule(search_body: AlertsSearchBody) -> AlertsByRuleResponse:
292 """
293 Retrieves alerts grouped by rule based on the provided search criteria.
294
295 Args:
296 search_body (AlertsSearchBody): The search criteria for retrieving alerts.
297
298 Returns:
299 AlertsByRuleResponse: The response containing the alerts grouped by rule.
300
301 """
302 aggregated_by_rule = await collect_and_aggregate_alerts(
303 ["rule_description"],
304 search_body,
305 )
306 alerts_by_rule_list: List[AlertsByRule] = [
307 AlertsByRule(
308 rule=rule[0],
309 number_of_alerts=count,
310 ) # rule[0] because rule is now a tuple
311 for rule, count in aggregated_by_rule.items()
312 ]
313 return AlertsByRuleResponse(
314 alerts_by_rule=alerts_by_rule_list,
315 success=bool(alerts_by_rule_list),
316 message="Successfully collected alerts by rule",
317 )
318
319
320 async def get_alerts_by_rule_per_host(
321 search_body: AlertsSearchBody,
322 ) -> AlertsByRulePerHostResponse:
323 """
324 Retrieves alerts grouped by rule per host based on the provided search criteria.
325
326 Args:
327 search_body (AlertsSearchBody): The search criteria for retrieving alerts.
328
329 Returns:
330 AlertsByRulePerHostResponse: The response containing the alerts grouped by rule per host.
331
332 """
333 aggregated_by_rule_per_host = await collect_and_aggregate_alerts(
334 ["agent_name", "rule_description"],
335 search_body,
336 )
337 alerts_by_rule_per_host_list: List[AlertsByRulePerHost] = [
338 AlertsByRulePerHost(agent_name=agent_name, rule=rule, number_of_alerts=count)
339 for (agent_name, rule), count in aggregated_by_rule_per_host.items()
340 ]
341
342 return AlertsByRulePerHostResponse(
343 alerts_by_rule_per_host=alerts_by_rule_per_host_list,
344 success=bool(alerts_by_rule_per_host_list),
345 message="Successfully collected alerts by rule per host",
346 )
347
348
349 def parse_timerange(timerange: str) -> str:
350 """
351 Parses the timerange string and returns the corresponding datetime string for Elasticsearch.
352 """
353 unit = timerange[-1]
354 amount = int(timerange[:-1])
355
356 if unit == "h":
357 delta = timedelta(hours=amount)
358 elif unit == "d":
359 delta = timedelta(days=amount)
360 elif unit == "w":
361 delta = timedelta(weeks=amount)
362 else:
363 raise HTTPException(
364 status_code=400,
365 detail="Invalid timerange unit. Must be one of 'h', 'd', or 'w'.",
366 )
367
368 start_time = datetime.utcnow() - delta
369 return start_time.isoformat() + "Z"
370
371
372 async def get_original_alert_id(origin_context: str) -> Tuple[str, str]:
373 """
374 Extracts the index name and id from the origin_context field of an alert.
375
376 Args:
377 origin_context (str): The origin_context field of an alert.
378
379 Returns:
380 Tuple[str, str]: A tuple containing the index name and id of the alert.
381 """
382 index_name, index_id = origin_context.split(":")[-2:]
383 return index_name, index_id
384
385
386 # ! THIS IS OLD WITH ASYNC OPERATIONS ! #
387 # async def get_single_alert_details(
388 # index_name: str,
389 # index_id: str,
390 # ) -> Dict:
391 # """
392 # Retrieves the details of a single alert based on the index name and id.
393
394 # Args:
395 # es_client: The Elasticsearch client.
396 # index_name (str): The name of the index to search for the alert.
397 # index_id (str): The id of the alert to retrieve.
398
399 # Returns:
400 # dict: The details of the alert.
401 # """
402 # es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
403 # try:
404 # alert = es_client.get(index=index_name, id=index_id)
405 # return alert
406 # except NotFoundError:
407 # logger.warning(f"Alert not found for index {index_name} and id {index_id}")
408 # return AlertNotFound(_index=index_name, _id=index_id, _source={"message": "alert not found"}).to_dict()
409
410
411 # async def fetch_alerts_from_graylog(index_prefix: str, size: int, timerange: str) -> List[Dict]:
412 # """
413 # Fetches alerts from the Graylog Alert Index.
414
415 # Args:
416 # es_client: The Elasticsearch client.
417 # index_prefix (str): The prefix of the index to search for alerts.
418 # size (int): The number of alerts to retrieve.
419
420 # Returns:
421 # List[Dict]: A list of alerts.
422 # """
423 # es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
424 # es_query = {
425 # "size": size,
426 # "query": {
427 # "bool": {
428 # "must": [
429 # {
430 # "range": {
431 # "timestamp": {
432 # "gte": parse_timerange(timerange),
433 # "format": "strict_date_optional_time",
434 # },
435 # },
436 # },
437 # ],
438 # },
439 # },
440 # }
441 # response = await es_client.search(index=index_prefix, body=es_query)
442 # logger.info(f"Graylog alerts response: {response}")
443 # return response["hits"]["hits"]
444
445
446 # async def process_alert_hits(hits: List[Dict]) -> List[Dict]:
447 # """
448 # Processes the hits from the Graylog alert search response.
449
450 # Args:
451 # es_client: The Elasticsearch client.
452 # hits (List[Dict]): The hits from the search response.
453
454 # Returns:
455 # List[Dict]: A list of detailed alerts.
456 # """
457 # alerts_dict = defaultdict(lambda: {"total_alerts": 0, "alerts": []})
458
459 # for hit in hits:
460 # origin_context = hit["_source"]["origin_context"]
461 # index_name, index_id = await get_original_alert_id(origin_context)
462 # logger.info(f"Fetching alert details for index {index_name} and id {index_id}")
463 # alert_details = await get_single_alert_details(index_name, index_id)
464 # logger.info(f"Alert details: {alert_details}")
465
466 # alerts_dict[index_name]["total_alerts"] += 1
467 # alerts_dict[index_name]["alerts"].append(alert_details)
468
469 # alerts = [
470 # Alert(index_name=index_name, total_alerts=data["total_alerts"], alerts=data["alerts"]) for index_name, data in alerts_dict.items()
471 # ]
472
473 # return alerts
474
475
476 # async def get_graylog_alerts(
477 # request: GraylogAlertsSearchBody,
478 # ) -> AlertsSearchResponse:
479 # """
480 # Retrieves alerts from the Graylog Alert Index.
481 # Looks up the actual alert details via the origin_context field.
482 # Strips out the index name and id from the origin_context field.
483 # Example: urn:graylog:message:es:huntress_00002_0:b4d2c721-f690-11ee-ac73-8600007a2218
484 # Looks up each alert by the index name and id.
485 # Adds each alert to the response.
486 # """
487 # logger.info(f"Fetching Graylog alerts for request: {request}")
488
489 # hits = await fetch_alerts_from_graylog(request.index_prefix, request.size, request.timerange)
490 # alerts = await process_alert_hits(hits)
491
492 # return alerts
493
494 # ! ^^THIS IS OLD WITH ASYNC OPERATIONS^^ ! #
495
496
497 async def get_single_alert_details(
498 es_client: AsyncElasticsearch,
499 index_name: str,
500 index_id: str,
501 ) -> Dict:
502 """
503 Retrieves the details of a single alert based on the index name and id.
504
505 Args:
506 es_client: The Elasticsearch client.
507 index_name (str): The name of the index to search for the alert.
508 index_id (str): The id of the alert to retrieve.
509
510 Returns:
511 dict: The details of the alert.
512 """
513 try:
514 alert = await es_client.get(index=index_name, id=index_id)
515 return alert
516 except NotFoundError:
517 logger.warning(f"Alert not found for index {index_name} and id {index_id}")
518 return AlertNotFound(_index=index_name, _id=index_id, _source={"message": "alert not found"}).to_dict()
519
520
521 async def fetch_alerts_from_graylog(index_prefix: str, size: int, timerange: str) -> List[Dict]:
522 """
523 Fetches alerts from the Graylog Alert Index.
524
525 Args:
526 es_client: The Elasticsearch client.
527 index_prefix (str): The prefix of the index to search for alerts.
528 size (int): The number of alerts to retrieve.
529
530 Returns:
531 List[Dict]: A list of alerts.
532 """
533 es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
534 es_query = {
535 "size": size,
536 "query": {
537 "bool": {
538 "must": [
539 {
540 "range": {
541 "timestamp": {
542 "gte": parse_timerange(timerange),
543 "format": "strict_date_optional_time",
544 },
545 },
546 },
547 ],
548 },
549 },
550 }
551 response = await es_client.search(index=index_prefix, body=es_query)
552 logger.info(f"Graylog alerts response: {response}")
553 return response["hits"]["hits"]
554
555
556 async def process_alert_hits(hits: List[Dict], es_client: AsyncElasticsearch) -> List[Dict]:
557 alerts_dict = defaultdict(lambda: {"total_alerts": 0, "alerts": []})
558
559 tasks = []
560 for hit in hits:
561 origin_context = hit.get("_source", {}).get("origin_context")
562 if not origin_context:
563 logger.warning(f"Skipping alert hit with missing or null origin_context: {hit.get('_id', 'unknown')}")
564 continue
565 try:
566 index_name, index_id = await get_original_alert_id(origin_context)
567 except Exception as e:
568 logger.warning(f"Skipping alert hit due to error parsing origin_context '{origin_context}': {e}")
569 continue
570 logger.info(f"Fetching alert details for index {index_name} and id {index_id}")
571 task = get_single_alert_details(es_client, index_name, index_id)
572 tasks.append(task)
573
574 alert_details_list = await asyncio.gather(*tasks)
575
576 for alert_details in alert_details_list:
577 index_name = alert_details["_index"]
578 alerts_dict[index_name]["total_alerts"] += 1
579 alerts_dict[index_name]["alerts"].append(alert_details)
580
581 alerts = [
582 {"index_name": index_name, "total_alerts": data["total_alerts"], "alerts": data["alerts"]}
583 for index_name, data in alerts_dict.items()
584 ]
585
586 logger.info(f"Processed alerts: {alerts}")
587
588 return alerts
589
590
591 async def get_graylog_alerts(
592 request: GraylogAlertsSearchBody,
593 ) -> AlertsSearchResponse:
594 """
595 Retrieves alerts from the Graylog Alert Index.
596 Looks up the actual alert details via the origin_context field.
597 Strips out the index name and id from the origin_context field.
598 Example: urn:graylog:message:es:huntress_00002_0:b4d2c721-f690-11ee-ac73-8600007a2218
599 Looks up each alert by the index name and id.
600 Adds each alert to the response.
601 """
602 logger.info(f"Fetching Graylog alerts for request: {request}")
603
604 hits = await fetch_alerts_from_graylog(request.index_prefix, request.size, request.timerange)
605 es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
606
607 return await process_alert_hits(hits, es_client)