Opensearch async test (#371)
* Implement asynchronous Elasticsearch client creation and update query execution to support async operations * Refactor run_active_sigma_queries_endpoint to execute queries concurrently using asyncio
taylor_socfortress committed
Dec 13, 2024 at 09:06 UTC
9db0af5ab1d4173737abb5e51ff26a064106c912
3 files changed
+104
-12
backend/app/connectors/wazuh_indexer/routes/sigma.py
+61
-8
@@ -7,6 +7,7 @@ from fastapi import File
7
from fastapi import HTTPException
8
from fastapi import Query
9
from fastapi import UploadFile
10
+import asyncio
11
from loguru import logger
12
from sqlalchemy.ext.asyncio import AsyncSession
13
@@ -272,6 +273,54 @@ async def deactivate_all_sigma_queries_endpoint(
273
)
274
275
276
+# @wazuh_indexer_sigma_router.post("/run-active-queries", response_model=SigmaQueryOutResponse)
277
+# async def run_active_sigma_queries_endpoint(
278
+# index_name: str = Query(default="wazuh*"),
279
+# db: AsyncSession = Depends(get_db),
280
+# ):
281
+# """
282
+# Runs the active Sigma queries.
283
+
284
+# Args:
285
+# db (AsyncSession): The database session.
286
+
287
+# Returns:
288
+# SigmaQueryOutResponse: The Sigma queries response.
289
+# """
290
+# active_sigma_queries = await list_active_sigma_queries(db)
291
+# for query in active_sigma_queries:
292
+# time_interval_delta = parse_time_interval(query.time_interval)
293
+# logger.info(f"Time interval delta: {time_interval_delta}")
294
+# current_time = datetime.now()
295
+# logger.info(f"Current time: {current_time}")
296
+# logger.info(f"Last execution time: {query.last_execution_time}")
297
+
298
+# # Check if the current time is less than the last execution time
299
+# if current_time < query.last_execution_time or current_time - query.last_execution_time >= time_interval_delta:
300
+# logger.info(f"Running Sigma query: {query.rule_name}")
301
+# await execute_query(
302
+# RunActiveSigmaQueries(
303
+# query=query.rule_query,
304
+# time_interval=query.time_interval,
305
+# last_execution_time=query.last_execution_time,
306
+# rule_name=query.rule_name,
307
+# index=index_name,
308
+# ),
309
+# session=db,
310
+# )
311
+# # Update the last execution time to the current time and commit the changes
312
+# # ! Remove commented out code after testing ! #
313
+# query.last_execution_time = current_time
314
+# await db.commit()
315
+# else:
316
+# time_comparison = current_time - query.last_execution_time
317
+# logger.info(f"Time comparison: {time_comparison}")
318
+# logger.info(f"Skipping Sigma query because the time interval has not passed: {query.rule_name}")
319
+# return SigmaQueryOutResponse(
320
+# success=True,
321
+# message="Successfully ran the active Sigma queries.",
322
+# )
323
+
324
@wazuh_indexer_sigma_router.post("/run-active-queries", response_model=SigmaQueryOutResponse)
325
async def run_active_sigma_queries_endpoint(
326
index_name: str = Query(default="wazuh*"),
@@ -287,6 +336,8 @@ async def run_active_sigma_queries_endpoint(
336
SigmaQueryOutResponse: The Sigma queries response.
337
"""
338
active_sigma_queries = await list_active_sigma_queries(db)
339
+ tasks = []
340
+
341
for query in active_sigma_queries:
342
time_interval_delta = parse_time_interval(query.time_interval)
343
logger.info(f"Time interval delta: {time_interval_delta}")
@@ -297,7 +348,7 @@ async def run_active_sigma_queries_endpoint(
348
# Check if the current time is less than the last execution time
349
if current_time < query.last_execution_time or current_time - query.last_execution_time >= time_interval_delta:
350
logger.info(f"Running Sigma query: {query.rule_name}")
300
- await execute_query(
351
+ task = execute_query(
352
RunActiveSigmaQueries(
353
query=query.rule_query,
354
time_interval=query.time_interval,
@@ -307,14 +358,16 @@ async def run_active_sigma_queries_endpoint(
358
),
359
session=db,
360
)
310
- # Update the last execution time to the current time and commit the changes
311
- # ! Remove commented out code after testing ! #
361
+ tasks.append(task)
362
+ # Update the last execution time to the current time
363
query.last_execution_time = current_time
313
- await db.commit()
314
- else:
315
- time_comparison = current_time - query.last_execution_time
316
- logger.info(f"Time comparison: {time_comparison}")
317
- logger.info(f"Skipping Sigma query because the time interval has not passed: {query.rule_name}")
364
+
365
+ # Run all tasks concurrently
366
+ results = await asyncio.gather(*tasks)
367
+
368
+ # Commit the changes to the database
369
+ await db.commit()
370
+
371
return SigmaQueryOutResponse(
372
success=True,
373
message="Successfully ran the active Sigma queries.",
backend/app/connectors/wazuh_indexer/services/sigma/execute_query.py
+4
-3
@@ -4,9 +4,10 @@ from typing import List
4
from fastapi import HTTPException
5
from loguru import logger
6
from sqlalchemy.ext.asyncio import AsyncSession
7
+import asyncio
8
9
from app.connectors.wazuh_indexer.schema.sigma import RunActiveSigmaQueries
9
-from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
10
+from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client, create_wazuh_indexer_client_async
11
from app.incidents.schema.incident_alert import CreatedAlertPayload
12
from app.incidents.services.incident_alert import add_asset_to_copilot_alert
13
from app.incidents.services.incident_alert import build_alert_context_payload
@@ -105,7 +106,7 @@ async def send_query_to_opensearch(
106
session: AsyncSession = None,
107
) -> List[dict]:
108
try:
108
- response = es_client.search(index=index, body=query)
109
+ response = await es_client.search(index=index, body=query)
110
logger.info(f"Response: {response}")
111
hits = response["hits"]["hits"]
112
return await process_hits(hits, rule_name, session)
@@ -148,7 +149,7 @@ async def process_hits(hits, rule_name, session: AsyncSession):
149
150
151
async def execute_query(payload: RunActiveSigmaQueries, session: AsyncSession = None):
151
- client = await create_wazuh_indexer_client()
152
+ client = await create_wazuh_indexer_client_async()
153
formatted_query = await format_opensearch_query(payload.query, payload.time_interval, payload.last_execution_time)
154
logger.info(f"Executing query: {formatted_query}")
155
results = await send_query_to_opensearch(client, formatted_query, payload.rule_name, index=payload.index, session=session)
backend/app/connectors/wazuh_indexer/utils/universal.py
+39
-1
@@ -6,7 +6,7 @@ from typing import Dict
6
from typing import Iterable
7
from typing import Tuple
8
9
-from elasticsearch7 import Elasticsearch
9
+from elasticsearch7 import Elasticsearch, AsyncElasticsearch
10
from fastapi import HTTPException
11
from loguru import logger
12
@@ -112,6 +112,44 @@ async def create_wazuh_indexer_client(connector_name: str = "Wazuh-Indexer") ->
112
detail=f"Failed to create Elasticsearch client: {e}",
113
)
114
115
+async def create_wazuh_indexer_client_async(connector_name: str = "Wazuh-Indexer") -> AsyncElasticsearch:
116
+ """
117
+ Returns an Elasticsearch client for the Wazuh Indexer service.
118
+
119
+ Returns:
120
+ Elasticsearch: Elasticsearch client for the Wazuh Indexer service.
121
+ """
122
+ # attributes = get_connector_info_from_db(connector_name)
123
+ async with get_db_session() as session: # This will correctly enter the context manager
124
+ attributes = await get_connector_info_from_db(connector_name, session)
125
+ if attributes is None:
126
+ raise HTTPException(
127
+ status_code=500,
128
+ detail=f"No {connector_name} connector found in the database",
129
+ )
130
+ if attributes["connector_url"] == "https://127.1.1.1:9200":
131
+ raise HTTPException(
132
+ status_code=500,
133
+ detail=f"Please update the {connector_name} connector URL",
134
+ )
135
+ try:
136
+ return AsyncElasticsearch(
137
+ [attributes["connector_url"]],
138
+ http_auth=(
139
+ attributes["connector_username"],
140
+ attributes["connector_password"],
141
+ ),
142
+ verify_certs=False,
143
+ timeout=15,
144
+ max_retries=10,
145
+ retry_on_timeout=False,
146
+ )
147
+ except Exception as e:
148
+ raise HTTPException(
149
+ status_code=500,
150
+ detail=f"Failed to create Elasticsearch client: {e}",
151
+ )
152
+
153
154
async def format_node_allocation(node_allocation):
155
"""