Implement asynchronous vulnerability collection and existence check (#373)
taylor_socfortress committed
Dec 13, 2024 at 10:01 UTC
1808374ca896fa1a6778b4af3449cc0e026dbd6c
1 file changed
+60
-7
backend/app/agents/wazuh/services/vulnerabilities.py
+60
-7
@@ -1,3 +1,4 @@
1
+import asyncio
2
from typing import List
3
4
from fastapi import HTTPException
@@ -7,6 +8,9 @@ from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities
8
from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
9
from app.connectors.wazuh_indexer.utils.universal import collect_indices
10
from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client
11
+from app.connectors.wazuh_indexer.utils.universal import (
12
+ create_wazuh_indexer_client_async,
13
+)
14
from app.connectors.wazuh_manager.utils.universal import send_get_request
15
from app.integrations.utils.event_shipper import event_shipper
16
from app.integrations.utils.schema import EventShipperPayload
@@ -194,6 +198,45 @@ async def collect_vulnerabilities_sync(es, vulnerabilities_indices, agent_name,
198
return agent_vulnerabilities
199
200
201
+async def collect_vulnerabilities_async(es, vulnerabilities_indices, agent_name, vulnerability_severity="All"):
202
+ agent_vulnerabilities = []
203
+ for index in vulnerabilities_indices:
204
+ if vulnerability_severity == "All":
205
+ query = {
206
+ "query": {
207
+ "bool": {
208
+ "must": [
209
+ {"match": {"agent.name": agent_name}},
210
+ {"terms": {"vulnerability.severity": ["Low", "Medium", "High", "Critical"]}},
211
+ ],
212
+ },
213
+ },
214
+ }
215
+ else:
216
+ query = {
217
+ "query": {
218
+ "bool": {
219
+ "must": [{"match": {"agent.name": agent_name}}, {"match": {"vulnerability.severity": vulnerability_severity}}],
220
+ },
221
+ },
222
+ }
223
+
224
+ page = await es.search(index=index, body=query, scroll="2m")
225
+ sid = page["_scroll_id"]
226
+ scroll_size = len(page["hits"]["hits"])
227
+
228
+ while scroll_size > 0:
229
+ for hit in page["hits"]["hits"]:
230
+ vulnerability = hit["_source"]
231
+ agent_vulnerabilities.append(vulnerability)
232
+
233
+ page = await es.scroll(scroll_id=sid, scroll="2m")
234
+ sid = page["_scroll_id"]
235
+ scroll_size = len(page["hits"]["hits"])
236
+
237
+ return agent_vulnerabilities
238
+
239
+
240
def process_agent_vulnerabilities_new(agent_vulnerabilities: List[dict]) -> List[WazuhAgentVulnerabilities]:
241
logger.info(f"Processing agent vulnerabilities: {agent_vulnerabilities}")
242
@@ -229,7 +272,7 @@ def ensure_list(value):
272
return value
273
274
232
-async def check_vulnerability_exists(es, vulnerability_cve, agent_name, index_prefix):
275
+async def check_vulnerability_exists_async(es, vulnerability_cve, agent_name, index_prefix):
276
query = {
277
"query": {
278
"bool": {
@@ -241,7 +284,7 @@ async def check_vulnerability_exists(es, vulnerability_cve, agent_name, index_pr
284
},
285
}
286
index_pattern = f"{index_prefix}*"
244
- response = es.search(index=index_pattern, body=query)
287
+ response = await es.search(index=index_pattern, body=query)
288
return response["hits"]["total"]["value"] > 0
289
290
@@ -255,12 +298,12 @@ async def sync_agent_vulnerabilities(agent_name: str, customer_code: str):
298
"""
299
logger.info(f"Syncing agent {agent_name} with customer code {customer_code} vulnerabilities")
300
258
- es = await create_wazuh_indexer_client("Wazuh-Indexer")
301
+ es = await create_wazuh_indexer_client_async("Wazuh-Indexer")
302
indices = await collect_indices(all_indices=True)
303
304
vulnerabilities_indices = filter_vulnerabilities_indices(indices.indices_list)
305
263
- agent_vulnerabilities = await collect_vulnerabilities_sync(es, vulnerabilities_indices, agent_name, vulnerability_severity="All")
306
+ agent_vulnerabilities = await collect_vulnerabilities_async(es, vulnerabilities_indices, agent_name, vulnerability_severity="All")
307
308
processed_vulnerabilities = process_agent_vulnerabilities_new(agent_vulnerabilities)
309
@@ -269,16 +312,26 @@ async def sync_agent_vulnerabilities(agent_name: str, customer_code: str):
312
313
if customer_vulnerabilities_indices:
314
logger.info("Customer vulnerabilities index already exists")
272
- # ! Check to see if the vulnerability exists in the customer's index and send to Graylog if it does not exist in the customer's index ! #
273
- for vulnerability in processed_vulnerabilities:
274
- vulnerability_exists = await check_vulnerability_exists(
315
+ # Create a list of tasks for checking vulnerabilities
316
+ tasks = [
317
+ check_vulnerability_exists_async(
318
es,
319
vulnerability_cve=vulnerability.cve,
320
agent_name=agent_name,
321
index_prefix=f"wazuh-vulnerabilities-{customer_code}",
322
)
323
+ for vulnerability in processed_vulnerabilities
324
+ ]
325
326
+ # Run all tasks concurrently
327
+ results = await asyncio.gather(*tasks)
328
+
329
+ # Process the results
330
+ for vulnerability, vulnerability_exists in zip(processed_vulnerabilities, results):
331
if not vulnerability_exists:
332
+ logger.info(
333
+ f"Vulnerability {vulnerability.cve} does not exist in customer index for agent {agent_name}, sending to Graylog",
334
+ )
335
await event_shipper(
336
EventShipperPayload(
337
integration="vulnerabilities",