| 1 | from typing import List |
| 2 | |
| 3 | from fastapi import HTTPException |
| 4 | from loguru import logger |
| 5 | from sqlalchemy.ext.asyncio import AsyncSession |
| 6 | from sqlalchemy.future import select |
| 7 | |
| 8 | from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse |
| 9 | from app.agents.schema.agents import OutdatedWazuhAgentsResponse |
| 10 | from app.connectors.velociraptor.utils.universal import UniversalService |
| 11 | from app.db.db_session import session |
| 12 | from app.db.universal_models import Agents |
| 13 | |
| 14 | |
| 15 | def get_agent(agent_id: str) -> List[Agents]: |
| 16 | """ |
| 17 | Retrieves a specific agent from the database using its ID. |
| 18 | |
| 19 | Args: |
| 20 | agent_id (str): The ID of the agent to retrieve. |
| 21 | |
| 22 | Returns: |
| 23 | AgentMetadata: The agent object if found, otherwise None. |
| 24 | """ |
| 25 | try: |
| 26 | return session.query(Agents).filter(Agents.agent_id == agent_id).first() |
| 27 | except Exception as e: |
| 28 | logger.error(f"Failed to fetch agent with agent_id {agent_id}: {e}") |
| 29 | raise HTTPException( |
| 30 | status_code=500, |
| 31 | detail=f"Failed to fetch agent with agent_id {agent_id}: {e}", |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | async def get_agents_by_customer_code(customer_code: str, session: AsyncSession) -> List[Agents]: |
| 36 | """ |
| 37 | Retrieves all agents associated with a specific customer code from the database asynchronously. |
| 38 | |
| 39 | Args: |
| 40 | customer_code (str): The customer code to filter agents by. |
| 41 | session (AsyncSession): The SQLAlchemy asynchronous session to use for the query. |
| 42 | |
| 43 | Returns: |
| 44 | List[Agents]: A list of agents associated with the customer code. |
| 45 | """ |
| 46 | try: |
| 47 | agents_result = await session.execute(select(Agents).filter(Agents.customer_code == customer_code)) |
| 48 | agents = agents_result.scalars().all() |
| 49 | |
| 50 | return agents |
| 51 | except Exception as e: |
| 52 | logger.error(f"Failed to fetch agents with customer_code {customer_code}: {e}") |
| 53 | raise HTTPException( |
| 54 | status_code=500, |
| 55 | detail=f"Failed to fetch agents with customer_code {customer_code}: {e}", |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | async def get_agent_os_by_id(agent_id: str, session: AsyncSession) -> str: |
| 60 | """ |
| 61 | Retrieves the operating system of a specific agent from the database using its ID. |
| 62 | |
| 63 | Args: |
| 64 | agent_id (str): The ID of the agent to retrieve. |
| 65 | session (AsyncSession): The SQLAlchemy asynchronous session to use for the query. |
| 66 | |
| 67 | Returns: |
| 68 | str: The operating system of the agent if found, otherwise None. |
| 69 | """ |
| 70 | try: |
| 71 | agent_result = await session.execute(select(Agents).filter(Agents.agent_id == agent_id)) |
| 72 | agent = agent_result.scalars().first() |
| 73 | |
| 74 | if agent is None: |
| 75 | logger.error(f"Agent with agent_id {agent_id} not found.") |
| 76 | raise HTTPException( |
| 77 | status_code=404, |
| 78 | detail=f"Agent with agent_id {agent_id} not found.", |
| 79 | ) |
| 80 | |
| 81 | return agent.os |
| 82 | except Exception as e: |
| 83 | logger.error(f"Failed to fetch agent with agent_id {agent_id}: {e}") |
| 84 | raise HTTPException( |
| 85 | status_code=500, |
| 86 | detail=f"Failed to fetch agent with agent_id {agent_id}: {e}", |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | async def get_outdated_agents_wazuh( |
| 91 | session: AsyncSession, |
| 92 | ) -> OutdatedWazuhAgentsResponse: |
| 93 | """ |
| 94 | Retrieves all agents with outdated Wazuh agent versions from the database asynchronously. |
| 95 | |
| 96 | Args: |
| 97 | session (AsyncSession): The SQLAlchemy asynchronous session to use for the query. |
| 98 | |
| 99 | Returns: |
| 100 | OutdatedWazuhAgentsResponse: Response object containing the outdated agents. |
| 101 | """ |
| 102 | try: |
| 103 | wazuh_manager_result = await session.execute( |
| 104 | select(Agents).filter(Agents.agent_id == "000"), |
| 105 | ) |
| 106 | wazuh_manager = wazuh_manager_result.scalars().first() |
| 107 | |
| 108 | if wazuh_manager is None: |
| 109 | logger.error("Wazuh Manager with agent_id '000' not found.") |
| 110 | raise HTTPException( |
| 111 | status_code=404, |
| 112 | detail="Wazuh Manager with agent_id '000' not found.", |
| 113 | ) |
| 114 | |
| 115 | outdated_agents_result = await session.execute( |
| 116 | select(Agents).filter( |
| 117 | Agents.agent_id != "000", |
| 118 | Agents.wazuh_agent_version != wazuh_manager.wazuh_agent_version, |
| 119 | ), |
| 120 | ) |
| 121 | outdated_wazuh_agents = outdated_agents_result.scalars().all() |
| 122 | |
| 123 | return OutdatedWazuhAgentsResponse( |
| 124 | message="Outdated Wazuh agents fetched successfully.", |
| 125 | success=True, |
| 126 | outdated_wazuh_agents=outdated_wazuh_agents, |
| 127 | ) |
| 128 | except Exception as e: |
| 129 | raise HTTPException( |
| 130 | status_code=500, |
| 131 | detail=f"Failed to fetch outdated Wazuh agents: {e}", |
| 132 | ) |
| 133 | |
| 134 | |
| 135 | async def get_outdated_agents_velociraptor( |
| 136 | session: AsyncSession, |
| 137 | ) -> OutdatedVelociraptorAgentsResponse: |
| 138 | """ |
| 139 | Retrieves all agents with outdated Velociraptor client versions from the database asynchronously. |
| 140 | |
| 141 | Args: |
| 142 | session (AsyncSession): The SQLAlchemy asynchronous session to use for the query. |
| 143 | |
| 144 | Returns: |
| 145 | OutdatedVelociraptorAgentsResponse: Response object containing the outdated agents. |
| 146 | """ |
| 147 | vql_server_version = "select * from config" |
| 148 | velociraptor_service = await UniversalService().create("Velociraptor") |
| 149 | |
| 150 | try: |
| 151 | # Assuming _get_server_version is an async function |
| 152 | server_version = await velociraptor_service._get_server_version( |
| 153 | vql_server_version, |
| 154 | ) |
| 155 | agents_result = await session.execute(select(Agents)) |
| 156 | agents = agents_result.scalars().all() |
| 157 | outdated_velociraptor_agents = [agent for agent in agents if agent.velociraptor_agent_version != server_version] |
| 158 | |
| 159 | return OutdatedVelociraptorAgentsResponse( |
| 160 | message="Outdated Velociraptor agents fetched successfully.", |
| 161 | success=True, |
| 162 | outdated_velociraptor_agents=outdated_velociraptor_agents, |
| 163 | ) |
| 164 | except Exception as e: |
| 165 | raise HTTPException( |
| 166 | status_code=500, |
| 167 | detail=f"Failed to fetch outdated Velociraptor agents: {e}", |
| 168 | ) |