Fetch more wazuh agents (#138)
* Update agent collection limit based on total affected items * Refactor agent collection logic and add delay before making another request * precommit fixes
taylor_socfortress committed
Feb 9, 2024 at 13:30 UTC
d7867fa1d18e5fa687a5b2ba80d78fe04201aa60
1 file changed
+18
-1
backend/app/agents/wazuh/services/agents.py
+18
-1
@@ -1,3 +1,5 @@
1
+import asyncio
2
+
3
from fastapi import HTTPException
4
from loguru import logger
5
@@ -18,8 +20,23 @@ async def collect_wazuh_agents() -> WazuhAgentsList:
20
logger.info("Collecting all agents from Wazuh Manager")
21
agents_collected = await send_get_request(
22
endpoint="/agents",
21
- params={"limit": 1000},
23
+ params={"limit": 500},
24
)
25
+ total_affected_items = agents_collected.get("data", {}).get("data", {}).get("total_affected_items", 0)
26
+
27
+ # If the number of agents is less than total_affected_items, make another request with the limit being total_affected_items
28
+ if len(agents_collected.get("data", {}).get("data", {}).get("affected_items", [])) < total_affected_items:
29
+ logger.info(
30
+ f"Total items: {total_affected_items}.\n"
31
+ f"Collected {len(agents_collected.get('data', {}).get('data', {}).get('affected_items', []))} agents.\n"
32
+ "Making another request.",
33
+ )
34
+ # sleep for 2 seconds before making another request
35
+ await asyncio.sleep(2)
36
+ agents_collected = await send_get_request(
37
+ endpoint="/agents",
38
+ params={"limit": total_affected_items},
39
+ )
40
41
if agents_collected.get("success") is False:
42
raise HTTPException(