| 1 | import asyncio |
| 2 | import csv |
| 3 | import io |
| 4 | import os |
| 5 | from datetime import datetime |
| 6 | from datetime import timedelta |
| 7 | from typing import List |
| 8 | from typing import Optional |
| 9 | |
| 10 | # from fastapi import BackgroundTasks |
| 11 | from fastapi import APIRouter |
| 12 | from fastapi import BackgroundTasks |
| 13 | from fastapi import Depends |
| 14 | from fastapi import Header |
| 15 | from fastapi import HTTPException |
| 16 | from fastapi import Path |
| 17 | from fastapi import Query |
| 18 | from fastapi import Security |
| 19 | from fastapi.responses import StreamingResponse |
| 20 | from loguru import logger |
| 21 | from packaging import version |
| 22 | from sqlalchemy import delete |
| 23 | from sqlalchemy.ext.asyncio import AsyncSession |
| 24 | from sqlalchemy.future import select |
| 25 | |
| 26 | from app.agents.schema.agents import AgentModifyResponse |
| 27 | from app.agents.schema.agents import AgentsResponse |
| 28 | from app.agents.schema.agents import AgentWazuhUpgradeResponse |
| 29 | from app.agents.schema.agents import BulkDeleteAgentRequest |
| 30 | from app.agents.schema.agents import BulkDeleteAgentResult |
| 31 | from app.agents.schema.agents import BulkDeleteAgentsResponse |
| 32 | from app.agents.schema.agents import BulkDeleteFilterRequest |
| 33 | from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse |
| 34 | from app.agents.schema.agents import OutdatedWazuhAgentsResponse |
| 35 | from app.agents.schema.agents import SyncedAgentsResponse |
| 36 | from app.agents.services.status import get_agents_by_customer_code |
| 37 | from app.agents.services.status import get_outdated_agents_velociraptor |
| 38 | from app.agents.services.status import get_outdated_agents_wazuh |
| 39 | from app.agents.services.sync import sync_agents_velociraptor |
| 40 | from app.agents.services.sync import sync_agents_wazuh |
| 41 | from app.agents.velociraptor.services.agents import delete_agent_velociraptor |
| 42 | from app.agents.wazuh.schema.agents import VulnSeverity |
| 43 | from app.agents.wazuh.schema.agents import WazuhAgentScaPolicyResultsResponse |
| 44 | from app.agents.wazuh.schema.agents import WazuhAgentScaResponse |
| 45 | from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse |
| 46 | from app.agents.wazuh.services.agents import delete_agent_wazuh |
| 47 | from app.agents.wazuh.services.agents import upgrade_wazuh_agent |
| 48 | from app.agents.wazuh.services.sca import collect_agent_sca |
| 49 | from app.agents.wazuh.services.sca import collect_agent_sca_policy_results |
| 50 | from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilities |
| 51 | from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilities_new |
| 52 | from app.agents.wazuh.services.vulnerabilities import sync_agent_vulnerabilities |
| 53 | from app.auth.models.users import User |
| 54 | |
| 55 | # App specific imports |
| 56 | from app.auth.routes.auth import AuthHandler |
| 57 | from app.connectors.wazuh_manager.utils.universal import send_get_request |
| 58 | from app.db.db_session import get_db |
| 59 | |
| 60 | # App specific imports |
| 61 | # from app.db.db_session import session |
| 62 | from app.db.universal_models import AgentDataStore |
| 63 | from app.db.universal_models import Agents |
| 64 | from app.incidents.schema.db_operations import CaseOutResponse |
| 65 | from app.incidents.services.db_operations import list_cases_by_asset_name |
| 66 | from app.middleware.customer_access import customer_access_handler |
| 67 | from app.threat_intel.schema.epss import EpssThreatIntelRequest |
| 68 | from app.threat_intel.services.epss import collect_epss_score |
| 69 | |
| 70 | |
| 71 | async def get_wazuh_manager_version() -> str: |
| 72 | """ |
| 73 | Fetches the version of the Wazuh Manager. |
| 74 | |
| 75 | Returns: |
| 76 | str: The version of the Wazuh Manager. |
| 77 | |
| 78 | Raises: |
| 79 | HTTPException: If there is an error fetching the version of the Wazuh Manager. |
| 80 | """ |
| 81 | try: |
| 82 | response = await send_get_request(endpoint="/") |
| 83 | logger.info(f"Fetched Wazuh Manager version: {response}") |
| 84 | return response["data"]["data"]["api_version"] |
| 85 | except Exception as e: |
| 86 | logger.error(f"Failed to fetch Wazuh Manager version: {e}") |
| 87 | raise HTTPException( |
| 88 | status_code=500, |
| 89 | detail=f"Failed to fetch Wazuh Manager version: {e}", |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | async def check_wazuh_manager_version() -> bool: |
| 94 | """ |
| 95 | Checks the version of the Wazuh Manager. |
| 96 | |
| 97 | Returns: |
| 98 | bool: True if the version of the Wazuh Manager is 4.8.0 or higher, False otherwise. |
| 99 | """ |
| 100 | try: |
| 101 | wazuh_manager_version = await get_wazuh_manager_version() |
| 102 | return version.parse(wazuh_manager_version) >= version.parse("4.8.0") |
| 103 | except Exception as e: |
| 104 | logger.error(f"Failed to check Wazuh Manager version: {e}") |
| 105 | return False |
| 106 | |
| 107 | |
| 108 | async def delete_single_agent( |
| 109 | db: AsyncSession, |
| 110 | agent_id: str, |
| 111 | agent: Agents, |
| 112 | ) -> BulkDeleteAgentResult: |
| 113 | """ |
| 114 | Delete a single agent and return the result. |
| 115 | This function handles errors gracefully and returns a result object. |
| 116 | |
| 117 | Args: |
| 118 | db (AsyncSession): The database session. |
| 119 | agent_id (str): The ID of the agent to delete. |
| 120 | agent (Agents): The agent object from the database. |
| 121 | |
| 122 | Returns: |
| 123 | BulkDeleteAgentResult: The result of the deletion attempt. |
| 124 | """ |
| 125 | errors = [] |
| 126 | |
| 127 | # Try to delete from Wazuh |
| 128 | try: |
| 129 | await delete_agent_wazuh(agent_id) |
| 130 | except Exception as e: |
| 131 | error_msg = f"Failed to delete from Wazuh: {str(e)}" |
| 132 | logger.warning(f"Agent {agent_id}: {error_msg}") |
| 133 | errors.append(error_msg) |
| 134 | |
| 135 | # Try to delete from Velociraptor if applicable |
| 136 | if agent.velociraptor_id and agent.velociraptor_id != "Unknown": |
| 137 | try: |
| 138 | await delete_agent_velociraptor(agent.velociraptor_id) |
| 139 | except Exception as e: |
| 140 | error_msg = f"Failed to delete from Velociraptor: {str(e)}" |
| 141 | logger.warning(f"Agent {agent_id}: {error_msg}") |
| 142 | errors.append(error_msg) |
| 143 | |
| 144 | # Delete from database |
| 145 | try: |
| 146 | await delete_agent_from_database(db=db, agent_id=agent_id) |
| 147 | except Exception as e: |
| 148 | error_msg = f"Failed to delete from database: {str(e)}" |
| 149 | logger.error(f"Agent {agent_id}: {error_msg}") |
| 150 | errors.append(error_msg) |
| 151 | return BulkDeleteAgentResult( |
| 152 | agent_id=agent_id, |
| 153 | success=False, |
| 154 | message=f"Failed to delete agent: {'; '.join(errors)}", |
| 155 | ) |
| 156 | |
| 157 | if errors: |
| 158 | return BulkDeleteAgentResult( |
| 159 | agent_id=agent_id, |
| 160 | success=True, |
| 161 | message=f"Agent deleted from database but with warnings: {'; '.join(errors)}", |
| 162 | ) |
| 163 | |
| 164 | return BulkDeleteAgentResult( |
| 165 | agent_id=agent_id, |
| 166 | success=True, |
| 167 | message="Agent deleted successfully", |
| 168 | ) |
| 169 | |
| 170 | |
| 171 | agents_router = APIRouter() |
| 172 | |
| 173 | |
| 174 | async def fetch_velociraptor_id(db: AsyncSession, agent_id: str) -> str: |
| 175 | """ |
| 176 | Fetches the velociraptor ID of an agent from the database. |
| 177 | |
| 178 | Args: |
| 179 | db (AsyncSession): The database session. |
| 180 | agent_id (str): The ID of the agent. |
| 181 | |
| 182 | Returns: |
| 183 | str: The velociraptor ID of the agent. |
| 184 | |
| 185 | Raises: |
| 186 | HTTPException: If the agent is not found in the database. |
| 187 | HTTPException: If there is an error fetching the agent from the database. |
| 188 | """ |
| 189 | try: |
| 190 | result = await db.execute(select(Agents).filter(Agents.agent_id == agent_id)) |
| 191 | agent = result.scalars().first() |
| 192 | if agent: |
| 193 | return agent.velociraptor_id |
| 194 | else: |
| 195 | raise HTTPException(status_code=404, detail=f"Agent {agent_id} not found") |
| 196 | except Exception as e: |
| 197 | logger.error(f"Failed to fetch agent {agent_id} from database: {e}") |
| 198 | raise HTTPException( |
| 199 | status_code=500, |
| 200 | detail=f"Failed to fetch agent {agent_id} from database: {e}", |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | async def delete_agent_from_database(db: AsyncSession, agent_id: str): |
| 205 | """ |
| 206 | Delete an agent from the database. |
| 207 | |
| 208 | Args: |
| 209 | db (AsyncSession): The async database session. |
| 210 | agent_id (str): The ID of the agent to be deleted. |
| 211 | |
| 212 | Raises: |
| 213 | HTTPException: If there is an error deleting the agent from the database. |
| 214 | |
| 215 | """ |
| 216 | try: |
| 217 | # First delete related records from agent_datastore |
| 218 | await db.execute(delete(AgentDataStore).filter(AgentDataStore.agent_id == agent_id)) |
| 219 | # Then delete the agent |
| 220 | await db.execute(delete(Agents).filter(Agents.agent_id == agent_id)) |
| 221 | await db.commit() |
| 222 | except Exception as e: |
| 223 | logger.error(f"Failed to delete agent {agent_id} from database: {e}") |
| 224 | await db.rollback() |
| 225 | raise HTTPException( |
| 226 | status_code=500, |
| 227 | detail=f"Failed to delete agent {agent_id} from database: {e}", |
| 228 | ) |
| 229 | |
| 230 | |
| 231 | @agents_router.get( |
| 232 | "", |
| 233 | response_model=AgentsResponse, |
| 234 | description="Get all agents currently synced to the database", |
| 235 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 236 | ) |
| 237 | async def get_agents( |
| 238 | customer_codes: Optional[List[str]] = Query(None, description="Optional subset of customer codes to scope the results to"), |
| 239 | current_user: User = Depends(AuthHandler().get_current_user), |
| 240 | db: AsyncSession = Depends(get_db), |
| 241 | ) -> AgentsResponse: |
| 242 | """ |
| 243 | Retrieve all agents currently synced to the database. |
| 244 | Results are filtered based on user's customer access permissions. |
| 245 | |
| 246 | Returns: |
| 247 | AgentsResponse: The response containing the list of agents, success status, and message. |
| 248 | |
| 249 | Raises: |
| 250 | HTTPException: If there is an error while fetching the agents. |
| 251 | """ |
| 252 | logger.info("Fetching all agents") |
| 253 | try: |
| 254 | # Apply customer access filtering (optionally narrowed to a requested subset) |
| 255 | base_query = select(Agents) |
| 256 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 257 | current_user, |
| 258 | db, |
| 259 | base_query, |
| 260 | Agents.customer_code, |
| 261 | requested_customers=customer_codes, |
| 262 | ) |
| 263 | |
| 264 | result = await db.execute(filtered_query) |
| 265 | agents = result.scalars().all() |
| 266 | return AgentsResponse( |
| 267 | agents=agents, |
| 268 | success=True, |
| 269 | message="Agents fetched successfully", |
| 270 | ) |
| 271 | except Exception as e: |
| 272 | logger.error(f"Failed to fetch agents: {e}") |
| 273 | raise HTTPException(status_code=500, detail=f"Failed to fetch agents: {e}") |
| 274 | |
| 275 | |
| 276 | # Function to validate the Grafana shared-secret header. |
| 277 | # Modeled on verify_graylog_header (app/active_response/routes/graylog.py), but |
| 278 | # fails closed: it does NOT fall back to a hardcoded default secret. A default |
| 279 | # baked into the open-source repo would let anyone reproduce it and bypass the |
| 280 | # check (the same class of bug as the JWT_SECRET default in GHSA-4gxj-hw3c-3x2x). |
| 281 | # GRAFANA_API_HEADER_VALUE must be set or the route is denied. See GHSA-xh98-w6qh-cr44. |
| 282 | async def verify_grafana_header(grafana: Optional[str] = Header(None)): |
| 283 | """Verify that a Grafana-invoked request carries the correct shared-secret header.""" |
| 284 | expected_header = os.getenv("GRAFANA_API_HEADER_VALUE") |
| 285 | if not expected_header: |
| 286 | logger.error("GRAFANA_API_HEADER_VALUE is not configured; denying Grafana dashboard request") |
| 287 | raise HTTPException(status_code=403, detail="Grafana header authentication is not configured") |
| 288 | if grafana != expected_header: |
| 289 | logger.error("Invalid or missing Grafana header") |
| 290 | raise HTTPException(status_code=403, detail="Invalid or missing Grafana header") |
| 291 | return grafana |
| 292 | |
| 293 | |
| 294 | @agents_router.get( |
| 295 | "/dashboard/agents", |
| 296 | response_model=AgentsResponse, |
| 297 | description="Get all Wazuh agents for a specific customer (Grafana dashboard use)", |
| 298 | dependencies=[Depends(verify_grafana_header)], |
| 299 | ) |
| 300 | async def get_customer_agents_for_dashboard( |
| 301 | customer_code: Optional[str] = Header(None, description="Customer code to filter agents by"), |
| 302 | db: AsyncSession = Depends(get_db), |
| 303 | ) -> AgentsResponse: |
| 304 | """ |
| 305 | Retrieve all agents for a specific customer for dashboard use. |
| 306 | This endpoint is designed specifically for integration with Grafana dashboards. |
| 307 | |
| 308 | Args: |
| 309 | customer_code (str, optional): The customer code from the request header. |
| 310 | db (AsyncSession): The database session. |
| 311 | |
| 312 | Returns: |
| 313 | AgentsResponse: The response containing the list of agents for the specified customer. |
| 314 | |
| 315 | Raises: |
| 316 | HTTPException: If the customer_code is not provided or if there's an error fetching the agents. |
| 317 | """ |
| 318 | if not customer_code: |
| 319 | logger.warning("Dashboard agent request made with no customer_code header") |
| 320 | return AgentsResponse( |
| 321 | agents=[], |
| 322 | success=False, |
| 323 | message="No customer_code header provided", |
| 324 | ) |
| 325 | |
| 326 | logger.info(f"Fetching agents for customer_code: {customer_code} (dashboard request)") |
| 327 | try: |
| 328 | # Query agents with the specified customer code |
| 329 | result = await db.execute(select(Agents).filter(Agents.customer_code == customer_code)) |
| 330 | agents = result.scalars().all() |
| 331 | |
| 332 | logger.info(f"Found {len(agents)} agents for customer_code: {customer_code}") |
| 333 | return AgentsResponse( |
| 334 | agents=agents, |
| 335 | success=True, |
| 336 | message=f"Agents for customer {customer_code} fetched successfully", |
| 337 | ) |
| 338 | except Exception as e: |
| 339 | logger.error(f"Failed to fetch agents for customer {customer_code}: {e}") |
| 340 | raise HTTPException(status_code=500, detail=f"Failed to fetch agents for customer {customer_code}") |
| 341 | |
| 342 | |
| 343 | @agents_router.post( |
| 344 | "/bulk/delete", |
| 345 | response_model=BulkDeleteAgentsResponse, |
| 346 | description="Delete multiple agents by their IDs", |
| 347 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 348 | ) |
| 349 | async def bulk_delete_agents( |
| 350 | request: BulkDeleteAgentRequest, |
| 351 | current_user: User = Depends(AuthHandler().get_current_user), |
| 352 | session: AsyncSession = Depends(get_db), |
| 353 | ) -> BulkDeleteAgentsResponse: |
| 354 | """ |
| 355 | Delete multiple agents by their IDs. |
| 356 | User must have access to each agent's customer. |
| 357 | If an agent fails to delete, the process continues with the remaining agents. |
| 358 | |
| 359 | Args: |
| 360 | request (BulkDeleteAgentRequest): The request containing agent IDs to delete. |
| 361 | current_user (User): The authenticated user. |
| 362 | session (AsyncSession): The database session. |
| 363 | |
| 364 | Returns: |
| 365 | BulkDeleteAgentsResponse: The response containing results for each deletion attempt. |
| 366 | """ |
| 367 | logger.info(f"Bulk deleting {len(request.agent_ids)} agents") |
| 368 | |
| 369 | results: List[BulkDeleteAgentResult] = [] |
| 370 | successful_count = 0 |
| 371 | failed_count = 0 |
| 372 | |
| 373 | for agent_id in request.agent_ids: |
| 374 | # Check customer access for each agent |
| 375 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 376 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 377 | current_user, |
| 378 | session, |
| 379 | base_query, |
| 380 | Agents.customer_code, |
| 381 | ) |
| 382 | |
| 383 | result = await session.execute(filtered_query) |
| 384 | agent = result.scalars().first() |
| 385 | |
| 386 | if not agent: |
| 387 | results.append( |
| 388 | BulkDeleteAgentResult( |
| 389 | agent_id=agent_id, |
| 390 | success=False, |
| 391 | message="Agent not found or access denied", |
| 392 | ), |
| 393 | ) |
| 394 | failed_count += 1 |
| 395 | continue |
| 396 | |
| 397 | # Delete the agent |
| 398 | delete_result = await delete_single_agent(db=session, agent_id=agent_id, agent=agent) |
| 399 | results.append(delete_result) |
| 400 | |
| 401 | if delete_result.success: |
| 402 | successful_count += 1 |
| 403 | else: |
| 404 | failed_count += 1 |
| 405 | |
| 406 | return BulkDeleteAgentsResponse( |
| 407 | success=failed_count == 0, |
| 408 | message=f"Bulk deletion completed: {successful_count} successful, {failed_count} failed", |
| 409 | total_requested=len(request.agent_ids), |
| 410 | successful_deletions=successful_count, |
| 411 | failed_deletions=failed_count, |
| 412 | results=results, |
| 413 | ) |
| 414 | |
| 415 | |
| 416 | @agents_router.post( |
| 417 | "/bulk/delete/filter", |
| 418 | response_model=BulkDeleteAgentsResponse, |
| 419 | description="Delete multiple agents based on filter conditions", |
| 420 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 421 | ) |
| 422 | async def bulk_delete_agents_by_filter( |
| 423 | request: BulkDeleteFilterRequest, |
| 424 | current_user: User = Depends(AuthHandler().get_current_user), |
| 425 | session: AsyncSession = Depends(get_db), |
| 426 | ) -> BulkDeleteAgentsResponse: |
| 427 | """ |
| 428 | Delete multiple agents based on filter conditions. |
| 429 | User must have access to each agent's customer. |
| 430 | If an agent fails to delete, the process continues with the remaining agents. |
| 431 | |
| 432 | Filter conditions: |
| 433 | - customer_code: Filter by customer code |
| 434 | - status: Filter by agent status ('disconnected', 'never_connected', 'active') |
| 435 | - disconnected_days: Filter agents disconnected for more than X days |
| 436 | |
| 437 | Args: |
| 438 | request (BulkDeleteFilterRequest): The filter conditions. |
| 439 | current_user (User): The authenticated user. |
| 440 | session (AsyncSession): The database session. |
| 441 | |
| 442 | Returns: |
| 443 | BulkDeleteAgentsResponse: The response containing results for each deletion attempt. |
| 444 | """ |
| 445 | logger.info(f"Bulk deleting agents with filters: {request}") |
| 446 | |
| 447 | # Safety check: require at least one filter to be specified |
| 448 | if not request.customer_code and not request.status and not request.disconnected_days: |
| 449 | raise HTTPException( |
| 450 | status_code=400, |
| 451 | detail="At least one filter condition must be specified to prevent accidental bulk deletion of all agents", |
| 452 | ) |
| 453 | |
| 454 | # Build the base query with filters |
| 455 | base_query = select(Agents) |
| 456 | |
| 457 | if request.customer_code: |
| 458 | base_query = base_query.filter(Agents.customer_code == request.customer_code) |
| 459 | |
| 460 | if request.status: |
| 461 | if request.status.lower() == "disconnected": |
| 462 | base_query = base_query.filter(Agents.wazuh_agent_status == "disconnected") |
| 463 | elif request.status.lower() == "never_connected": |
| 464 | base_query = base_query.filter(Agents.wazuh_agent_status == "never_connected") |
| 465 | elif request.status.lower() == "active": |
| 466 | base_query = base_query.filter(Agents.wazuh_agent_status == "active") |
| 467 | |
| 468 | if request.disconnected_days: |
| 469 | cutoff_date = datetime.utcnow() - timedelta(days=request.disconnected_days) |
| 470 | # Filter agents whose last keep alive is older than the cutoff |
| 471 | base_query = base_query.filter(Agents.wazuh_last_seen < cutoff_date) |
| 472 | |
| 473 | # Apply customer access filtering |
| 474 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 475 | current_user, |
| 476 | session, |
| 477 | base_query, |
| 478 | Agents.customer_code, |
| 479 | ) |
| 480 | |
| 481 | result = await session.execute(filtered_query) |
| 482 | agents_to_delete = result.scalars().all() |
| 483 | |
| 484 | if not agents_to_delete: |
| 485 | return BulkDeleteAgentsResponse( |
| 486 | success=True, |
| 487 | message="No agents found matching the filter criteria", |
| 488 | total_requested=0, |
| 489 | successful_deletions=0, |
| 490 | failed_deletions=0, |
| 491 | results=[], |
| 492 | ) |
| 493 | |
| 494 | logger.info(f"Found {len(agents_to_delete)} agents matching filter criteria") |
| 495 | |
| 496 | results: List[BulkDeleteAgentResult] = [] |
| 497 | successful_count = 0 |
| 498 | failed_count = 0 |
| 499 | |
| 500 | for agent in agents_to_delete: |
| 501 | delete_result = await delete_single_agent(db=session, agent_id=agent.agent_id, agent=agent) |
| 502 | results.append(delete_result) |
| 503 | |
| 504 | if delete_result.success: |
| 505 | successful_count += 1 |
| 506 | else: |
| 507 | failed_count += 1 |
| 508 | |
| 509 | return BulkDeleteAgentsResponse( |
| 510 | success=failed_count == 0, |
| 511 | message=f"Bulk deletion completed: {successful_count} successful, {failed_count} failed", |
| 512 | total_requested=len(agents_to_delete), |
| 513 | successful_deletions=successful_count, |
| 514 | failed_deletions=failed_count, |
| 515 | results=results, |
| 516 | ) |
| 517 | |
| 518 | |
| 519 | @agents_router.get( |
| 520 | "/{agent_id}", |
| 521 | response_model=AgentsResponse, |
| 522 | description="Get agent by agent_id", |
| 523 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 524 | ) |
| 525 | async def get_agent( |
| 526 | agent_id: str, |
| 527 | current_user: User = Depends(AuthHandler().get_current_user), |
| 528 | db: AsyncSession = Depends(get_db), |
| 529 | ) -> AgentsResponse: |
| 530 | """ |
| 531 | Retrieve an agent by agent_id. |
| 532 | Results are filtered based on user's customer access permissions. |
| 533 | |
| 534 | Args: |
| 535 | agent_id (str): The ID of the agent to retrieve. |
| 536 | current_user (User): The authenticated user. |
| 537 | db (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 538 | |
| 539 | Returns: |
| 540 | AgentsResponse: The response containing the agent information. |
| 541 | |
| 542 | Raises: |
| 543 | HTTPException: If the agent with the specified agent_id is not found or if there is an error fetching the agent. |
| 544 | """ |
| 545 | logger.info(f"Fetching agent with agent_id: {agent_id}") |
| 546 | try: |
| 547 | # Apply customer access filtering |
| 548 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 549 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, db, base_query, Agents.customer_code) |
| 550 | |
| 551 | result = await db.execute(filtered_query) |
| 552 | agent = result.scalars().first() |
| 553 | if agent: |
| 554 | return AgentsResponse( |
| 555 | agents=[agent], |
| 556 | success=True, |
| 557 | message="Agent fetched successfully", |
| 558 | ) |
| 559 | else: |
| 560 | raise HTTPException( |
| 561 | status_code=404, |
| 562 | detail=f"Agent with agent_id {agent_id} not found or access denied", |
| 563 | ) |
| 564 | except Exception as e: |
| 565 | logger.error( |
| 566 | f"Failed to fetch agent: {agent_id} with error {e}. Does it exist?", |
| 567 | ) |
| 568 | raise HTTPException( |
| 569 | status_code=500, |
| 570 | detail=f"Failed to fetch agent: {agent_id}. Does it exist?", |
| 571 | ) |
| 572 | |
| 573 | |
| 574 | @agents_router.get( |
| 575 | "/hostname/{hostname}", |
| 576 | response_model=AgentsResponse, |
| 577 | description="Get agent by hostname", |
| 578 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 579 | ) |
| 580 | async def get_agent_by_hostname( |
| 581 | hostname: str, |
| 582 | current_user: User = Depends(AuthHandler().get_current_user), |
| 583 | db: AsyncSession = Depends(get_db), |
| 584 | ) -> AgentsResponse: |
| 585 | """ |
| 586 | Retrieve an agent by its hostname. |
| 587 | Results are filtered based on user's customer access permissions. |
| 588 | |
| 589 | Args: |
| 590 | hostname (str): The hostname of the agent. |
| 591 | current_user (User): The authenticated user. |
| 592 | db (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 593 | |
| 594 | Returns: |
| 595 | AgentsResponse: The response containing the agent information. |
| 596 | |
| 597 | Raises: |
| 598 | HTTPException: If the agent with the specified hostname is not found or if there is an error fetching the agent. |
| 599 | """ |
| 600 | logger.info(f"Fetching agent with hostname: {hostname}") |
| 601 | try: |
| 602 | # Apply customer access filtering |
| 603 | base_query = select(Agents).filter(Agents.hostname == hostname) |
| 604 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, db, base_query, Agents.customer_code) |
| 605 | |
| 606 | result = await db.execute(filtered_query) |
| 607 | agent = result.scalars().first() |
| 608 | if agent: |
| 609 | return AgentsResponse( |
| 610 | agents=[agent], |
| 611 | success=True, |
| 612 | message="Agent fetched successfully", |
| 613 | ) |
| 614 | else: |
| 615 | raise HTTPException( |
| 616 | status_code=404, |
| 617 | detail=f"Agent with hostname {hostname} not found or access denied", |
| 618 | ) |
| 619 | except Exception as e: |
| 620 | logger.error(f"Failed to fetch agent: {e}") |
| 621 | # The exception message should not be exposed directly, especially in production |
| 622 | raise HTTPException(status_code=500, detail="Failed to fetch agent") |
| 623 | |
| 624 | |
| 625 | @agents_router.post( |
| 626 | "/sync", |
| 627 | response_model=SyncedAgentsResponse, |
| 628 | description="Sync agents from Wazuh Manager", |
| 629 | dependencies=[ |
| 630 | Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler")), |
| 631 | ], |
| 632 | ) |
| 633 | async def sync_all_agents() -> SyncedAgentsResponse: |
| 634 | """ |
| 635 | Sync all agents from Wazuh Manager. |
| 636 | |
| 637 | This endpoint triggers the synchronization of all agents from the Wazuh Manager. |
| 638 | It requires authentication with any of the following scopes: "admin", "analyst", "scheduler". |
| 639 | |
| 640 | Parameters: |
| 641 | - backgroud_tasks (BackgroundTasks): The background tasks object used to add the sync_agents task. |
| 642 | - session (AsyncSession, optional): The async session object used to interact with the database. Defaults to Depends(get_db). |
| 643 | |
| 644 | Returns: |
| 645 | - SyncedAgentsResponse: The response model indicating the success of the sync operation. |
| 646 | |
| 647 | """ |
| 648 | logger.info("Syncing agents as part of scheduled job") |
| 649 | loop = asyncio.get_event_loop() |
| 650 | await loop.create_task(sync_agents_wazuh()) |
| 651 | await loop.create_task(sync_agents_velociraptor()) |
| 652 | return SyncedAgentsResponse( |
| 653 | success=True, |
| 654 | message="Agents synced started successfully", |
| 655 | ) |
| 656 | |
| 657 | |
| 658 | @agents_router.post( |
| 659 | "/{agent_id}/critical", |
| 660 | response_model=AgentModifyResponse, |
| 661 | description="Mark agent as critical", |
| 662 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 663 | ) |
| 664 | async def mark_agent_as_critical( |
| 665 | agent_id: str, |
| 666 | current_user: User = Depends(AuthHandler().get_current_user), |
| 667 | session: AsyncSession = Depends(get_db), |
| 668 | ) -> AgentModifyResponse: |
| 669 | """ |
| 670 | Marks the specified agent as critical. |
| 671 | User must have access to the agent's customer. |
| 672 | |
| 673 | Args: |
| 674 | agent_id (str): The ID of the agent to mark as critical. |
| 675 | current_user (User): The authenticated user. |
| 676 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 677 | |
| 678 | Returns: |
| 679 | AgentModifyResponse: The response indicating the success or failure of marking the agent as critical. |
| 680 | """ |
| 681 | logger.info(f"Marking agent {agent_id} as critical") |
| 682 | try: |
| 683 | # Check customer access - first find the agent |
| 684 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 685 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 686 | current_user, |
| 687 | session, |
| 688 | base_query, |
| 689 | Agents.customer_code, |
| 690 | ) |
| 691 | |
| 692 | result = await session.execute(filtered_query) |
| 693 | agent = result.scalars().first() |
| 694 | |
| 695 | if not agent: |
| 696 | raise HTTPException( |
| 697 | status_code=404, |
| 698 | detail=f"Agent with agent_id {agent_id} not found or access denied", |
| 699 | ) |
| 700 | |
| 701 | agent.critical_asset = True |
| 702 | await session.commit() |
| 703 | |
| 704 | return AgentModifyResponse( |
| 705 | success=True, |
| 706 | message=f"Agent {agent_id} marked as critical: {True}", |
| 707 | ) |
| 708 | except Exception as e: |
| 709 | session.rollback() # Roll back the session in case of error |
| 710 | raise HTTPException( |
| 711 | status_code=500, |
| 712 | detail=f"Failed to mark agent as critical: {str(e)}", |
| 713 | ) |
| 714 | |
| 715 | |
| 716 | @agents_router.post( |
| 717 | "/{agent_id}/noncritical", |
| 718 | response_model=AgentModifyResponse, |
| 719 | description="Mark agent as not critical", |
| 720 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 721 | ) |
| 722 | async def mark_agent_as_not_critical( |
| 723 | agent_id: str, |
| 724 | current_user: User = Depends(AuthHandler().get_current_user), |
| 725 | session: AsyncSession = Depends(get_db), |
| 726 | ) -> AgentModifyResponse: |
| 727 | """ |
| 728 | Marks the specified agent as not critical. |
| 729 | User must have access to the agent's customer. |
| 730 | |
| 731 | Args: |
| 732 | agent_id (str): The ID of the agent to mark as not critical. |
| 733 | current_user (User): The authenticated user. |
| 734 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 735 | |
| 736 | Returns: |
| 737 | AgentModifyResponse: The response indicating the success or failure of the operation. |
| 738 | |
| 739 | Raises: |
| 740 | HTTPException: If the agent with the specified ID is not found or if there is an error marking the agent as not critical. |
| 741 | """ |
| 742 | logger.info(f"Marking agent {agent_id} as not critical") |
| 743 | try: |
| 744 | # Check customer access - first find the agent |
| 745 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 746 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 747 | current_user, |
| 748 | session, |
| 749 | base_query, |
| 750 | Agents.customer_code, |
| 751 | ) |
| 752 | |
| 753 | result = await session.execute(filtered_query) |
| 754 | agent = result.scalars().first() |
| 755 | |
| 756 | if not agent: |
| 757 | raise HTTPException( |
| 758 | status_code=404, |
| 759 | detail=f"Agent with agent_id {agent_id} not found or access denied", |
| 760 | ) |
| 761 | |
| 762 | agent.critical_asset = False |
| 763 | await session.commit() |
| 764 | |
| 765 | return AgentModifyResponse( |
| 766 | success=True, |
| 767 | message=f"Agent {agent_id} marked as not critical", |
| 768 | ) |
| 769 | except Exception as e: |
| 770 | await session.rollback() # Roll back the session in case of error |
| 771 | raise HTTPException( |
| 772 | status_code=500, |
| 773 | detail=f"Failed to mark agent as not critical: {str(e)}", |
| 774 | ) |
| 775 | |
| 776 | |
| 777 | @agents_router.post( |
| 778 | "/{agent_id}/wazuh/upgrade", |
| 779 | response_model=AgentModifyResponse, |
| 780 | description="Upgrade wazuh agent", |
| 781 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 782 | ) |
| 783 | async def upgrade_wazuh_agent_route( |
| 784 | agent_id: str, |
| 785 | current_user: User = Depends(AuthHandler().get_current_user), |
| 786 | session: AsyncSession = Depends(get_db), |
| 787 | ) -> AgentWazuhUpgradeResponse: |
| 788 | """ |
| 789 | Upgrade Wazuh agent. |
| 790 | User must have access to the agent's customer. |
| 791 | |
| 792 | Args: |
| 793 | agent_id (str): The ID of the agent to be upgraded. |
| 794 | current_user (User): The authenticated user. |
| 795 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 796 | |
| 797 | Returns: |
| 798 | AgentModifyResponse: The response indicating the success or failure of the operation. |
| 799 | """ |
| 800 | logger.info(f"Upgrading Wazuh agent {agent_id}") |
| 801 | try: |
| 802 | # Check customer access - first find the agent |
| 803 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 804 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 805 | current_user, |
| 806 | session, |
| 807 | base_query, |
| 808 | Agents.customer_code, |
| 809 | ) |
| 810 | |
| 811 | result = await session.execute(filtered_query) |
| 812 | agent = result.scalars().first() |
| 813 | |
| 814 | if not agent: |
| 815 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 816 | return await upgrade_wazuh_agent(agent_id) |
| 817 | return AgentWazuhUpgradeResponse( |
| 818 | success=True, |
| 819 | message=f"Agent {agent_id} upgraded successfully started. Upgrade may take a few minutes to complete.", |
| 820 | ) |
| 821 | except Exception as e: |
| 822 | logger.error(f"Failed to upgrade Wazuh agent {agent_id}: {e}") |
| 823 | raise HTTPException( |
| 824 | status_code=500, |
| 825 | detail=f"Failed to upgrade Wazuh agent {agent_id}: {e}", |
| 826 | ) |
| 827 | |
| 828 | |
| 829 | @agents_router.get( |
| 830 | "/{agent_id}/vulnerabilities/{vulnerability_severity}", |
| 831 | response_model=WazuhAgentVulnerabilitiesResponse, |
| 832 | description="Get agent vulnerabilities", |
| 833 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 834 | ) |
| 835 | async def get_agent_vulnerabilities( |
| 836 | agent_id: str, |
| 837 | vulnerability_severity: VulnSeverity = Path(..., description="The severity of the vulnerabilities to fetch."), |
| 838 | current_user: User = Depends(AuthHandler().get_current_user), |
| 839 | session: AsyncSession = Depends(get_db), |
| 840 | ) -> WazuhAgentVulnerabilitiesResponse: |
| 841 | """ |
| 842 | Fetches the vulnerabilities of a specific agent. |
| 843 | User must have access to the agent's customer. |
| 844 | |
| 845 | Args: |
| 846 | agent_id (str): The ID of the agent. |
| 847 | vulnerability_severity: The severity level filter. |
| 848 | current_user (User): The authenticated user. |
| 849 | session (AsyncSession): The database session. |
| 850 | |
| 851 | Returns: |
| 852 | WazuhAgentVulnerabilitiesResponse: The response containing the agent vulnerabilities. |
| 853 | """ |
| 854 | logger.info(f"Fetching agent {agent_id} vulnerabilities") |
| 855 | |
| 856 | # Check customer access - first verify user has access to this agent |
| 857 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 858 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 859 | |
| 860 | result = await session.execute(filtered_query) |
| 861 | agent = result.scalars().first() |
| 862 | |
| 863 | if not agent: |
| 864 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 865 | |
| 866 | wazuh_new = await check_wazuh_manager_version() |
| 867 | if wazuh_new is True: |
| 868 | logger.info("Wazuh Manager version is 4.8.0 or higher. Fetching vulnerabilities using new API") |
| 869 | return await collect_agent_vulnerabilities_new(agent_id, vulnerability_severity.value) |
| 870 | return await collect_agent_vulnerabilities(agent_id, vulnerability_severity.value) |
| 871 | |
| 872 | |
| 873 | @agents_router.get( |
| 874 | "/{agent_id}/csv/vulnerabilities/{vulnerability_severity}", |
| 875 | description="Get agent vulnerabilities as CSV", |
| 876 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 877 | ) |
| 878 | async def get_agent_vulnerabilities_csv( |
| 879 | agent_id: str, |
| 880 | vulnerability_severity: VulnSeverity = Path(...), |
| 881 | current_user: User = Depends(AuthHandler().get_current_user), |
| 882 | session: AsyncSession = Depends(get_db), |
| 883 | ) -> StreamingResponse: |
| 884 | """ |
| 885 | Fetches the vulnerabilities of a specific agent and returns them as a CSV file. |
| 886 | User must have access to the agent's customer. |
| 887 | |
| 888 | Args: |
| 889 | agent_id (str): The ID of the agent. |
| 890 | vulnerability_severity: The severity level filter. |
| 891 | current_user (User): The authenticated user. |
| 892 | session (AsyncSession): The database session. |
| 893 | |
| 894 | Returns: |
| 895 | StreamingResponse: The response containing the agent vulnerabilities in CSV format. |
| 896 | """ |
| 897 | logger.info(f"Fetching agent {agent_id} vulnerabilities as CSV") |
| 898 | |
| 899 | # Check customer access - first verify user has access to this agent |
| 900 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 901 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 902 | |
| 903 | result = await session.execute(filtered_query) |
| 904 | agent = result.scalars().first() |
| 905 | |
| 906 | if not agent: |
| 907 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 908 | |
| 909 | wazuh_new = await check_wazuh_manager_version() |
| 910 | if wazuh_new is True: |
| 911 | logger.info("Wazuh Manager version is 4.8.0 or higher. Fetching vulnerabilities using new API") |
| 912 | vulnerabilities = ( |
| 913 | await collect_agent_vulnerabilities_new(agent_id, vulnerability_severity=vulnerability_severity.value) |
| 914 | ).vulnerabilities |
| 915 | else: |
| 916 | vulnerabilities = ( |
| 917 | await collect_agent_vulnerabilities(agent_id, vulnerability_severity=vulnerability_severity.value) |
| 918 | ).vulnerabilities |
| 919 | # Create a CSV file |
| 920 | logger.info(f"Creating CSV file for agent {agent_id} with {len(vulnerabilities)} vulnerabilities") |
| 921 | logger.info(f"Vulnerabilities: {vulnerabilities}") |
| 922 | output = io.StringIO() |
| 923 | writer = csv.writer(output) |
| 924 | # Write the header |
| 925 | writer.writerow( |
| 926 | [ |
| 927 | "Severity", |
| 928 | "Version", |
| 929 | "Type", |
| 930 | "Name", |
| 931 | "External References", |
| 932 | "Detection Time", |
| 933 | "CVSS3 Score", |
| 934 | "Published", |
| 935 | "Architecture", |
| 936 | "CVE", |
| 937 | "Status", |
| 938 | "Title", |
| 939 | "EPSS Score", |
| 940 | ], |
| 941 | ) |
| 942 | # Write the rows |
| 943 | for vulnerability in vulnerabilities: |
| 944 | epss_score = await collect_epss_score(EpssThreatIntelRequest(cve=vulnerability.cve)) |
| 945 | writer.writerow( |
| 946 | [ |
| 947 | vulnerability.severity, |
| 948 | vulnerability.version, |
| 949 | vulnerability.type, |
| 950 | vulnerability.name, |
| 951 | ", ".join(vulnerability.external_references) if vulnerability.external_references else "", |
| 952 | vulnerability.detection_time, |
| 953 | vulnerability.cvss3_score, |
| 954 | vulnerability.published, |
| 955 | vulnerability.architecture, |
| 956 | vulnerability.cve, |
| 957 | vulnerability.status, |
| 958 | vulnerability.title, |
| 959 | epss_score.data[0].epss if epss_score.data else "", |
| 960 | ], |
| 961 | ) |
| 962 | # Return the CSV file as a streaming response |
| 963 | output.seek(0) |
| 964 | return StreamingResponse( |
| 965 | output, # Use the StringIO object directly |
| 966 | media_type="text/csv", |
| 967 | headers={"Content-Disposition": f"attachment; filename={agent_id}_vulnerabilities.csv"}, |
| 968 | ) |
| 969 | |
| 970 | |
| 971 | @agents_router.get( |
| 972 | "/{agent_id}/sca", |
| 973 | response_model=WazuhAgentScaResponse, |
| 974 | description="Get agent sca results", |
| 975 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 976 | ) |
| 977 | async def get_agent_sca( |
| 978 | agent_id: str, |
| 979 | current_user: User = Depends(AuthHandler().get_current_user), |
| 980 | session: AsyncSession = Depends(get_db), |
| 981 | ) -> WazuhAgentScaResponse: |
| 982 | """ |
| 983 | Fetches the sca results of a specific agent. |
| 984 | User must have access to the agent's customer. |
| 985 | |
| 986 | Args: |
| 987 | agent_id (str): The ID of the agent. |
| 988 | current_user (User): The authenticated user. |
| 989 | session (AsyncSession): The database session. |
| 990 | |
| 991 | Returns: |
| 992 | WazuhAgentScaResponse: The response containing the agent sca. |
| 993 | """ |
| 994 | logger.info(f"Fetching agent {agent_id} sca") |
| 995 | |
| 996 | # Check customer access - first verify user has access to this agent |
| 997 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 998 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 999 | |
| 1000 | result = await session.execute(filtered_query) |
| 1001 | agent = result.scalars().first() |
| 1002 | |
| 1003 | if not agent: |
| 1004 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 1005 | |
| 1006 | return await collect_agent_sca(agent_id) |
| 1007 | |
| 1008 | |
| 1009 | @agents_router.get( |
| 1010 | "/{agent_id}/sca/{policy_id}", |
| 1011 | response_model=WazuhAgentScaPolicyResultsResponse, |
| 1012 | description="Get agent sca results", |
| 1013 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 1014 | ) |
| 1015 | async def get_agent_sca_policy_results( |
| 1016 | agent_id: str, |
| 1017 | policy_id: str, |
| 1018 | current_user: User = Depends(AuthHandler().get_current_user), |
| 1019 | session: AsyncSession = Depends(get_db), |
| 1020 | ) -> WazuhAgentScaPolicyResultsResponse: |
| 1021 | """ |
| 1022 | Fetches the sca results of a specific agent. |
| 1023 | User must have access to the agent's customer. |
| 1024 | |
| 1025 | Args: |
| 1026 | agent_id (str): The ID of the agent. |
| 1027 | policy_id (str): The ID of the policy. |
| 1028 | current_user (User): The authenticated user. |
| 1029 | session (AsyncSession): The database session. |
| 1030 | |
| 1031 | Returns: |
| 1032 | WazuhAgentScaPolicyResultsResponse: The response containing the agent sca. |
| 1033 | """ |
| 1034 | logger.info(f"Fetching agent {agent_id} sca policy results") |
| 1035 | |
| 1036 | # Check customer access - first verify user has access to this agent |
| 1037 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 1038 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 1039 | |
| 1040 | result = await session.execute(filtered_query) |
| 1041 | agent = result.scalars().first() |
| 1042 | |
| 1043 | if not agent: |
| 1044 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 1045 | |
| 1046 | return await collect_agent_sca_policy_results(agent_id, policy_id) |
| 1047 | |
| 1048 | |
| 1049 | @agents_router.get( |
| 1050 | "/{agent_id}/csv/sca/{policy_id}", |
| 1051 | description="Get agent sca results as CSV", |
| 1052 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 1053 | ) |
| 1054 | async def get_agent_sca_policy_results_csv( |
| 1055 | agent_id: str, |
| 1056 | policy_id: str, |
| 1057 | current_user: User = Depends(AuthHandler().get_current_user), |
| 1058 | session: AsyncSession = Depends(get_db), |
| 1059 | ) -> StreamingResponse: |
| 1060 | """ |
| 1061 | Fetches the sca results of a specific agent and returns them as a CSV file. |
| 1062 | User must have access to the agent's customer. |
| 1063 | |
| 1064 | Args: |
| 1065 | agent_id (str): The ID of the agent. |
| 1066 | policy_id (str): The ID of the policy. |
| 1067 | current_user (User): The authenticated user. |
| 1068 | session (AsyncSession): The database session. |
| 1069 | |
| 1070 | Returns: |
| 1071 | StreamingResponse: The response containing the agent sca in CSV format. |
| 1072 | """ |
| 1073 | logger.info(f"Fetching agent {agent_id} sca policy results as CSV") |
| 1074 | |
| 1075 | # Check customer access - first verify user has access to this agent |
| 1076 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 1077 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 1078 | |
| 1079 | result = await session.execute(filtered_query) |
| 1080 | agent = result.scalars().first() |
| 1081 | |
| 1082 | if not agent: |
| 1083 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 1084 | |
| 1085 | sca_results = (await collect_agent_sca_policy_results(agent_id, policy_id)).sca_policy_results |
| 1086 | # Create a CSV file |
| 1087 | logger.info(f"Creating CSV file for agent {agent_id} with {len(sca_results)} sca policy results") |
| 1088 | output = io.StringIO() |
| 1089 | writer = csv.writer(output) |
| 1090 | # Write the header |
| 1091 | writer.writerow( |
| 1092 | [ |
| 1093 | "Description", |
| 1094 | "Policy ID", |
| 1095 | "Reason", |
| 1096 | "Command", |
| 1097 | "Rationale", |
| 1098 | "Condition", |
| 1099 | "Title", |
| 1100 | "Result", |
| 1101 | "Remediation", |
| 1102 | "Compliance", |
| 1103 | "Rules", |
| 1104 | ], |
| 1105 | ) |
| 1106 | # Write the rows |
| 1107 | for sca_result in sca_results: |
| 1108 | writer.writerow( |
| 1109 | [ |
| 1110 | sca_result.description, |
| 1111 | sca_result.policy_id, |
| 1112 | sca_result.reason, |
| 1113 | sca_result.command, |
| 1114 | sca_result.rationale, |
| 1115 | sca_result.condition, |
| 1116 | sca_result.title, |
| 1117 | sca_result.result, |
| 1118 | sca_result.remediation, |
| 1119 | ", ".join([f"{compliance.key}: {compliance.value}" for compliance in sca_result.compliance]) |
| 1120 | if sca_result.compliance |
| 1121 | else "", |
| 1122 | ", ".join([f"{rule.type}: {rule.rule}" for rule in sca_result.rules]) if sca_result.rules else "", |
| 1123 | ], |
| 1124 | ) |
| 1125 | # Return the CSV file as a streaming response |
| 1126 | output.seek(0) |
| 1127 | return StreamingResponse( |
| 1128 | output, # Use the StringIO object directly |
| 1129 | media_type="text/csv", |
| 1130 | headers={"Content-Disposition": f"attachment; filename={agent_id}_sca_policy_results.csv"}, |
| 1131 | ) |
| 1132 | |
| 1133 | |
| 1134 | @agents_router.get( |
| 1135 | "/{agent_hostname}/cases", |
| 1136 | response_model=CaseOutResponse, |
| 1137 | description="Get cases for agent", |
| 1138 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 1139 | ) |
| 1140 | async def get_agent_soc_cases( |
| 1141 | agent_hostname: str, |
| 1142 | current_user: User = Depends(AuthHandler().get_current_user), |
| 1143 | session: AsyncSession = Depends(get_db), |
| 1144 | ): |
| 1145 | """ |
| 1146 | Fetches the SOC cases of a specific agent. |
| 1147 | User must have access to the agent's customer. |
| 1148 | |
| 1149 | Args: |
| 1150 | agent_hostname (str): The hostname of the agent. |
| 1151 | current_user (User): The authenticated user. |
| 1152 | session (AsyncSession): The database session. |
| 1153 | |
| 1154 | Returns: |
| 1155 | SocCasesResponse: The response containing the agent SOC cases. |
| 1156 | """ |
| 1157 | logger.info(f"Fetching agent {agent_hostname} cases") |
| 1158 | |
| 1159 | # Check customer access - first verify user has access to this agent |
| 1160 | base_query = select(Agents).filter(Agents.hostname == agent_hostname) |
| 1161 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 1162 | |
| 1163 | result = await session.execute(filtered_query) |
| 1164 | agent = result.scalars().first() |
| 1165 | |
| 1166 | if not agent: |
| 1167 | raise HTTPException(status_code=404, detail=f"Agent with hostname {agent_hostname} not found or access denied") |
| 1168 | |
| 1169 | return CaseOutResponse( |
| 1170 | cases=await list_cases_by_asset_name(asset_name=agent_hostname, db=session), |
| 1171 | success=True, |
| 1172 | message="Cases retrieved successfully", |
| 1173 | ) |
| 1174 | |
| 1175 | |
| 1176 | @agents_router.get( |
| 1177 | "/wazuh/outdated", |
| 1178 | response_model=OutdatedWazuhAgentsResponse, |
| 1179 | description="Get all outdated Wazuh agents", |
| 1180 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1181 | ) |
| 1182 | async def get_outdated_wazuh_agents( |
| 1183 | session: AsyncSession = Depends(get_db), |
| 1184 | ) -> OutdatedWazuhAgentsResponse: |
| 1185 | """ |
| 1186 | Retrieve all outdated Wazuh agents. |
| 1187 | |
| 1188 | This endpoint requires the user to have either the "admin" or "analyst" scope. |
| 1189 | |
| 1190 | Returns: |
| 1191 | OutdatedWazuhAgentsResponse: The response containing the outdated Wazuh agents. |
| 1192 | """ |
| 1193 | logger.info("Fetching all outdated Wazuh agents") |
| 1194 | return await get_outdated_agents_wazuh(session) |
| 1195 | |
| 1196 | |
| 1197 | @agents_router.get( |
| 1198 | "/velociraptor/outdated", |
| 1199 | response_model=OutdatedVelociraptorAgentsResponse, |
| 1200 | description="Get all outdated Velociraptor agents", |
| 1201 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1202 | ) |
| 1203 | async def get_outdated_velociraptor_agents( |
| 1204 | session: AsyncSession = Depends(get_db), |
| 1205 | ) -> OutdatedVelociraptorAgentsResponse: |
| 1206 | """ |
| 1207 | Fetches all outdated Velociraptor agents. |
| 1208 | |
| 1209 | Parameters: |
| 1210 | - session: The database session. |
| 1211 | |
| 1212 | Returns: |
| 1213 | - OutdatedVelociraptorAgentsResponse: The response containing the outdated Velociraptor agents. |
| 1214 | """ |
| 1215 | logger.info("Fetching all outdated Velociraptor agents") |
| 1216 | return await get_outdated_agents_velociraptor(session) |
| 1217 | |
| 1218 | |
| 1219 | @agents_router.put( |
| 1220 | "/{agent_id}/update", |
| 1221 | response_model=AgentModifyResponse, |
| 1222 | description="Update agent", |
| 1223 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1224 | ) |
| 1225 | async def update_agent( |
| 1226 | agent_id: str, |
| 1227 | velociraptor_id: str, |
| 1228 | current_user: User = Depends(AuthHandler().get_current_user), |
| 1229 | session: AsyncSession = Depends(get_db), |
| 1230 | ) -> AgentModifyResponse: |
| 1231 | """ |
| 1232 | Updates an agent's velociraptor_id |
| 1233 | User must have access to the agent's customer. |
| 1234 | |
| 1235 | Args: |
| 1236 | agent_id (str): The ID of the agent to be updated. |
| 1237 | velociraptor_id (str): The new velociraptor_id of the agent. |
| 1238 | current_user (User): The authenticated user. |
| 1239 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 1240 | |
| 1241 | Returns: |
| 1242 | AgentModifyResponse: The response indicating the success or failure of the update. |
| 1243 | """ |
| 1244 | logger.info(f"Updating agent {agent_id} with Velociraptor ID: {velociraptor_id}") |
| 1245 | try: |
| 1246 | # Check customer access - first find the agent |
| 1247 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 1248 | filtered_query = await customer_access_handler.filter_query_by_customer_access( |
| 1249 | current_user, |
| 1250 | session, |
| 1251 | base_query, |
| 1252 | Agents.customer_code, |
| 1253 | ) |
| 1254 | |
| 1255 | result = await session.execute(filtered_query) |
| 1256 | agent = result.scalars().first() |
| 1257 | |
| 1258 | if not agent: |
| 1259 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 1260 | |
| 1261 | agent.velociraptor_id = velociraptor_id |
| 1262 | await session.commit() |
| 1263 | logger.info(f"Agent {agent_id} updated with Velociraptor ID: {velociraptor_id}") |
| 1264 | return AgentModifyResponse( |
| 1265 | success=True, |
| 1266 | message=f"Agent {agent_id} updated with Velociraptor ID: {velociraptor_id}", |
| 1267 | ) |
| 1268 | except Exception as e: |
| 1269 | logger.error(f"Failed to update agent {agent_id} with Velociraptor ID: {velociraptor_id}: {e}") |
| 1270 | raise HTTPException( |
| 1271 | status_code=500, |
| 1272 | detail=f"Failed to update agent {agent_id} with Velociraptor ID: {velociraptor_id}: {e}", |
| 1273 | ) |
| 1274 | |
| 1275 | |
| 1276 | @agents_router.delete( |
| 1277 | "/{agent_id}/delete", |
| 1278 | response_model=AgentModifyResponse, |
| 1279 | description="Delete agent", |
| 1280 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1281 | ) |
| 1282 | async def delete_agent( |
| 1283 | agent_id: str, |
| 1284 | current_user: User = Depends(AuthHandler().get_current_user), |
| 1285 | session: AsyncSession = Depends(get_db), |
| 1286 | ) -> AgentModifyResponse: |
| 1287 | """ |
| 1288 | Delete an agent. |
| 1289 | User must have access to the agent's customer. |
| 1290 | |
| 1291 | Args: |
| 1292 | agent_id (str): The ID of the agent to be deleted. |
| 1293 | current_user (User): The authenticated user. |
| 1294 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 1295 | |
| 1296 | Returns: |
| 1297 | AgentModifyResponse: The response indicating the success or failure of the deletion. |
| 1298 | """ |
| 1299 | logger.info(f"Deleting agent {agent_id}") |
| 1300 | |
| 1301 | # Check customer access - first find the agent to verify access |
| 1302 | base_query = select(Agents).filter(Agents.agent_id == agent_id) |
| 1303 | filtered_query = await customer_access_handler.filter_query_by_customer_access(current_user, session, base_query, Agents.customer_code) |
| 1304 | |
| 1305 | result = await session.execute(filtered_query) |
| 1306 | agent = result.scalars().first() |
| 1307 | |
| 1308 | if not agent: |
| 1309 | raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found or access denied") |
| 1310 | |
| 1311 | upstream_errors: list[str] = [] |
| 1312 | |
| 1313 | try: |
| 1314 | await delete_agent_wazuh(agent_id) |
| 1315 | except Exception as e: |
| 1316 | logger.error(f"Wazuh delete failed for agent {agent_id}: {e}") |
| 1317 | upstream_errors.append(f"Wazuh: {e}") |
| 1318 | |
| 1319 | try: |
| 1320 | client_id = await fetch_velociraptor_id(db=session, agent_id=agent_id) |
| 1321 | logger.info(f"Client ID: {client_id}") |
| 1322 | if client_id != "Unknown": |
| 1323 | await delete_agent_velociraptor(client_id) |
| 1324 | except Exception as e: |
| 1325 | logger.error(f"Velociraptor delete failed for agent {agent_id}: {e}") |
| 1326 | upstream_errors.append(f"Velociraptor: {e}") |
| 1327 | |
| 1328 | await delete_agent_from_database(db=session, agent_id=agent_id) |
| 1329 | |
| 1330 | if upstream_errors: |
| 1331 | return AgentModifyResponse( |
| 1332 | success=True, |
| 1333 | message=f"Agent {agent_id} deleted from CoPilot; upstream services reported: {'; '.join(upstream_errors)}", |
| 1334 | ) |
| 1335 | return AgentModifyResponse( |
| 1336 | success=True, |
| 1337 | message=f"Agent {agent_id} deleted successfully", |
| 1338 | ) |
| 1339 | |
| 1340 | |
| 1341 | @agents_router.get( |
| 1342 | "/sync/vulnerabilities", |
| 1343 | description="Sync agent vulnerabilities", |
| 1344 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1345 | ) |
| 1346 | async def sync_vulnerabilities_route( |
| 1347 | session: AsyncSession = Depends(get_db), |
| 1348 | ): |
| 1349 | """ |
| 1350 | Only applies to Wazuh Manager Version 4.8.1 or higher. |
| 1351 | 1. Loops through all agents in the database to collect their agent_name and customer code. |
| 1352 | 2. Queries the `wazuh-states-vulnerabilities-*` index in Wazuh Indexer to get vulnerabilities based on the agent_name. |
| 1353 | 3. Checks the `wazuh-vulnerabilities-*customer_code*` index in Wazuh Indexer to get vulnerabilities based on the |
| 1354 | agent_name and checks to see if a vulnerability_id already exists. |
| 1355 | 4. If the vulnerability_id does not exist, it is sent to the Graylog GELF Input. |
| 1356 | """ |
| 1357 | logger.info("Syncing agent vulnerabilities") |
| 1358 | agents = await get_agents(session) |
| 1359 | for agent in agents.agents: |
| 1360 | if agent.customer_code is None: |
| 1361 | logger.info(f"Skipping agent {agent.hostname} due to missing customer code") |
| 1362 | continue |
| 1363 | await sync_agent_vulnerabilities(agent.hostname, agent.customer_code) |
| 1364 | return {"success": True, "message": "Agent vulnerabilities synced successfully"} |
| 1365 | |
| 1366 | |
| 1367 | @agents_router.post( |
| 1368 | "/sync/vulnerabilities/{customer_code}", |
| 1369 | description="Sync agent vulnerabilities", |
| 1370 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1371 | ) |
| 1372 | async def sync_vulnerabilities_customer_code_route( |
| 1373 | customer_code: str, |
| 1374 | background_tasks: BackgroundTasks, |
| 1375 | session: AsyncSession = Depends(get_db), |
| 1376 | ): |
| 1377 | logger.info("Syncing agent vulnerabilities") |
| 1378 | agents = await get_agents_by_customer_code(customer_code, session) |
| 1379 | for agent in agents: |
| 1380 | if agent.customer_code is None: |
| 1381 | logger.info(f"Skipping agent {agent.hostname} due to missing customer code") |
| 1382 | continue |
| 1383 | background_tasks.add_task(sync_agent_vulnerabilities, agent.hostname, customer_code) |
| 1384 | return {"success": True, "message": "Agent vulnerabilities sync initiated successfully"} |
| 1385 | |
| 1386 | |
| 1387 | # ! TODO: CURRENTLY UPDATES IN THE DB BUT NEED TO UPDATE IN WAZUH # ! |
| 1388 | # @agents_router.put( |
| 1389 | # "/{agent_id}/update-customer-code", |
| 1390 | # response_model=AgentUpdateCustomerCodeResponse, |
| 1391 | # description="Update `agent` customer code", |
| 1392 | # dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 1393 | # ) |
| 1394 | # async def update_agent_customer_code(agent_id: str, body: AgentUpdateCustomerCodeBody, db: AsyncSession = Depends(get_db)) -> AgentUpdateCustomerCodeResponse: |
| 1395 | # logger.info(f"Updating agent {agent_id} customer code to {body.customer_code}") |
| 1396 | # try: |
| 1397 | # result = await db.execute(select(Agents).filter(Agents.agent_id == agent_id)) |
| 1398 | # agent = result.scalars().first() |
| 1399 | # if not agent: |
| 1400 | # raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found") |
| 1401 | # agent.customer_code = body.customer_code |
| 1402 | # await db.commit() |
| 1403 | # logger.info(f"Agent {agent_id} customer code updated to {body.customer_code}") |
| 1404 | # return {"success": True, "message": f"Agent {agent_id} customer code updated to {body.customer_code}"} |
| 1405 | # except Exception as e: |
| 1406 | # if not agent: |
| 1407 | # raise HTTPException(status_code=404, detail=f"Agent with agent_id {agent_id} not found") |