| 1 | from datetime import datetime |
| 2 | from datetime import timedelta |
| 3 | from typing import Optional |
| 4 | from typing import Type |
| 5 | |
| 6 | from fastapi import HTTPException |
| 7 | from loguru import logger |
| 8 | |
| 9 | from app.connectors.wazuh_indexer.utils.universal import LogsQueryBuilder |
| 10 | from app.connectors.wazuh_indexer.utils.universal import collect_indices |
| 11 | from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client |
| 12 | from app.healthchecks.agents.schema.agents import AgentHealthCheckResponse |
| 13 | from app.healthchecks.agents.schema.agents import AgentModel |
| 14 | from app.healthchecks.agents.schema.agents import CollectLogsResponse |
| 15 | from app.healthchecks.agents.schema.agents import ExtendedAgentModel |
| 16 | from app.healthchecks.agents.schema.agents import HostLogsSearchBody |
| 17 | from app.healthchecks.agents.schema.agents import HostLogsSearchResponse |
| 18 | from app.healthchecks.agents.schema.agents import LogsSearchBody |
| 19 | from app.healthchecks.agents.schema.agents import TimeCriteriaModel |
| 20 | |
| 21 | |
| 22 | def is_wazuh_agent_unhealthy( |
| 23 | agent: AgentModel, |
| 24 | time_criteria: TimeCriteriaModel, |
| 25 | ) -> ExtendedAgentModel: |
| 26 | """ |
| 27 | Checks if a Wazuh agent is unhealthy based on the last seen time and time criteria. |
| 28 | |
| 29 | Args: |
| 30 | agent (AgentModel): The agent to check. |
| 31 | time_criteria (TimeCriteriaModel): The time criteria for determining agent health. |
| 32 | |
| 33 | Returns: |
| 34 | ExtendedAgentModel: An extended agent model with the unhealthy status updated. |
| 35 | """ |
| 36 | # If wazuh_last_seen is None, consider it unhealthy |
| 37 | if agent.wazuh_last_seen is None: |
| 38 | logger.info(f"Agent {agent.hostname} (ID: {agent.agent_id}) has no Wazuh last seen time - marking as unhealthy") |
| 39 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_wazuh_agent=True) |
| 40 | |
| 41 | current_time = datetime.now() |
| 42 | wazuh_last_seen = agent.wazuh_last_seen |
| 43 | |
| 44 | if wazuh_last_seen > current_time: |
| 45 | logger.info( |
| 46 | f"Agent {agent} has a wazuh_last_seen time in the future: {wazuh_last_seen}", |
| 47 | ) |
| 48 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_wazuh_agent=True) |
| 49 | |
| 50 | # Calculate the total time delta based on the criteria |
| 51 | total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 |
| 52 | time_delta = timedelta(minutes=total_minutes) |
| 53 | |
| 54 | is_unhealthy = (current_time - wazuh_last_seen) > time_delta |
| 55 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_wazuh_agent=is_unhealthy) |
| 56 | |
| 57 | |
| 58 | def is_velociraptor_agent_unhealthy( |
| 59 | agent: AgentModel, |
| 60 | time_criteria: TimeCriteriaModel, |
| 61 | ) -> ExtendedAgentModel: |
| 62 | """ |
| 63 | Checks if a velociraptor agent is unhealthy based on the last seen time and time criteria. |
| 64 | |
| 65 | Args: |
| 66 | agent (AgentModel): The agent to check. |
| 67 | time_criteria (TimeCriteriaModel): The time criteria for determining agent health. |
| 68 | |
| 69 | Returns: |
| 70 | ExtendedAgentModel: An extended agent model with the unhealthy_velociraptor_agent flag set. |
| 71 | """ |
| 72 | # If velociraptor_id is None or velociraptor_last_seen is None, consider it unhealthy |
| 73 | if agent.velociraptor_id is None or agent.velociraptor_last_seen is None: |
| 74 | logger.info(f"Agent {agent.hostname} (ID: {agent.agent_id}) has no Velociraptor ID or last seen time - marking as unhealthy") |
| 75 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_velociraptor_agent=True) |
| 76 | |
| 77 | current_time = datetime.now() |
| 78 | velociraptor_last_seen = agent.velociraptor_last_seen |
| 79 | |
| 80 | if velociraptor_last_seen > current_time: |
| 81 | logger.info( |
| 82 | f"Agent {agent} has a velociraptor_last_seen time in the future: {velociraptor_last_seen}", |
| 83 | ) |
| 84 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_velociraptor_agent=True) |
| 85 | |
| 86 | # Calculate the total time delta based on the criteria |
| 87 | total_minutes = time_criteria.minutes + time_criteria.hours * 60 + time_criteria.days * 24 * 60 |
| 88 | time_delta = timedelta(minutes=total_minutes) |
| 89 | |
| 90 | is_unhealthy = (current_time - velociraptor_last_seen) > time_delta |
| 91 | return ExtendedAgentModel(**agent.model_dump(), unhealthy_velociraptor_agent=is_unhealthy) |
| 92 | |
| 93 | |
| 94 | async def wazuh_agents_healthcheck( |
| 95 | agents: list, |
| 96 | time_criteria: TimeCriteriaModel, |
| 97 | ) -> AgentHealthCheckResponse: |
| 98 | """ |
| 99 | Perform a health check on Wazuh agents. |
| 100 | |
| 101 | Args: |
| 102 | agents (list): List of agents to perform health check on. |
| 103 | time_criteria (TimeCriteriaModel): Time criteria for determining agent health. |
| 104 | |
| 105 | Returns: |
| 106 | AgentHealthCheckResponse: Response object containing the results of the health check. |
| 107 | """ |
| 108 | healthy_wazuh_agents = [] |
| 109 | unhealthy_wazuh_agents = [] |
| 110 | for agent in agents: |
| 111 | # If agent_id is `000` skip it because this is the Wazuh manager |
| 112 | if agent.agent_id == "000": |
| 113 | continue |
| 114 | logger.info(f"Checking agent {agent} for health") |
| 115 | extended_agent = is_wazuh_agent_unhealthy(agent, time_criteria) |
| 116 | logger.info(f"Extended agent: {extended_agent}") |
| 117 | if extended_agent.unhealthy_wazuh_agent: |
| 118 | unhealthy_wazuh_agents.append(extended_agent) |
| 119 | else: |
| 120 | healthy_wazuh_agents.append(extended_agent) |
| 121 | |
| 122 | return AgentHealthCheckResponse( |
| 123 | healthy_wazuh_agents=healthy_wazuh_agents, |
| 124 | unhealthy_wazuh_agents=unhealthy_wazuh_agents, |
| 125 | success=True, |
| 126 | message="Wazuh agent healthcheck fetched successfully", |
| 127 | ) |
| 128 | |
| 129 | |
| 130 | async def wazuh_agent_healthcheck( |
| 131 | agent: AgentModel, |
| 132 | time_criteria: TimeCriteriaModel, |
| 133 | ) -> AgentHealthCheckResponse: |
| 134 | """ |
| 135 | Performs a health check on a Wazuh agent. |
| 136 | |
| 137 | Args: |
| 138 | agent (AgentModel): The agent to perform the health check on. |
| 139 | time_criteria (TimeCriteriaModel): The time criteria for the health check. |
| 140 | |
| 141 | Returns: |
| 142 | AgentHealthCheckResponse: The health check response containing the results of the health check. |
| 143 | """ |
| 144 | extended_agent = is_wazuh_agent_unhealthy(agent, time_criteria) |
| 145 | if extended_agent.unhealthy_wazuh_agent: |
| 146 | return AgentHealthCheckResponse( |
| 147 | healthy_wazuh_agents=[], |
| 148 | unhealthy_wazuh_agents=[extended_agent], |
| 149 | success=True, |
| 150 | message="Wazuh agent healthcheck fetched successfully", |
| 151 | ) |
| 152 | else: |
| 153 | return AgentHealthCheckResponse( |
| 154 | healthy_wazuh_agents=[extended_agent], |
| 155 | unhealthy_wazuh_agents=[], |
| 156 | success=True, |
| 157 | message="Wazuh agent healthcheck fetched successfully", |
| 158 | ) |
| 159 | |
| 160 | |
| 161 | async def velociraptor_agents_healthcheck( |
| 162 | agents: list, |
| 163 | time_criteria: TimeCriteriaModel, |
| 164 | ) -> AgentHealthCheckResponse: |
| 165 | """ |
| 166 | Perform health check on Velociraptor agents. |
| 167 | |
| 168 | Args: |
| 169 | agents (list): List of agents to perform health check on. |
| 170 | time_criteria (TimeCriteriaModel): Time criteria for determining agent health. |
| 171 | |
| 172 | Returns: |
| 173 | AgentHealthCheckResponse: Response object containing healthy and unhealthy agents. |
| 174 | """ |
| 175 | healthy_velociraptor_agents = [] |
| 176 | unhealthy_velociraptor_agents = [] |
| 177 | |
| 178 | for agent in agents: |
| 179 | # If agent_id is `000` skip it because this is the Wazuh manager |
| 180 | if agent.agent_id == "000": |
| 181 | continue |
| 182 | logger.info(f"Checking agent {agent} for health") |
| 183 | extended_agent = is_velociraptor_agent_unhealthy(agent, time_criteria) |
| 184 | logger.info(f"Extended agent: {extended_agent}") |
| 185 | if extended_agent.unhealthy_velociraptor_agent: |
| 186 | unhealthy_velociraptor_agents.append(extended_agent) |
| 187 | else: |
| 188 | healthy_velociraptor_agents.append(extended_agent) |
| 189 | |
| 190 | return AgentHealthCheckResponse( |
| 191 | healthy_velociraptor_agents=healthy_velociraptor_agents, |
| 192 | unhealthy_velociraptor_agents=unhealthy_velociraptor_agents, |
| 193 | success=True, |
| 194 | message="Velociraptor agent healthcheck fetched successfully", |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | async def velociraptor_agent_healthcheck( |
| 199 | agent: AgentModel, |
| 200 | time_criteria: TimeCriteriaModel, |
| 201 | ) -> AgentHealthCheckResponse: |
| 202 | """ |
| 203 | Perform a health check on a Velociraptor agent. |
| 204 | |
| 205 | Args: |
| 206 | agent (AgentModel): The agent to perform the health check on. |
| 207 | time_criteria (TimeCriteriaModel): The time criteria for the health check. |
| 208 | |
| 209 | Returns: |
| 210 | AgentHealthCheckResponse: The health check response containing the status of the Velociraptor agent. |
| 211 | """ |
| 212 | extended_agent = is_velociraptor_agent_unhealthy(agent, time_criteria) |
| 213 | if extended_agent.unhealthy_velociraptor_agent: |
| 214 | return AgentHealthCheckResponse( |
| 215 | healthy_velociraptor_agents=[], |
| 216 | unhealthy_velociraptor_agents=[extended_agent], |
| 217 | success=True, |
| 218 | message="Velociraptor agent healthcheck fetched successfully", |
| 219 | ) |
| 220 | else: |
| 221 | return AgentHealthCheckResponse( |
| 222 | healthy_velociraptor_agents=[extended_agent], |
| 223 | unhealthy_velociraptor_agents=[], |
| 224 | success=True, |
| 225 | message="Velociraptor agent healthcheck fetched successfully", |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | async def host_logs(search_body: HostLogsSearchBody) -> HostLogsSearchResponse: |
| 230 | """ |
| 231 | Search for host logs based on the provided search criteria. |
| 232 | |
| 233 | Args: |
| 234 | search_body (HostLogsSearchBody): The search criteria for host logs. |
| 235 | |
| 236 | Returns: |
| 237 | HostLogsSearchResponse: The response containing the search result and status information. |
| 238 | """ |
| 239 | result = await get_logs_generic(search_body, is_host_specific=True) |
| 240 | logger.info(f"Host logs search result: {result}") |
| 241 | |
| 242 | # Initialize variable to keep track of total logs |
| 243 | total_logs = 0 |
| 244 | |
| 245 | # Loop through each item in logs_summary to count total logs |
| 246 | for log_summary in result["logs_summary"]: |
| 247 | total_logs += log_summary["total_logs"] |
| 248 | |
| 249 | # Check if there are any logs |
| 250 | if total_logs > 0: |
| 251 | return HostLogsSearchResponse( |
| 252 | success=True, |
| 253 | healthy=True, |
| 254 | message=f"Host is healthy. At least one log was found within the specified time range of {search_body.timerange}", |
| 255 | ) |
| 256 | else: |
| 257 | return HostLogsSearchResponse( |
| 258 | success=True, |
| 259 | healthy=False, |
| 260 | message=f"Host is unhealthy. No logs were found within the specified time range of {search_body.timerange}", |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | async def get_logs_generic( |
| 265 | search_body: Type[LogsSearchBody], |
| 266 | is_host_specific: bool = False, |
| 267 | index_name: Optional[str] = None, |
| 268 | ): |
| 269 | """ |
| 270 | Retrieves logs based on the provided search criteria. |
| 271 | |
| 272 | Args: |
| 273 | search_body (Type[LogsSearchBody]): The search criteria for the logs. |
| 274 | is_host_specific (bool, optional): Specifies if the search is host-specific. Defaults to False. |
| 275 | index_name (str, optional): The name of the index to search in. If not provided, all indices will be searched. |
| 276 | |
| 277 | Returns: |
| 278 | dict: A dictionary containing the logs summary, success status, and message. |
| 279 | """ |
| 280 | logger.info( |
| 281 | f"Collecting Wazuh Indexer alerts for host {search_body.agent_name if is_host_specific else ''}", |
| 282 | ) |
| 283 | logs_summary = [] |
| 284 | indices = await collect_indices() |
| 285 | index_list = [index_name] if index_name else indices.indices_list # Use the provided index_name or get all indices |
| 286 | |
| 287 | for index_name in index_list: |
| 288 | try: |
| 289 | logs = await collect_logs_generic( |
| 290 | index_name, |
| 291 | body=search_body, |
| 292 | is_host_specific=is_host_specific, |
| 293 | ) |
| 294 | if logs.success and len(logs.logs) > 0: |
| 295 | logs_summary.append( |
| 296 | { |
| 297 | "index_name": index_name, |
| 298 | "total_logs": len(logs.logs), |
| 299 | "logs": logs.logs, |
| 300 | }, |
| 301 | ) |
| 302 | break # Only collect logs from the first index that has logs |
| 303 | except HTTPException as e: |
| 304 | logger.warning( |
| 305 | f"An error occurred while processing index {index_name}: {e.detail}", |
| 306 | ) |
| 307 | |
| 308 | if len(logs_summary) == 0: |
| 309 | message = "No logs found" |
| 310 | else: |
| 311 | message = f"Succesfully collected top {search_body.size} logs for each index" |
| 312 | |
| 313 | return { |
| 314 | "logs_summary": logs_summary, |
| 315 | "success": len(logs_summary) > 0, |
| 316 | "message": message, |
| 317 | } |
| 318 | |
| 319 | |
| 320 | async def collect_logs_generic( |
| 321 | index_name: str, |
| 322 | body: LogsSearchBody, |
| 323 | is_host_specific: bool = False, |
| 324 | ) -> CollectLogsResponse: |
| 325 | """ |
| 326 | Collects logs from Elasticsearch based on the specified parameters. |
| 327 | |
| 328 | Args: |
| 329 | index_name (str): The name of the Elasticsearch index to search. |
| 330 | body (LogsSearchBody): The search body containing the timerange, log field, log value, and other parameters. |
| 331 | is_host_specific (bool, optional): Specifies whether the search should be limited to a specific agent. Defaults to False. |
| 332 | |
| 333 | Returns: |
| 334 | CollectLogsResponse: The response object containing the collected logs, success status, and message. |
| 335 | """ |
| 336 | es_client = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 337 | query_builder = LogsQueryBuilder() |
| 338 | query_builder.add_time_range( |
| 339 | timerange=body.timerange, |
| 340 | timestamp_field=body.timestamp_field, |
| 341 | ) |
| 342 | query_builder.add_matches(matches=[(body.log_field, body.log_value)]) |
| 343 | query_builder.add_sort(body.timestamp_field) |
| 344 | |
| 345 | if is_host_specific: |
| 346 | query_builder.add_match_phrase(matches=[("agent_name", body.agent_name)]) |
| 347 | |
| 348 | query = query_builder.build() |
| 349 | |
| 350 | try: |
| 351 | logs = es_client.search(index=index_name, body=query, size=body.size) |
| 352 | logger.info(f"logs collected: {logs}") |
| 353 | logs_list = [log for log in logs["hits"]["hits"]] |
| 354 | logger.info(f"logs collected: {logs_list}") |
| 355 | return CollectLogsResponse( |
| 356 | logs=logs_list, |
| 357 | success=True, |
| 358 | message="logs collected successfully", |
| 359 | ) |
| 360 | except Exception as e: |
| 361 | logger.debug(f"Failed to collect logs: {e}") |
| 362 | return CollectLogsResponse( |
| 363 | logs=[], |
| 364 | success=False, |
| 365 | message=f"Failed to collect logs: {e}", |
| 366 | ) |