376 wazuh indexer async (#377)
* Refactor alert handling to use asynchronous Elasticsearch client for improved performance and responsiveness * Refactor alert collection and incident alert services to use asynchronous Elasticsearch client for improved performance * precommit fixes
taylor_socfortress committed
Dec 13, 2024 at 15:29 UTC
30f3ec51566e859b2b8f33f81de898b2fe34352c
3 files changed
+157
-28
backend/app/connectors/wazuh_indexer/services/alerts.py
+135
-11
@@ -1,3 +1,4 @@
1
+import asyncio
2
from collections import defaultdict
3
from datetime import datetime
4
from datetime import timedelta
@@ -7,12 +8,13 @@ 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
15
-from app.connectors.wazuh_indexer.schema.alerts import Alert
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
@@ -32,6 +34,9 @@ from app.connectors.wazuh_indexer.schema.alerts import SkippableWazuhIndexerClie
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(
@@ -378,7 +383,119 @@ async def get_original_alert_id(origin_context: str) -> Tuple[str, str]:
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:
@@ -393,9 +510,8 @@ async def get_single_alert_details(
510
Returns:
511
dict: The details of the alert.
512
"""
396
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
513
try:
398
- alert = es_client.get(index=index_name, id=index_id)
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}")
@@ -414,7 +530,7 @@ async def fetch_alerts_from_graylog(index_prefix: str, size: int, timerange: str
530
Returns:
531
List[Dict]: A list of alerts.
532
"""
417
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
533
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
534
es_query = {
535
"size": size,
536
"query": {
@@ -432,12 +548,12 @@ async def fetch_alerts_from_graylog(index_prefix: str, size: int, timerange: str
548
},
549
},
550
}
435
- response = es_client.search(index=index_prefix, body=es_query)
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
440
-async def process_alert_hits(hits: List[Dict]) -> List[Dict]:
556
+async def process_alert_hits(hits: List[Dict], es_client: AsyncElasticsearch) -> List[Dict]:
557
"""
558
Processes the hits from the Graylog alert search response.
559
@@ -450,20 +566,28 @@ async def process_alert_hits(hits: List[Dict]) -> List[Dict]:
566
"""
567
alerts_dict = defaultdict(lambda: {"total_alerts": 0, "alerts": []})
568
569
+ tasks = []
570
for hit in hits:
571
origin_context = hit["_source"]["origin_context"]
572
index_name, index_id = await get_original_alert_id(origin_context)
573
logger.info(f"Fetching alert details for index {index_name} and id {index_id}")
457
- alert_details = await get_single_alert_details(index_name, index_id)
458
- logger.info(f"Alert details: {alert_details}")
574
+ task = get_single_alert_details(es_client, index_name, index_id)
575
+ tasks.append(task)
576
577
+ alert_details_list = await asyncio.gather(*tasks)
578
+
579
+ for alert_details in alert_details_list:
580
+ index_name = alert_details["_index"]
581
alerts_dict[index_name]["total_alerts"] += 1
582
alerts_dict[index_name]["alerts"].append(alert_details)
583
584
alerts = [
464
- Alert(index_name=index_name, total_alerts=data["total_alerts"], alerts=data["alerts"]) for index_name, data in alerts_dict.items()
585
+ {"index_name": index_name, "total_alerts": data["total_alerts"], "alerts": data["alerts"]}
586
+ for index_name, data in alerts_dict.items()
587
]
588
589
+ logger.info(f"Processed alerts: {alerts}")
590
+
591
return alerts
592
593
@@ -481,6 +605,6 @@ async def get_graylog_alerts(
605
logger.info(f"Fetching Graylog alerts for request: {request}")
606
607
hits = await fetch_alerts_from_graylog(request.index_prefix, request.size, request.timerange)
484
- alerts = await process_alert_hits(hits)
608
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
609
486
- return alerts
610
+ return await process_alert_hits(hits, es_client)
backend/app/incidents/services/alert_collection.py
+12
-10
@@ -1,6 +1,8 @@
1
from loguru import logger
2
3
-from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
3
+from app.connectors.wazuh_indexer.utils.universal import (
4
+ create_wazuh_indexer_client_async,
5
+)
6
from app.connectors.wazuh_indexer.utils.universal import (
7
return_graylog_events_index_names,
8
)
@@ -33,7 +35,7 @@ async def fetch_alerts_for_index(es_client, index, query):
35
Fetches alerts for a given index that match the query using the Elasticsearch scroll API.
36
"""
37
# Start the initial search request
36
- response = es_client.search(
38
+ response = await es_client.search(
39
index=index,
40
body=query,
41
scroll="2m",
@@ -44,13 +46,13 @@ async def fetch_alerts_for_index(es_client, index, query):
46
47
# Keep fetching results while there are still results to fetch
48
while len(response["hits"]["hits"]):
47
- response = es_client.scroll(scroll_id=scroll_id, scroll="2m") # Extend the scroll context for another 2 minutes
49
+ response = await es_client.scroll(scroll_id=scroll_id, scroll="2m") # Extend the scroll context for another 2 minutes
50
# Update the scroll ID in case it changes
51
scroll_id = response["_scroll_id"]
52
hits.extend(response["hits"]["hits"])
53
54
# Close the scroll context
53
- es_client.clear_scroll(scroll_id=scroll_id)
55
+ await es_client.clear_scroll(scroll_id=scroll_id)
56
57
return [AlertPayloadItem(**hit) for hit in hits]
58
@@ -61,7 +63,7 @@ async def get_alerts_not_created_in_copilot() -> AlertsPayload:
63
"""
64
indices = await return_graylog_events_index_names()
65
logger.info(f"Indices: {indices}")
64
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
66
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
67
query = await construct_query()
68
69
alerts_not_created = []
@@ -99,10 +101,10 @@ async def add_copilot_alert_id(index_data: CreateAlertRequest, alert_id: int):
101
"""
102
Add the CoPilot alert ID to the Graylog event.
103
"""
102
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
104
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
105
body = {"doc": {"fields": {"COPILOT_ALERT_ID": f"{alert_id}"}}}
106
try:
105
- es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body)
107
+ await es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body)
108
logger.info(f"Added CoPilot alert ID {alert_id} to Graylog event {index_data.alert_id} in index {index_data.index_name}")
109
except Exception as e:
110
logger.error(
@@ -111,17 +113,17 @@ async def add_copilot_alert_id(index_data: CreateAlertRequest, alert_id: int):
113
114
# Attempt to remove read-only block
115
try:
114
- es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": None})
116
+ await es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": None})
117
logger.info(f"Removed read-only block from index {index_data.index_name}. Retrying update.")
118
119
# Retry the update operation
118
- es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body)
120
+ await es_client.update(index=index_data.index_name, id=index_data.alert_id, body=body)
121
logger.info(
122
f"Added CoPilot alert ID {alert_id} to Graylog event {index_data.alert_id} in index {index_data.index_name} after removing read-only block",
123
)
124
125
# Re-enable the write block
124
- es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": True})
126
+ await es_client.indices.put_settings(index=index_data.index_name, body={"index.blocks.write": True})
127
except Exception as e2:
128
logger.error(f"Failed to remove read-only block from index {index_data.index_name}: {e2}")
129
backend/app/incidents/services/incident_alert.py
+10
-7
@@ -15,6 +15,9 @@ from sqlalchemy.future import select
15
from app.connectors.shuffle.schema.integrations import ExecuteWorkflowRequest
16
from app.connectors.shuffle.services.integrations import execute_workflow
17
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
18
+from app.connectors.wazuh_indexer.utils.universal import (
19
+ create_wazuh_indexer_client_async,
20
+)
21
from app.db.universal_models import Agents
22
from app.incidents.models import Alert
23
from app.incidents.models import AlertContext
@@ -112,9 +115,9 @@ async def get_single_alert_details(
115
logger.info(
116
f"Fetching alert details for alert {alert_details.alert_id} in index {alert_details.index_name}",
117
)
115
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
118
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
119
try:
117
- alert = es_client.get(index=alert_details.index_name, id=alert_details.alert_id)
120
+ alert = await es_client.get(index=alert_details.index_name, id=alert_details.alert_id)
121
source_model = GenericSourceModel(**alert["_source"])
122
syslog_type = getattr(source_model, "syslog_type", None)
123
if syslog_type is None:
@@ -236,9 +239,9 @@ async def add_alert_to_document(
239
Returns:
240
- True if the update is successful, False otherwise.
241
"""
239
- es_client = await create_wazuh_indexer_client("Wazuh-Indexer")
242
+ es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
243
try:
241
- es_client.update(
244
+ await es_client.update(
245
index=alert.index_name,
246
id=alert.alert_id,
247
body={"doc": {"alert_id": soc_alert_id}},
@@ -253,7 +256,7 @@ async def add_alert_to_document(
256
)
257
# Attempt to remove read-only block
258
try:
256
- es_client.indices.put_settings(
259
+ await es_client.indices.put_settings(
260
index=alert.index_name,
261
body={"index.blocks.write": None},
262
)
@@ -262,7 +265,7 @@ async def add_alert_to_document(
265
)
266
267
# Retry the update operation
265
- es_client.update(
268
+ await es_client.update(
269
index=alert.index_name,
270
id=alert.alert_id,
271
body={"doc": {"alert_id": soc_alert_id}},
@@ -272,7 +275,7 @@ async def add_alert_to_document(
275
)
276
277
# Reenable the write block
275
- es_client.indices.put_settings(
278
+ await es_client.indices.put_settings(
279
index=alert.index_name,
280
body={"index.blocks.write": True},
281
)