6
from asyncio import Semaphore
7
from datetime import datetime
8
from typing import Any
9
+from typing import AsyncGenerator
10
from typing import Dict
11
from typing import List
12
from typing import Optional
67
raise HTTPException(status_code=500, detail=f"Failed to fetch agents: {e}")
68
69
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
-
70
async def collect_sca_for_single_agent(
71
agent: Agents,
72
semaphore: Semaphore,
993
except Exception as e:
994
logger.error(f"Error deleting SCA report: {e}")
995
raise HTTPException(status_code=500, detail=f"Failed to delete report: {e}")
996
+
997
+
998
+async def stream_sca_for_all_agents(
999
+ db_session: AsyncSession,
1000
+ customer_code: Optional[str] = None,
1001
+ agent_name: Optional[str] = None,
1002
+ policy_id: Optional[str] = None,
1003
+ policy_name: Optional[str] = None,
1004
+ min_score: Optional[int] = None,
1005
+ max_score: Optional[int] = None,
1006
+ max_concurrent_requests: int = DEFAULT_MAX_CONCURRENT_REQUESTS,
1007
+) -> AsyncGenerator[Dict[str, Any], None]:
1008
+ """
1009
+ Stream SCA results for all agents as they are collected.
1010
+
1011
+ Yields SSE-formatted events as results come in from each agent.
1012
+
1013
+ Args:
1014
+ db_session: Database session to use
1015
+ customer_code: Optional customer code filter
1016
+ agent_name: Optional agent name filter
1017
+ policy_id: Optional policy ID filter
1018
+ policy_name: Optional policy name filter (partial matching)
1019
+ min_score: Optional minimum score filter
1020
+ max_score: Optional maximum score filter
1021
+ max_concurrent_requests: Maximum concurrent API requests
1022
+
1023
+ Yields:
1024
+ Dict with 'event' type and 'data' payload
1025
+ """
1026
+ try:
1027
+ # Get agents from database
1028
+ agents = await get_all_agents_from_db(db_session, customer_code)
1029
+
1030
+ if not agents:
1031
+ yield {
1032
+ "event": "complete",
1033
+ "data": {
1034
+ "total_results": 0,
1035
+ "total_agents": 0,
1036
+ "message": "No agents found",
1037
+ },
1038
+ }
1039
+ return
1040
+
1041
+ # Filter agents by name if specified
1042
+ if agent_name:
1043
+ agents = [a for a in agents if a.hostname == agent_name]
1044
+ if not agents:
1045
+ yield {
1046
+ "event": "complete",
1047
+ "data": {
1048
+ "total_results": 0,
1049
+ "total_agents": 0,
1050
+ "message": f"No agents found matching hostname: {agent_name}",
1051
+ },
1052
+ }
1053
+ return
1054
+
1055
+ total_agents = len(agents)
1056
+
1057
+ # Send start event
1058
+ yield {
1059
+ "event": "start",
1060
+ "data": {
1061
+ "total_agents": total_agents,
1062
+ "message": f"Starting SCA collection for {total_agents} agents...",
1063
+ },
1064
+ }
1065
+
1066
+ # Create semaphore for rate limiting
1067
+ semaphore = Semaphore(max_concurrent_requests)
1068
+
1069
+ # Track statistics
1070
+ all_results: List[AgentScaOverviewItem] = []
1071
+ processed_count = 0
1072
+ successful_count = 0
1073
+ failed_count = 0
1074
+
1075
+ # Create a queue to receive results as they complete
1076
+ result_queue: asyncio.Queue = asyncio.Queue()
1077
+
1078
+ async def collect_and_queue(agent: Agents):
1079
+ """Collect SCA for an agent and put result in queue"""
1080
+ result = await collect_sca_for_single_agent(
1081
+ agent=agent,
1082
+ semaphore=semaphore,
1083
+ policy_id=policy_id,
1084
+ policy_name=policy_name,
1085
+ min_score=min_score,
1086
+ max_score=max_score,
1087
+ )
1088
+ await result_queue.put((agent, result))
1089
+
1090
+ # Start all tasks
1091
+ tasks = [asyncio.create_task(collect_and_queue(agent)) for agent in agents]
1092
+
1093
+ # Process results as they come in
1094
+ for _ in range(total_agents):
1095
+ try:
1096
+ # Wait for next result with timeout
1097
+ agent, results = await asyncio.wait_for(result_queue.get(), timeout=60.0) # 60 second timeout per agent
1098
+
1099
+ processed_count += 1
1100
+
1101
+ if results:
1102
+ successful_count += 1
1103
+ all_results.extend(results)
1104
+
1105
+ # Yield agent results
1106
+ yield {
1107
+ "event": "agent_result",
1108
+ "data": {
1109
+ "agent_id": agent.agent_id,
1110
+ "agent_name": agent.hostname,
1111
+ "customer_code": agent.customer_code,
1112
+ "policy_count": len(results),
1113
+ "policies": [
1114
+ {
1115
+ "policy_id": r.policy_id,
1116
+ "policy_name": r.policy_name,
1117
+ "description": r.description,
1118
+ "total_checks": r.total_checks,
1119
+ "pass_count": r.pass_count,
1120
+ "fail_count": r.fail_count,
1121
+ "invalid_count": r.invalid_count,
1122
+ "score": r.score,
1123
+ "start_scan": r.start_scan,
1124
+ "end_scan": r.end_scan,
1125
+ "references": r.references,
1126
+ "hash_file": r.hash_file,
1127
+ }
1128
+ for r in results
1129
+ ],
1130
+ },
1131
+ }
1132
+ else:
1133
+ # Agent had no SCA data (not necessarily an error)
1134
+ yield {
1135
+ "event": "agent_empty",
1136
+ "data": {
1137
+ "agent_id": agent.agent_id,
1138
+ "agent_name": agent.hostname,
1139
+ "message": "No SCA data available",
1140
+ },
1141
+ }
1142
+
1143
+ # Yield progress update every 5 agents or on last agent
1144
+ if processed_count % 5 == 0 or processed_count == total_agents:
1145
+ yield {
1146
+ "event": "progress",
1147
+ "data": {
1148
+ "processed": processed_count,
1149
+ "total": total_agents,
1150
+ "successful": successful_count,
1151
+ "failed": failed_count,
1152
+ "results_so_far": len(all_results),
1153
+ "percent_complete": round((processed_count / total_agents) * 100, 1),
1154
+ },
1155
+ }
1156
+
1157
+ except asyncio.TimeoutError:
1158
+ failed_count += 1
1159
+ processed_count += 1
1160
+ yield {
1161
+ "event": "agent_error",
1162
+ "data": {
1163
+ "agent_id": "unknown",
1164
+ "message": "Timeout waiting for agent response",
1165
+ },
1166
+ }
1167
+ except Exception as e:
1168
+ failed_count += 1
1169
+ processed_count += 1
1170
+ logger.error(f"Error processing agent result: {e}")
1171
+ yield {
1172
+ "event": "agent_error",
1173
+ "data": {
1174
+ "message": str(e),
1175
+ },
1176
+ }
1177
+
1178
+ # Wait for all tasks to complete (cleanup)
1179
+ await asyncio.gather(*tasks, return_exceptions=True)
1180
+
1181
+ # Calculate final statistics
1182
+ unique_agents = set(item.agent_id for item in all_results)
1183
+ unique_policies = set(item.policy_id for item in all_results)
1184
+
1185
+ total_checks = sum(item.total_checks for item in all_results)
1186
+ total_passes = sum(item.pass_count for item in all_results)
1187
+ total_fails = sum(item.fail_count for item in all_results)
1188
+ total_invalid = sum(item.invalid_count for item in all_results)
1189
+
1190
+ average_score = sum(item.score for item in all_results) / len(all_results) if all_results else 0.0
1191
+
1192
+ # Yield completion event
1193
+ yield {
1194
+ "event": "complete",
1195
+ "data": {
1196
+ "total_results": len(all_results),
1197
+ "total_agents": len(unique_agents),
1198
+ "total_policies": len(unique_policies),
1199
+ "average_score": round(average_score, 2),
1200
+ "total_checks": total_checks,
1201
+ "total_passes": total_passes,
1202
+ "total_fails": total_fails,
1203
+ "total_invalid": total_invalid,
1204
+ "agents_processed": processed_count,
1205
+ "agents_successful": successful_count,
1206
+ "agents_failed": failed_count,
1207
+ "message": f"Completed SCA collection: {len(all_results)} results from {len(unique_agents)} agents",
1208
+ },
1209
+ }
1210
+
1211
+ except Exception as e:
1212
+ logger.error(f"Error in SCA streaming: {e}")
1213
+ yield {
1214
+ "event": "error",
1215
+ "data": {
1216
+ "error": str(e),
1217
+ "message": "Fatal error during SCA collection",
1218
+ },
1219
+ }