@cryptotaxi247 / CoPilot / commits / a3c0055a

feat: implement concurrency control for SCA data collection and add t… (#643)

* feat: implement concurrency control for SCA data collection and add token caching for Wazuh Manager * precommit-fixes

taylor_socfortress committed Jan 30, 2026 at 09:14 UTC a3c0055ac0ff90a1e08a6dcade0b633d7a591abd
4 files changed +606 -108
backend/app/agents/sca/services/sca.py
+289 -56
@@ -1,7 +1,9 @@
1 +import asyncio
2 import csv
3 import hashlib
4 import io
5 import json
6 +from asyncio import Semaphore
7 from datetime import datetime
8 from typing import Any
9 from typing import Dict
@@ -30,6 +32,9 @@ from app.db.universal_models import Agents
32 from app.db.universal_models import SCAReport
33 from app.middleware.customer_access import customer_access_handler
34
35 +# Default concurrency limit for parallel API requests
36 +DEFAULT_MAX_CONCURRENT_REQUESTS = 100
37 +
38
39 async def get_all_agents_from_db(
40 db_session: AsyncSession,
@@ -61,6 +66,242 @@ async def get_all_agents_from_db(
66 raise HTTPException(status_code=500, detail=f"Failed to fetch agents: {e}")
67
68
69 +# async def collect_sca_for_all_agents(
70 +# db_session: AsyncSession,
71 +# customer_code: Optional[str] = None,
72 +# agent_name: Optional[str] = None,
73 +# policy_id: Optional[str] = None,
74 +# policy_name: Optional[str] = None,
75 +# min_score: Optional[int] = None,
76 +# max_score: Optional[int] = None,
77 +# ) -> List[AgentScaOverviewItem]:
78 +# """
79 +# Collect SCA results for all agents from Wazuh Manager
80 +
81 +# Args:
82 +# db_session: Database session to use
83 +# customer_code: Optional customer code filter
84 +# agent_name: Optional agent name filter
85 +# policy_id: Optional policy ID filter
86 +# policy_name: Optional policy name filter (partial matching)
87 +# min_score: Optional minimum score filter
88 +# max_score: Optional maximum score filter
89 +
90 +# Returns:
91 +# List of AgentScaOverviewItem objects
92 +# """
93 +# try:
94 +# # Get agents from database
95 +# agents = await get_all_agents_from_db(db_session, customer_code)
96 +
97 +# if not agents:
98 +# logger.warning("No agents found" + (f" for customer {customer_code}" if customer_code else ""))
99 +# return []
100 +
101 +# all_sca_results = []
102 +
103 +# for agent in agents:
104 +# # Skip if agent name filter is specified and doesn't match
105 +# if agent_name and agent.hostname != agent_name:
106 +# continue
107 +
108 +# try:
109 +# logger.info(f"Collecting SCA results for agent: {agent.hostname}")
110 +
111 +# # Collect SCA data from Wazuh Manager for this agent
112 +# sca_response = await collect_agent_sca(agent.agent_id)
113 +
114 +# if not sca_response.success or not sca_response.sca:
115 +# logger.warning(f"No SCA data found for agent {agent.hostname}")
116 +# continue
117 +
118 +# # Process each SCA policy result for this agent
119 +# for sca_result in sca_response.sca:
120 +# # Apply filters
121 +# if policy_id and sca_result.policy_id != policy_id:
122 +# continue
123 +
124 +# if policy_name and policy_name.lower() not in sca_result.name.lower():
125 +# continue
126 +
127 +# if min_score is not None and sca_result.score < min_score:
128 +# continue
129 +
130 +# if max_score is not None and sca_result.score > max_score:
131 +# continue
132 +
133 +# # Create overview item
134 +# overview_item = AgentScaOverviewItem(
135 +# agent_id=agent.agent_id,
136 +# agent_name=agent.hostname,
137 +# customer_code=agent.customer_code,
138 +# policy_id=sca_result.policy_id,
139 +# policy_name=sca_result.name,
140 +# description=sca_result.description,
141 +# total_checks=sca_result.total_checks,
142 +# pass_count=sca_result.pass_count,
143 +# fail_count=sca_result.fail,
144 +# invalid_count=sca_result.invalid,
145 +# score=sca_result.score,
146 +# start_scan=sca_result.start_scan,
147 +# end_scan=sca_result.end_scan,
148 +# references=sca_result.references,
149 +# hash_file=sca_result.hash_file,
150 +# )
151 +
152 +# all_sca_results.append(overview_item)
153 +
154 +# except Exception as e:
155 +# logger.error(f"Error collecting SCA for agent {agent.hostname}: {e}")
156 +# # Continue with other agents even if one fails
157 +# continue
158 +
159 +# logger.info(f"Collected SCA results for {len(all_sca_results)} policy results across agents")
160 +# return all_sca_results
161 +
162 +# except Exception as e:
163 +# logger.error(f"Error collecting SCA for all agents: {e}")
164 +# raise HTTPException(status_code=500, detail=f"Failed to collect SCA results: {e}")
165 +
166 +
167 +# async def search_sca_overview(
168 +# db_session: AsyncSession,
169 +# customer_code: Optional[str] = None,
170 +# agent_name: Optional[str] = None,
171 +# policy_id: Optional[str] = None,
172 +# policy_name: Optional[str] = None,
173 +# min_score: Optional[int] = None,
174 +# max_score: Optional[int] = None,
175 +# page: int = 1,
176 +# page_size: int = 50,
177 +# ) -> ScaOverviewResponse:
178 +# """
179 +# Search SCA results across all agents with filtering and pagination
180 +
181 +# Args:
182 +# db_session: Database session for agent lookup
183 +# customer_code: Optional customer code filter
184 +# agent_name: Optional agent hostname filter
185 +# policy_id: Optional policy ID filter
186 +# policy_name: Optional policy name filter (partial matching)
187 +# min_score: Optional minimum score filter
188 +# max_score: Optional maximum score filter
189 +# page: Page number for pagination
190 +# page_size: Number of results per page
191 +
192 +# Returns:
193 +# ScaOverviewResponse with paginated results and statistics
194 +# """
195 +# logger.info(
196 +# f"Searching SCA overview with filters: customer_code={customer_code}, "
197 +# f"agent_name={agent_name}, policy_id={policy_id}, policy_name={policy_name}, "
198 +# f"min_score={min_score}, max_score={max_score}, page={page}, page_size={page_size}",
199 +# )
200 +
201 +# # Build filters applied dict for response
202 +# filters_applied = {}
203 +# if customer_code:
204 +# filters_applied["customer_code"] = customer_code
205 +# if agent_name:
206 +# filters_applied["agent_name"] = agent_name
207 +# if policy_id:
208 +# filters_applied["policy_id"] = policy_id
209 +# if policy_name:
210 +# filters_applied["policy_name"] = policy_name
211 +# if min_score is not None:
212 +# filters_applied["min_score"] = min_score
213 +# if max_score is not None:
214 +# filters_applied["max_score"] = max_score
215 +
216 +# try:
217 +# # Collect all SCA results with filtering
218 +# all_sca_results = await collect_sca_for_all_agents(
219 +# db_session=db_session,
220 +# customer_code=customer_code,
221 +# agent_name=agent_name,
222 +# policy_id=policy_id,
223 +# policy_name=policy_name,
224 +# min_score=min_score,
225 +# max_score=max_score,
226 +# )
227 +
228 +
229 +async def collect_sca_for_single_agent(
230 + agent: Agents,
231 + semaphore: Semaphore,
232 + policy_id: Optional[str] = None,
233 + policy_name: Optional[str] = None,
234 + min_score: Optional[int] = None,
235 + max_score: Optional[int] = None,
236 +) -> List[AgentScaOverviewItem]:
237 + """
238 + Collect SCA results for a single agent with semaphore-based rate limiting.
239 +
240 + Args:
241 + agent: Agent to collect SCA data for
242 + semaphore: Semaphore to limit concurrent requests
243 + policy_id: Optional policy ID filter
244 + policy_name: Optional policy name filter (partial matching)
245 + min_score: Optional minimum score filter
246 + max_score: Optional maximum score filter
247 +
248 + Returns:
249 + List of AgentScaOverviewItem objects for this agent
250 + """
251 + async with semaphore:
252 + try:
253 + logger.debug(f"Collecting SCA results for agent: {agent.hostname} (ID: {agent.agent_id})")
254 +
255 + # Collect SCA data from Wazuh Manager for this agent
256 + sca_response = await collect_agent_sca(agent.agent_id)
257 +
258 + if not sca_response.success or not sca_response.sca:
259 + logger.debug(f"No SCA data for agent {agent.hostname}")
260 + return []
261 +
262 + results = []
263 +
264 + # Process each SCA policy result for this agent
265 + for sca_result in sca_response.sca:
266 + # Apply filters
267 + if policy_id and sca_result.policy_id != policy_id:
268 + continue
269 + if policy_name and policy_name.lower() not in sca_result.name.lower():
270 + continue
271 + if min_score is not None and sca_result.score < min_score:
272 + continue
273 + if max_score is not None and sca_result.score > max_score:
274 + continue
275 +
276 + # Create overview item
277 + overview_item = AgentScaOverviewItem(
278 + agent_id=agent.agent_id,
279 + agent_name=agent.hostname,
280 + customer_code=agent.customer_code,
281 + policy_id=sca_result.policy_id,
282 + policy_name=sca_result.name,
283 + description=sca_result.description,
284 + total_checks=sca_result.total_checks,
285 + pass_count=sca_result.pass_count,
286 + fail_count=sca_result.fail,
287 + invalid_count=sca_result.invalid,
288 + score=sca_result.score,
289 + start_scan=sca_result.start_scan,
290 + end_scan=sca_result.end_scan,
291 + references=sca_result.references,
292 + hash_file=sca_result.hash_file,
293 + )
294 + results.append(overview_item)
295 +
296 + logger.debug(f"Collected {len(results)} SCA results for agent {agent.hostname}")
297 + return results
298 +
299 + except Exception as e:
300 + logger.error(f"Error collecting SCA for agent {agent.hostname}: {e}")
301 + # Return empty list instead of raising - allows other agents to continue
302 + return []
303 +
304 +
305 async def collect_sca_for_all_agents(
306 db_session: AsyncSession,
307 customer_code: Optional[str] = None,
@@ -69,9 +310,10 @@ async def collect_sca_for_all_agents(
310 policy_name: Optional[str] = None,
311 min_score: Optional[int] = None,
312 max_score: Optional[int] = None,
313 + max_concurrent_requests: int = DEFAULT_MAX_CONCURRENT_REQUESTS,
314 ) -> List[AgentScaOverviewItem]:
315 """
74 - Collect SCA results for all agents from Wazuh Manager
316 + Collect SCA results for all agents from Wazuh Manager using parallel requests.
317
318 Args:
319 db_session: Database session to use
@@ -81,6 +323,7 @@ async def collect_sca_for_all_agents(
323 policy_name: Optional policy name filter (partial matching)
324 min_score: Optional minimum score filter
325 max_score: Optional maximum score filter
326 + max_concurrent_requests: Maximum number of concurrent API requests (default: 10)
327
328 Returns:
329 List of AgentScaOverviewItem objects
@@ -93,65 +336,52 @@ async def collect_sca_for_all_agents(
336 logger.warning("No agents found" + (f" for customer {customer_code}" if customer_code else ""))
337 return []
338
96 - all_sca_results = []
97 -
98 - for agent in agents:
99 - # Skip if agent name filter is specified and doesn't match
100 - if agent_name and agent.hostname != agent_name:
101 - continue
102 -
103 - try:
104 - logger.info(f"Collecting SCA results for agent: {agent.hostname}")
339 + # Filter agents by name if specified (do this before parallel processing)
340 + if agent_name:
341 + agents = [a for a in agents if a.hostname == agent_name]
342 + if not agents:
343 + logger.info(f"No agents found matching hostname: {agent_name}")
344 + return []
345 +
346 + logger.info(f"Collecting SCA data for {len(agents)} agents " f"(max concurrent: {max_concurrent_requests})")
347 +
348 + # Create semaphore to limit concurrent requests to Wazuh Manager
349 + semaphore = Semaphore(max_concurrent_requests)
350 +
351 + # Create tasks for parallel execution
352 + tasks = [
353 + collect_sca_for_single_agent(
354 + agent=agent,
355 + semaphore=semaphore,
356 + policy_id=policy_id,
357 + policy_name=policy_name,
358 + min_score=min_score,
359 + max_score=max_score,
360 + )
361 + for agent in agents
362 + ]
363
106 - # Collect SCA data from Wazuh Manager for this agent
107 - sca_response = await collect_agent_sca(agent.agent_id)
364 + # Execute all tasks concurrently with gather
365 + # return_exceptions=True prevents one failure from canceling all tasks
366 + results = await asyncio.gather(*tasks, return_exceptions=True)
367
109 - if not sca_response.success or not sca_response.sca:
110 - logger.warning(f"No SCA data found for agent {agent.hostname}")
111 - continue
368 + # Flatten results and filter out exceptions
369 + all_sca_results = []
370 + successful_agents = 0
371 + failed_agents = 0
372
113 - # Process each SCA policy result for this agent
114 - for sca_result in sca_response.sca:
115 - # Apply filters
116 - if policy_id and sca_result.policy_id != policy_id:
117 - continue
118 -
119 - if policy_name and policy_name.lower() not in sca_result.name.lower():
120 - continue
121 -
122 - if min_score is not None and sca_result.score < min_score:
123 - continue
124 -
125 - if max_score is not None and sca_result.score > max_score:
126 - continue
127 -
128 - # Create overview item
129 - overview_item = AgentScaOverviewItem(
130 - agent_id=agent.agent_id,
131 - agent_name=agent.hostname,
132 - customer_code=agent.customer_code,
133 - policy_id=sca_result.policy_id,
134 - policy_name=sca_result.name,
135 - description=sca_result.description,
136 - total_checks=sca_result.total_checks,
137 - pass_count=sca_result.pass_count,
138 - fail_count=sca_result.fail,
139 - invalid_count=sca_result.invalid,
140 - score=sca_result.score,
141 - start_scan=sca_result.start_scan,
142 - end_scan=sca_result.end_scan,
143 - references=sca_result.references,
144 - hash_file=sca_result.hash_file,
145 - )
146 -
147 - all_sca_results.append(overview_item)
148 -
149 - except Exception as e:
150 - logger.error(f"Error collecting SCA for agent {agent.hostname}: {e}")
151 - # Continue with other agents even if one fails
373 + for i, result in enumerate(results):
374 + if isinstance(result, Exception):
375 + logger.error(f"Task for agent {agents[i].hostname} failed: {result}")
376 + failed_agents += 1
377 continue
378 + if isinstance(result, list):
379 + if result: # Only count as successful if we got data
380 + successful_agents += 1
381 + all_sca_results.extend(result)
382 +
383 + logger.info(f"Collected {len(all_sca_results)} SCA policy results from " f"{successful_agents} agents ({failed_agents} failed)")
384
154 - logger.info(f"Collected SCA results for {len(all_sca_results)} policy results across agents")
385 return all_sca_results
386
387 except Exception as e:
@@ -169,6 +399,7 @@ async def search_sca_overview(
399 max_score: Optional[int] = None,
400 page: int = 1,
401 page_size: int = 50,
402 + max_concurrent_requests: int = DEFAULT_MAX_CONCURRENT_REQUESTS,
403 ) -> ScaOverviewResponse:
404 """
405 Search SCA results across all agents with filtering and pagination
@@ -183,6 +414,7 @@ async def search_sca_overview(
414 max_score: Optional maximum score filter
415 page: Page number for pagination
416 page_size: Number of results per page
417 + max_concurrent_requests: Maximum concurrent API requests (default: 10)
418
419 Returns:
420 ScaOverviewResponse with paginated results and statistics
@@ -209,7 +441,7 @@ async def search_sca_overview(
441 filters_applied["max_score"] = max_score
442
443 try:
212 - # Collect all SCA results with filtering
444 + # Collect all SCA results with filtering (now uses parallel requests)
445 all_sca_results = await collect_sca_for_all_agents(
446 db_session=db_session,
447 customer_code=customer_code,
@@ -218,6 +450,7 @@ async def search_sca_overview(
450 policy_name=policy_name,
451 min_score=min_score,
452 max_score=max_score,
453 + max_concurrent_requests=max_concurrent_requests,
454 )
455
456 # Sort results by agent's minimum score (lowest first)
backend/app/connectors/wazuh_indexer/routes/sigma.py
+33 -33
@@ -1,4 +1,3 @@
1 -import asyncio
1 import os
2 from datetime import datetime
3
@@ -336,38 +335,39 @@ async def run_active_sigma_queries_endpoint(
335 Returns:
336 SigmaQueryOutResponse: The Sigma queries response.
337 """
339 - active_sigma_queries = await list_active_sigma_queries(db)
340 - tasks = []
341 -
342 - for query in active_sigma_queries:
343 - time_interval_delta = parse_time_interval(query.time_interval)
344 - logger.info(f"Time interval delta: {time_interval_delta}")
345 - current_time = datetime.now()
346 - logger.info(f"Current time: {current_time}")
347 - logger.info(f"Last execution time: {query.last_execution_time}")
348 -
349 - # Check if the current time is less than the last execution time
350 - if current_time < query.last_execution_time or current_time - query.last_execution_time >= time_interval_delta:
351 - logger.info(f"Running Sigma query: {query.rule_name}")
352 - task = execute_query(
353 - RunActiveSigmaQueries(
354 - query=query.rule_query,
355 - time_interval=query.time_interval,
356 - last_execution_time=query.last_execution_time,
357 - rule_name=query.rule_name,
358 - index=index_name,
359 - ),
360 - session=db,
361 - )
362 - tasks.append(task)
363 - # Update the last execution time to the current time
364 - query.last_execution_time = current_time
365 -
366 - # Run all tasks concurrently
367 - await asyncio.gather(*tasks)
368 -
369 - # Commit the changes to the database
370 - await db.commit()
338 + # ! Commenting Out for now, will revisit later if needed ! #
339 + # active_sigma_queries = await list_active_sigma_queries(db)
340 + # tasks = []
341 +
342 + # for query in active_sigma_queries:
343 + # time_interval_delta = parse_time_interval(query.time_interval)
344 + # logger.info(f"Time interval delta: {time_interval_delta}")
345 + # current_time = datetime.now()
346 + # logger.info(f"Current time: {current_time}")
347 + # logger.info(f"Last execution time: {query.last_execution_time}")
348 +
349 + # # Check if the current time is less than the last execution time
350 + # if current_time < query.last_execution_time or current_time - query.last_execution_time >= time_interval_delta:
351 + # logger.info(f"Running Sigma query: {query.rule_name}")
352 + # task = execute_query(
353 + # RunActiveSigmaQueries(
354 + # query=query.rule_query,
355 + # time_interval=query.time_interval,
356 + # last_execution_time=query.last_execution_time,
357 + # rule_name=query.rule_name,
358 + # index=index_name,
359 + # ),
360 + # session=db,
361 + # )
362 + # tasks.append(task)
363 + # # Update the last execution time to the current time
364 + # query.last_execution_time = current_time
365 +
366 + # # Run all tasks concurrently
367 + # await asyncio.gather(*tasks)
368 +
369 + # # Commit the changes to the database
370 + # await db.commit()
371
372 return SigmaQueryOutResponse(
373 success=True,
backend/app/connectors/wazuh_manager/utils/universal.py
+283 -19
@@ -1,4 +1,8 @@
1 +import asyncio
2 import json
3 +from dataclasses import dataclass
4 +from datetime import datetime
5 +from datetime import timedelta
6 from typing import Any
7 from typing import Dict
8 from typing import Optional
@@ -10,6 +14,109 @@ from app.connectors.utils import get_connector_info_from_db
14 from app.db.db_session import AsyncSessionLocal
15 from app.db.db_session import get_db_session
16
17 +# =============================================================================
18 +# Token Cache Implementation
19 +# =============================================================================
20 +
21 +
22 +@dataclass
23 +class CachedToken:
24 + """Cached authentication token with expiration"""
25 +
26 + token: str
27 + expires_at: datetime
28 + connector_url: str
29 +
30 +
31 +class WazuhTokenCache:
32 + """
33 + Thread-safe cache for Wazuh Manager authentication tokens.
34 +
35 + Caches tokens per connector name to support multiple Wazuh Manager instances.
36 + Default TTL is 10 minutes (Wazuh tokens typically expire after 15-30 minutes).
37 + """
38 +
39 + def __init__(self, default_ttl_minutes: int = 10):
40 + self._cache: Dict[str, CachedToken] = {}
41 + self._lock = asyncio.Lock()
42 + self._default_ttl = timedelta(minutes=default_ttl_minutes)
43 +
44 + async def get(self, connector_name: str) -> Optional[Dict[str, str]]:
45 + """
46 + Get cached token headers if valid.
47 +
48 + Returns:
49 + Dict with Authorization header if token is valid, None otherwise
50 + """
51 + async with self._lock:
52 + if connector_name not in self._cache:
53 + return None
54 +
55 + cached = self._cache[connector_name]
56 +
57 + # Check if token is expired (with 30 second buffer)
58 + if datetime.utcnow() >= (cached.expires_at - timedelta(seconds=30)):
59 + logger.debug(f"Cached token for {connector_name} has expired")
60 + del self._cache[connector_name]
61 + return None
62 +
63 + logger.debug(f"Using cached token for {connector_name}")
64 + return {"Authorization": f"Bearer {cached.token}"}
65 +
66 + async def set(self, connector_name: str, token: str, connector_url: str, ttl_minutes: Optional[int] = None):
67 + """Cache a new token"""
68 + async with self._lock:
69 + ttl = timedelta(minutes=ttl_minutes) if ttl_minutes else self._default_ttl
70 + expires_at = datetime.utcnow() + ttl
71 +
72 + self._cache[connector_name] = CachedToken(
73 + token=token,
74 + expires_at=expires_at,
75 + connector_url=connector_url,
76 + )
77 + logger.debug(f"Cached token for {connector_name}, expires at {expires_at.isoformat()}")
78 +
79 + async def invalidate(self, connector_name: str):
80 + """Remove cached token for a connector"""
81 + async with self._lock:
82 + if connector_name in self._cache:
83 + del self._cache[connector_name]
84 + logger.debug(f"Invalidated cached token for {connector_name}")
85 +
86 + async def clear(self):
87 + """Clear all cached tokens"""
88 + async with self._lock:
89 + self._cache.clear()
90 + logger.debug("Cleared all cached Wazuh tokens")
91 +
92 +
93 +# Global token cache instance
94 +_token_cache = WazuhTokenCache(default_ttl_minutes=10)
95 +
96 +
97 +# =============================================================================
98 +# Public Cache Management Functions
99 +# =============================================================================
100 +
101 +
102 +async def invalidate_wazuh_token_cache(connector_name: str = "Wazuh-Manager"):
103 + """
104 + Invalidate cached token for a connector.
105 +
106 + Call this when credentials are updated or if you receive auth errors.
107 + """
108 + await _token_cache.invalidate(connector_name)
109 +
110 +
111 +# ============================================================================
112 +# Existing Wazuh Manager Utility Functions
113 +# ============================================================================
114 +
115 +
116 +async def clear_all_wazuh_token_caches():
117 + """Clear all cached Wazuh tokens"""
118 + await _token_cache.clear()
119 +
120
121 async def verify_wazuh_manager_credentials(
122 attributes: Dict[str, Any],
@@ -76,25 +183,82 @@ async def verify_wazuh_manager_connection(connector_name: str) -> str:
183 return await verify_wazuh_manager_credentials(attributes)
184
185
79 -async def create_wazuh_manager_client(connector_name: str) -> str:
186 +# async def create_wazuh_manager_client(connector_name: str) -> str:
187 +# """
188 +# Returns the authentication token for the Wazuh manager service.
189 +
190 +# Returns:
191 +# str: Authentication token for the Wazuh manager service.
192 +# """
193 +# logger.info("Getting Wazuh Manager authentication token")
194 +# # attributes = get_connector_info_from_db(connector_name)
195 +# async with AsyncSessionLocal() as session:
196 +# attributes = await get_connector_info_from_db(connector_name, session)
197 +# if attributes is None:
198 +# logger.error("No Wazuh Manager connector found in the database")
199 +# return None
200 +# logger.info(
201 +# f"Verifying the wazuh-manager connection to {attributes['connector_url']}",
202 +# )
203 +# try:
204 +# wazuh_auth_token = requests.get(
205 +# f"{attributes['connector_url']}/security/user/authenticate",
206 +# auth=(
207 +# attributes["connector_username"],
208 +# attributes["connector_password"],
209 +# ),
210 +# verify=False,
211 +# )
212 +
213 +# if wazuh_auth_token.status_code == 200:
214 +# logger.debug("Wazuh Authentication Token successful")
215 +# wazuh_auth_token = wazuh_auth_token.json()
216 +# wazuh_auth_token = wazuh_auth_token["data"]["token"]
217 +
218 +# return {"Authorization": f"Bearer {wazuh_auth_token}"}
219 +# else:
220 +# logger.error(
221 +# f"Connection to {attributes['connector_url']} failed with error: {wazuh_auth_token.text}",
222 +# )
223 +
224 +# return None
225 +# except Exception as e:
226 +# logger.error(
227 +# f"Connection to {attributes['connector_url']} failed with error: {e}",
228 +# )
229 +
230 +# return None
231 +
232 +
233 +async def create_wazuh_manager_client(connector_name: str) -> Optional[Dict[str, str]]:
234 """
81 - Returns the authentication token for the Wazuh manager service.
235 + Returns the authentication token headers for the Wazuh manager service.
236 +
237 + Uses cached token if available and valid, otherwise fetches a new one.
238
239 Returns:
84 - str: Authentication token for the Wazuh manager service.
240 + Dict with Authorization header, or None if authentication fails
241 """
86 - logger.info("Getting Wazuh Manager authentication token")
87 - # attributes = get_connector_info_from_db(connector_name)
242 + # Check cache first
243 + cached_headers = await _token_cache.get(connector_name)
244 + if cached_headers is not None:
245 + return cached_headers
246 +
247 + logger.info(f"Fetching new Wazuh Manager authentication token for {connector_name}")
248 +
249 async with AsyncSessionLocal() as session:
250 attributes = await get_connector_info_from_db(connector_name, session)
251 +
252 if attributes is None:
253 logger.error("No Wazuh Manager connector found in the database")
254 return None
255 +
256 logger.info(
94 - f"Verifying the wazuh-manager connection to {attributes['connector_url']}",
257 + f"Authenticating to wazuh-manager at {attributes['connector_url']}",
258 )
259 +
260 try:
97 - wazuh_auth_token = requests.get(
261 + response = requests.get(
262 f"{attributes['connector_url']}/security/user/authenticate",
263 auth=(
264 attributes["connector_username"],
@@ -103,26 +267,96 @@ async def create_wazuh_manager_client(connector_name: str) -> str:
267 verify=False,
268 )
269
106 - if wazuh_auth_token.status_code == 200:
270 + if response.status_code == 200:
271 logger.debug("Wazuh Authentication Token successful")
108 - wazuh_auth_token = wazuh_auth_token.json()
109 - wazuh_auth_token = wazuh_auth_token["data"]["token"]
272 + token_data = response.json()
273 + token = token_data["data"]["token"]
274 +
275 + # Cache the token
276 + await _token_cache.set(
277 + connector_name=connector_name,
278 + token=token,
279 + connector_url=attributes["connector_url"],
280 + )
281
111 - return {"Authorization": f"Bearer {wazuh_auth_token}"}
282 + return {"Authorization": f"Bearer {token}"}
283 else:
284 logger.error(
114 - f"Connection to {attributes['connector_url']} failed with error: {wazuh_auth_token.text}",
285 + f"Connection to {attributes['connector_url']} failed with error: {response.text}",
286 )
116 -
287 return None
288 +
289 except Exception as e:
290 logger.error(
291 f"Connection to {attributes['connector_url']} failed with error: {e}",
292 )
122 -
293 return None
294
295
296 +# async def send_get_request(
297 +# endpoint: str,
298 +# params: Optional[Dict[str, Any]] = None,
299 +# connector_name: str = "Wazuh-Manager",
300 +# ) -> Dict[str, Any]:
301 +# """
302 +# Sends a GET request to the Wazuh Manager service.
303 +
304 +# Args:
305 +# endpoint (str): The endpoint to send the GET request to.
306 +# params (Optional[Dict[str, Any]], optional): The parameters to send with the GET request. Defaults to None.
307 +# connector_name (str, optional): The name of the connector to use. Defaults to "Wazuh-Manager".
308 +
309 +# Returns:
310 +# Dict[str, Any]: The response from the GET request.
311 +# """
312 +# logger.info(f"Sending GET request to {endpoint}")
313 +# wazuh_manager_client = await create_wazuh_manager_client(connector_name)
314 +# # attributes = get_connector_info_from_db(connector_name)
315 +# async with AsyncSessionLocal() as session:
316 +# attributes = await get_connector_info_from_db(connector_name, session)
317 +
318 +# if attributes is None:
319 +# logger.error("No Wazuh Manager connector found in the database")
320 +# return None
321 +# try:
322 +# # Check if raw response is requested - support both old and new ways
323 +# # Old way: params == {"raw": True} (exact match for backward compatibility)
324 +# # New way: params contains "raw": True (for requests with multiple parameters)
325 +# is_raw_request = (params == {"raw": True}) or (params and params.get("raw", False))
326 +
327 +# if is_raw_request:
328 +# response = requests.get(
329 +# f"{attributes['connector_url']}{endpoint}",
330 +# headers=wazuh_manager_client,
331 +# params=params,
332 +# verify=False,
333 +# )
334 +# response.raise_for_status()
335 +# return {
336 +# "data": response.text,
337 +# "success": True,
338 +# "message": "Successfully retrieved data",
339 +# }
340 +# response = requests.get(
341 +# f"{attributes['connector_url']}{endpoint}",
342 +# headers=wazuh_manager_client,
343 +# params=params,
344 +# verify=False,
345 +# )
346 +# response.raise_for_status()
347 +# return {
348 +# "data": response.json(),
349 +# "success": True,
350 +# "message": "Successfully retrieved data",
351 +# }
352 +# except Exception as e:
353 +# logger.error(f"Failed to send GET request to {endpoint} with error: {e}")
354 +# return {
355 +# "success": False,
356 +# "message": f"Failed to send GET request to {endpoint} with error: {e}",
357 +# }
358 +
359 +
360 async def send_get_request(
361 endpoint: str,
362 params: Optional[Dict[str, Any]] = None,
@@ -141,17 +375,25 @@ async def send_get_request(
375 """
376 logger.info(f"Sending GET request to {endpoint}")
377 wazuh_manager_client = await create_wazuh_manager_client(connector_name)
144 - # attributes = get_connector_info_from_db(connector_name)
378 +
379 + if wazuh_manager_client is None:
380 + logger.error("Failed to get Wazuh Manager client")
381 + return {
382 + "success": False,
383 + "message": "Failed to authenticate with Wazuh Manager",
384 + }
385 +
386 async with AsyncSessionLocal() as session:
387 attributes = await get_connector_info_from_db(connector_name, session)
388
389 if attributes is None:
390 logger.error("No Wazuh Manager connector found in the database")
150 - return None
391 + return {
392 + "success": False,
393 + "message": "No Wazuh Manager connector found in the database",
394 + }
395 +
396 try:
152 - # Check if raw response is requested - support both old and new ways
153 - # Old way: params == {"raw": True} (exact match for backward compatibility)
154 - # New way: params contains "raw": True (for requests with multiple parameters)
397 is_raw_request = (params == {"raw": True}) or (params and params.get("raw", False))
398
399 if is_raw_request:
@@ -167,12 +409,34 @@ async def send_get_request(
409 "success": True,
410 "message": "Successfully retrieved data",
411 }
412 +
413 response = requests.get(
414 f"{attributes['connector_url']}{endpoint}",
415 headers=wazuh_manager_client,
416 params=params,
417 verify=False,
418 )
419 +
420 + # Handle 401 Unauthorized - token may have expired on server side
421 + if response.status_code == 401:
422 + logger.warning("Received 401 Unauthorized, invalidating cached token and retrying")
423 + await _token_cache.invalidate(connector_name)
424 +
425 + # Retry with fresh token
426 + wazuh_manager_client = await create_wazuh_manager_client(connector_name)
427 + if wazuh_manager_client is None:
428 + return {
429 + "success": False,
430 + "message": "Failed to re-authenticate with Wazuh Manager",
431 + }
432 +
433 + response = requests.get(
434 + f"{attributes['connector_url']}{endpoint}",
435 + headers=wazuh_manager_client,
436 + params=params,
437 + verify=False,
438 + )
439 +
440 response.raise_for_status()
441 return {
442 "data": response.json(),
backend/app/schedulers/scheduler.py
+1
@@ -215,6 +215,7 @@ async def schedule_enabled_jobs(scheduler):
215 "wazuh_index_fields_resize",
216 "invoke_huntress_integration_collection",
217 "invoke_cato_integration_collect",
218 + "invoke_sigma_queries_collect",
219 ]
220
221 # Disable each job in the list