| 1 | from datetime import datetime |
| 2 | from datetime import timezone |
| 3 | from typing import List |
| 4 | |
| 5 | from fastapi import HTTPException |
| 6 | from loguru import logger |
| 7 | from sqlalchemy.ext.asyncio import AsyncSession |
| 8 | from sqlalchemy.future import select |
| 9 | |
| 10 | import app.agents.velociraptor.services.agents as velociraptor_services |
| 11 | import app.agents.wazuh.services.agents as wazuh_services |
| 12 | from app.agents.schema.agents import SyncedAgentsResponse |
| 13 | from app.agents.schema.agents import SyncedWazuhAgent |
| 14 | from app.agents.velociraptor.schema.agents import VelociraptorAgent |
| 15 | from app.agents.velociraptor.schema.agents import VelociraptorClients |
| 16 | from app.agents.velociraptor.schema.agents import VelociraptorOrganizations |
| 17 | from app.agents.wazuh.schema.agents import WazuhAgent |
| 18 | from app.agents.wazuh.schema.agents import WazuhAgentsList |
| 19 | from app.connectors.models import Connectors |
| 20 | from app.db.db_session import get_db_session |
| 21 | from app.db.universal_models import Agents |
| 22 | |
| 23 | |
| 24 | async def fetch_wazuh_agents() -> WazuhAgentsList: |
| 25 | """ |
| 26 | Fetch agents from Wazuh service. |
| 27 | |
| 28 | This function retrieves a list of agents from the Wazuh service. |
| 29 | It calls the `collect_wazuh_agents` function from the `wazuh_services` module |
| 30 | and returns a `WazuhAgentsList` object containing the collected agents. |
| 31 | |
| 32 | Returns: |
| 33 | WazuhAgentsList: A list of agents retrieved from the Wazuh service. |
| 34 | |
| 35 | """ |
| 36 | collected_wazuh_agents = await wazuh_services.collect_wazuh_agents() |
| 37 | return WazuhAgentsList( |
| 38 | agents=collected_wazuh_agents.agents, |
| 39 | success=collected_wazuh_agents.success, |
| 40 | message=collected_wazuh_agents.message, |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | async def fetch_velociraptor_clients(org_id: str) -> VelociraptorClients: |
| 45 | """ |
| 46 | Fetches clients from Velociraptor service. |
| 47 | |
| 48 | Args: |
| 49 | None |
| 50 | |
| 51 | Returns: |
| 52 | VelociraptorClientsList: The fetched clients. |
| 53 | """ |
| 54 | collected_velociraptor_agents = await velociraptor_services.collect_velociraptor_clients(org_id=org_id) |
| 55 | return VelociraptorClients( |
| 56 | clients=collected_velociraptor_agents, |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | async def fetch_velociraptor_organizations() -> VelociraptorOrganizations: |
| 61 | """ |
| 62 | Fetches organizations from Velociraptor service. |
| 63 | |
| 64 | Args: |
| 65 | None |
| 66 | |
| 67 | Returns: |
| 68 | VelociraptorOrgsList: The fetched orgs. |
| 69 | """ |
| 70 | collected_velociraptor_orgs = await velociraptor_services.collect_velociraptor_organizations() |
| 71 | logger.info(f"Collected Velociraptor Orgs: {collected_velociraptor_orgs}") |
| 72 | return VelociraptorOrganizations( |
| 73 | organizations=collected_velociraptor_orgs, |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | async def fetch_velociraptor_agent(agent_name: str) -> VelociraptorAgent: |
| 78 | """ |
| 79 | Fetches agent details from Velociraptor service. |
| 80 | |
| 81 | Args: |
| 82 | agent_name (str): The name of the agent to fetch. |
| 83 | |
| 84 | Returns: |
| 85 | VelociraptorAgent: The fetched agent details. |
| 86 | """ |
| 87 | return await velociraptor_services.collect_velociraptor_agent(agent_name) |
| 88 | |
| 89 | |
| 90 | async def fetch_velociraptor_agent_via_client_id(client_id: str) -> VelociraptorAgent: |
| 91 | """ |
| 92 | Fetches agent details from Velociraptor service. |
| 93 | |
| 94 | Args: |
| 95 | client_id (str): The client_id of the agent to fetch. |
| 96 | |
| 97 | Returns: |
| 98 | VelociraptorAgent: The fetched agent details. |
| 99 | """ |
| 100 | return await velociraptor_services.collect_velociraptor_agent_via_client_id(client_id) |
| 101 | |
| 102 | |
| 103 | async def add_wazuh_agent_in_db( |
| 104 | session: AsyncSession, |
| 105 | agent: WazuhAgent, |
| 106 | customer_code: str, |
| 107 | ): |
| 108 | """Add new agent to database. |
| 109 | |
| 110 | Args: |
| 111 | session (AsyncSession): The asynchronous session object for database operations. |
| 112 | agent (WazuhAgent): The Wazuh agent object to be added. |
| 113 | customer_code (str): The customer code for the agent. |
| 114 | |
| 115 | Returns: |
| 116 | None |
| 117 | |
| 118 | """ |
| 119 | new_agent = Agents.create_wazuh_agent_from_model(agent, customer_code) |
| 120 | session.add(new_agent) |
| 121 | logger.info(f"Adding agent {agent.agent_name} to the database") |
| 122 | try: |
| 123 | await session.commit() # Use the await keyword to commit asynchronously |
| 124 | except Exception as e: |
| 125 | logger.error(f"Failed to add agent {agent.agent_name} to the database: {e}") |
| 126 | await session.rollback() |
| 127 | raise HTTPException(status_code=500, detail=str(e)) |
| 128 | logger.info(f"Agent {agent.agent_name} added to the database") |
| 129 | |
| 130 | |
| 131 | async def update_wazuh_agent_in_db( |
| 132 | session: AsyncSession, |
| 133 | existing_agent: Agents, |
| 134 | agent: WazuhAgent, |
| 135 | customer_code: str, |
| 136 | ): |
| 137 | """Update existing agent in database. |
| 138 | |
| 139 | Args: |
| 140 | session (AsyncSession): The async session object for database operations. |
| 141 | existing_agent (Agents): The existing agent object in the database. |
| 142 | agent (WazuhAgent): The updated agent object. |
| 143 | customer_code (str): The customer code associated with the agent. |
| 144 | |
| 145 | Returns: |
| 146 | None |
| 147 | |
| 148 | """ |
| 149 | existing_agent.update_wazuh_agent_from_model(agent, customer_code) |
| 150 | await session.commit() # Use the await keyword to commit asynchronously |
| 151 | logger.info(f"Agent {agent.agent_name} updated in the database") |
| 152 | |
| 153 | |
| 154 | def extract_customer_code(customer_code: str): |
| 155 | """Extracts the customer code from the agent label. |
| 156 | |
| 157 | Args: |
| 158 | customer_code (str): The agent label containing the customer code. |
| 159 | |
| 160 | Returns: |
| 161 | str: The extracted customer code, or None if the agent label is invalid. |
| 162 | """ |
| 163 | parts = customer_code.split("_") |
| 164 | return parts[1] if len(parts) > 1 else None |
| 165 | |
| 166 | |
| 167 | async def get_velociraptor_connector(session): |
| 168 | """ |
| 169 | Retrieves the Velociraptor connector from the database. |
| 170 | |
| 171 | Args: |
| 172 | session: The database session. |
| 173 | |
| 174 | Returns: |
| 175 | The first result of the query as a scalar value. |
| 176 | """ |
| 177 | connector_query = select(Connectors).filter( |
| 178 | Connectors.connector_name == "Velociraptor", |
| 179 | ) |
| 180 | result = await session.execute(connector_query) |
| 181 | return result.scalars().first() |
| 182 | |
| 183 | |
| 184 | async def get_velociraptor_agent(agent_name): |
| 185 | """ |
| 186 | Retrieves a Velociraptor agent with the specified name. |
| 187 | |
| 188 | Args: |
| 189 | agent_name (str): The name of the agent to retrieve. |
| 190 | |
| 191 | Returns: |
| 192 | VelociraptorAgent: The retrieved Velociraptor agent, or None if retrieval fails. |
| 193 | """ |
| 194 | try: |
| 195 | return await fetch_velociraptor_agent(agent_name) |
| 196 | except Exception as e: |
| 197 | logger.error(f"Failed to collect Velociraptor Agent for {agent_name}: {e}") |
| 198 | return None |
| 199 | |
| 200 | |
| 201 | async def get_velociraptor_agent_by_client_id(client_id): |
| 202 | """ |
| 203 | Retrieves a Velociraptor agent with the specified client_id. |
| 204 | |
| 205 | Args: |
| 206 | client_id (str): The client_id of the agent to retrieve. |
| 207 | |
| 208 | Returns: |
| 209 | VelociraptorAgent: The retrieved Velociraptor agent, or None if retrieval fails. |
| 210 | """ |
| 211 | try: |
| 212 | return await fetch_velociraptor_agent_via_client_id(client_id) |
| 213 | except Exception as e: |
| 214 | logger.error(f"Failed to collect Velociraptor Agent for {client_id}: {e}") |
| 215 | return None |
| 216 | |
| 217 | |
| 218 | async def process_velociraptor_agent(session, agent, client_id=None): |
| 219 | """ |
| 220 | Process the Velociraptor agent for a given Wazuh agent. |
| 221 | |
| 222 | Args: |
| 223 | session (object): The session object for the connection. |
| 224 | wazuh_agent (object): The Wazuh agent object. |
| 225 | |
| 226 | Returns: |
| 227 | object: The Velociraptor agent object if successful, None otherwise. |
| 228 | """ |
| 229 | try: |
| 230 | velociraptor_connector = await get_velociraptor_connector(session) |
| 231 | if velociraptor_connector.connector_verified: |
| 232 | velociraptor_agent = await get_velociraptor_agent(agent) |
| 233 | if client_id is not None: |
| 234 | velociraptor_agent = await get_velociraptor_agent_by_client_id(client_id) |
| 235 | else: |
| 236 | velociraptor_agent = VelociraptorAgent( |
| 237 | client_id="Unknown", |
| 238 | client_last_seen="1970-01-01T00:00:00+00:00", |
| 239 | client_version="Unknown", |
| 240 | ) |
| 241 | return velociraptor_agent |
| 242 | except Exception as e: |
| 243 | logger.error(f"Failed to process agent {agent}: {e}") |
| 244 | return None |
| 245 | |
| 246 | |
| 247 | async def sync_agents_wazuh() -> SyncedAgentsResponse: |
| 248 | wazuh_agents_list = await fetch_wazuh_agents() |
| 249 | logger.info(f"Collected Wazuh Agents: {wazuh_agents_list}") |
| 250 | |
| 251 | agents_added_list: List[WazuhAgent] = [] |
| 252 | |
| 253 | async with get_db_session() as session: # Create a new session here |
| 254 | for wazuh_agent in wazuh_agents_list.agents: |
| 255 | customer_code = extract_customer_code(wazuh_agent.agent_label) |
| 256 | |
| 257 | existing_agent_query = select(Agents).filter( |
| 258 | Agents.hostname == wazuh_agent.agent_name, |
| 259 | ) |
| 260 | result = await session.execute(existing_agent_query) |
| 261 | existing_agent = result.scalars().first() |
| 262 | |
| 263 | if existing_agent: |
| 264 | await update_wazuh_agent_in_db(session, existing_agent, wazuh_agent, customer_code) |
| 265 | else: |
| 266 | await add_wazuh_agent_in_db(session, wazuh_agent, customer_code) |
| 267 | |
| 268 | synced_wazuh_agent = SyncedWazuhAgent(**wazuh_agent.model_dump()) |
| 269 | agents_added_list.append(synced_wazuh_agent) |
| 270 | |
| 271 | logger.info(f"Agents Added List: {agents_added_list}") |
| 272 | |
| 273 | # Close the session |
| 274 | await session.close() |
| 275 | |
| 276 | return SyncedAgentsResponse( |
| 277 | success=True, |
| 278 | message="Agents synced successfully", |
| 279 | ) |
| 280 | |
| 281 | |
| 282 | async def update_agent_with_velociraptor_in_db( |
| 283 | session: AsyncSession, |
| 284 | agent: Agents, |
| 285 | velociraptor_agent: VelociraptorAgent, |
| 286 | ): |
| 287 | """Update existing agent in database with Velociraptor details. |
| 288 | |
| 289 | Args: |
| 290 | session (AsyncSession): The async session object for database operations. |
| 291 | agent (Agents): The existing agent object in the database. |
| 292 | client (VelociraptorAgent): The updated client object. |
| 293 | |
| 294 | Returns: |
| 295 | None |
| 296 | |
| 297 | """ |
| 298 | logger.info(f"Updating agent {agent.hostname} with Velociraptor details in the database") |
| 299 | agent.update_velociraptor_details(velociraptor_agent) |
| 300 | session.add(agent) # Add the updated agent back to the session |
| 301 | await session.commit() # Use the await keyword to commit asynchronously |
| 302 | logger.info("Agent updated with Velociraptor details in the database") |
| 303 | |
| 304 | |
| 305 | async def sync_agents_velociraptor() -> SyncedAgentsResponse: |
| 306 | """ |
| 307 | Syncronizes the agents with Velociraptor. This function retrieves all the |
| 308 | agents from the `Agents` table and invokes the velociraptor API with the |
| 309 | hostname. If the hostname cannot be found within Velociraptor, and the agent's |
| 310 | `velociraptor_id` is not None, invoke the Velociraptor API and pass it the |
| 311 | `velociraptor_id`. |
| 312 | |
| 313 | :param session: The database session to use for querying and updating agents. |
| 314 | :type session: AsyncSession |
| 315 | :return: The response indicating the success of the synchronization operation and the list of agents added. |
| 316 | :rtype: SyncedAgentsResponse |
| 317 | """ |
| 318 | agents_added_list: List[VelociraptorAgent] = [] |
| 319 | velo_orgs = await fetch_velociraptor_organizations() |
| 320 | logger.info(f"Collected Velociraptor Orgs: {velo_orgs}") |
| 321 | for org in velo_orgs.organizations: |
| 322 | velociraptor_clients = await fetch_velociraptor_clients(org_id=org.OrgId) |
| 323 | logger.info(f"Collected Velociraptor Clients: {velociraptor_clients}") |
| 324 | velociraptor_clients = velociraptor_clients.clients if hasattr(velociraptor_clients, "clients") else [] |
| 325 | |
| 326 | async with get_db_session() as session: # Create a new session here |
| 327 | existing_agents_query = select(Agents) |
| 328 | result = await session.execute(existing_agents_query) |
| 329 | existing_agents = result.scalars().all() |
| 330 | |
| 331 | for agent in existing_agents: |
| 332 | logger.info(f"Collecting Velociraptor Agent for {agent.hostname}") |
| 333 | |
| 334 | try: |
| 335 | # Build the velociraptor_agent where the hostname or `client_id` is that equal to the `agents` |
| 336 | velociraptor_agent = next( |
| 337 | ( |
| 338 | client |
| 339 | for client in velociraptor_clients |
| 340 | if client.os_info.hostname == agent.hostname or client.client_id == agent.velociraptor_id |
| 341 | ), |
| 342 | None, |
| 343 | ) |
| 344 | # Convert Unix epoch timestamp to datetime |
| 345 | last_seen_at = datetime.fromtimestamp( |
| 346 | int(velociraptor_agent.last_seen_at) / 1e6, |
| 347 | ) # Divide by 1e6 to convert from microseconds to seconds |
| 348 | # Convert datetime to ISO 8601 format without fractional seconds |
| 349 | last_seen_at_iso = last_seen_at.replace(tzinfo=timezone.utc).isoformat(timespec="seconds") |
| 350 | velociraptor_agent = VelociraptorAgent( |
| 351 | velociraptor_id=velociraptor_agent.client_id, |
| 352 | velociraptor_last_seen=last_seen_at_iso, |
| 353 | velociraptor_agent_version=velociraptor_agent.agent_information.version, |
| 354 | velociraptor_org=org.OrgId, |
| 355 | ) |
| 356 | |
| 357 | except Exception as e: |
| 358 | logger.error( |
| 359 | f"Failed to collect Velociraptor Agent for {agent.hostname}: {e}", |
| 360 | ) |
| 361 | continue |
| 362 | |
| 363 | if velociraptor_agent: |
| 364 | # Update the agent with the Velociraptor client's details |
| 365 | await update_agent_with_velociraptor_in_db(session, agent, velociraptor_agent) |
| 366 | agents_added_list.append(velociraptor_agent) |
| 367 | |
| 368 | # Close the session |
| 369 | await session.close() |
| 370 | |
| 371 | logger.info(f"Agents Added List: {agents_added_list}") |
| 372 | return SyncedAgentsResponse( |
| 373 | success=True, |
| 374 | message="Agents synced successfully", |
| 375 | agents_added=agents_added_list, |
| 376 | ) |