| 1 | import asyncio |
| 2 | |
| 3 | from fastapi import HTTPException |
| 4 | from loguru import logger |
| 5 | |
| 6 | from app.agents.schema.agents import AgentModifyResponse |
| 7 | from app.agents.schema.agents import AgentWazuhUpgradeResponse |
| 8 | from app.agents.wazuh.schema.agents import WazuhAgent |
| 9 | from app.agents.wazuh.schema.agents import WazuhAgentsList |
| 10 | from app.connectors.wazuh_manager.utils.universal import send_delete_request |
| 11 | from app.connectors.wazuh_manager.utils.universal import send_get_request |
| 12 | from app.connectors.wazuh_manager.utils.universal import send_put_request |
| 13 | |
| 14 | |
| 15 | async def collect_wazuh_agents() -> WazuhAgentsList: |
| 16 | """ |
| 17 | Collects all agents from Wazuh Manager. |
| 18 | |
| 19 | Returns: |
| 20 | WazuhAgentsList: A list of WazuhAgent objects representing the collected agents. |
| 21 | """ |
| 22 | logger.info("Collecting all agents from Wazuh Manager") |
| 23 | agents_collected = await send_get_request( |
| 24 | endpoint="/agents", |
| 25 | params={"limit": 500}, |
| 26 | ) |
| 27 | total_affected_items = agents_collected.get("data", {}).get("data", {}).get("total_affected_items", 0) |
| 28 | |
| 29 | # If the number of agents is less than total_affected_items, make another request with the limit being total_affected_items |
| 30 | if len(agents_collected.get("data", {}).get("data", {}).get("affected_items", [])) < total_affected_items: |
| 31 | logger.info( |
| 32 | f"Total items: {total_affected_items}.\n" |
| 33 | f"Collected {len(agents_collected.get('data', {}).get('data', {}).get('affected_items', []))} agents.\n" |
| 34 | "Making another request.", |
| 35 | ) |
| 36 | # sleep for 2 seconds before making another request |
| 37 | await asyncio.sleep(2) |
| 38 | agents_collected = await send_get_request( |
| 39 | endpoint="/agents", |
| 40 | params={"limit": total_affected_items}, |
| 41 | ) |
| 42 | |
| 43 | if agents_collected.get("success") is False: |
| 44 | raise HTTPException( |
| 45 | status_code=500, |
| 46 | detail=agents_collected.get("message", "Unknown error"), |
| 47 | ) |
| 48 | try: |
| 49 | if agents_collected.get("success"): |
| 50 | wazuh_agents_list = [] |
| 51 | for agent in agents_collected.get("data", {}).get("data", {}).get("affected_items", []): |
| 52 | os_name = agent.get("os", {}).get("name", "Unknown") |
| 53 | last_keep_alive = agent.get("lastKeepAlive", "Unknown") |
| 54 | agent_group_list = agent.get("group", []) |
| 55 | agent_group = agent_group_list[0] if agent_group_list else "Unknown" |
| 56 | |
| 57 | wazuh_agent = WazuhAgent( |
| 58 | agent_id=agent.get("id", "Unknown"), |
| 59 | agent_name=agent.get("name", "Unknown"), |
| 60 | agent_ip=agent.get("ip", "Unknown"), |
| 61 | agent_os=os_name, |
| 62 | agent_label=agent_group, |
| 63 | agent_last_seen=last_keep_alive, |
| 64 | wazuh_agent_version=agent.get("version", "n/a"), |
| 65 | wazuh_agent_status=agent.get("status", "n/a"), |
| 66 | ) |
| 67 | wazuh_agents_list.append(wazuh_agent) |
| 68 | |
| 69 | return WazuhAgentsList( |
| 70 | agents=wazuh_agents_list, |
| 71 | success=True, |
| 72 | message="Agents collected successfully", |
| 73 | ) |
| 74 | |
| 75 | except (KeyError, IndexError, HTTPException) as e: |
| 76 | # Handle or log the error as needed |
| 77 | logger.error(f"An error occurred: {e}") |
| 78 | raise HTTPException( |
| 79 | status_code=500, |
| 80 | detail=f"Failed to collect agents: {e}", |
| 81 | ) |
| 82 | |
| 83 | except Exception as e: |
| 84 | # Catch-all for other exceptions |
| 85 | logger.error(f"An unexpected error occurred: {e}") |
| 86 | raise HTTPException( |
| 87 | status_code=500, |
| 88 | detail=f"Failed to collect agents: {e}", |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def handle_agent_deletion_response(agent_deleted: dict, agent_id: str): |
| 93 | """ |
| 94 | Handles the response of agent deletion from the Wazuh Manager. |
| 95 | |
| 96 | Args: |
| 97 | agent_deleted (dict): The response of agent deletion. |
| 98 | agent_id (str): The ID of the agent. |
| 99 | |
| 100 | Returns: |
| 101 | AgentModifyResponse: An instance of AgentModifyResponse if the agent is deleted successfully. |
| 102 | |
| 103 | Raises: |
| 104 | HTTPException: If the agent deletion fails, an HTTPException is raised with the appropriate error message. |
| 105 | """ |
| 106 | if agent_deleted["success"]: |
| 107 | return AgentModifyResponse(success=True, message="Agent deleted successfully") |
| 108 | else: |
| 109 | raise HTTPException( |
| 110 | status_code=400, |
| 111 | detail=f"Failed to delete agent {agent_id} from Wazuh Manager: {agent_deleted.get('message', 'Unknown error')}", |
| 112 | ) |
| 113 | |
| 114 | |
| 115 | async def delete_agent_wazuh(agent_id: str) -> AgentModifyResponse: |
| 116 | """Delete agent from Wazuh Manager. |
| 117 | |
| 118 | Args: |
| 119 | agent_id (str): The ID of the agent to be deleted. |
| 120 | |
| 121 | Returns: |
| 122 | AgentModifyResponse: The response indicating the status of the agent deletion. |
| 123 | |
| 124 | Raises: |
| 125 | HTTPException: If there is an HTTP error during the deletion process. |
| 126 | """ |
| 127 | logger.info(f"Deleting agent {agent_id} from Wazuh Manager") |
| 128 | |
| 129 | params = { |
| 130 | "purge": True, |
| 131 | "agents_list": [agent_id], |
| 132 | "status": "all", |
| 133 | "older_than": "0s", |
| 134 | } |
| 135 | |
| 136 | try: |
| 137 | agent_deleted = await send_delete_request(endpoint="/agents", params=params) |
| 138 | return handle_agent_deletion_response(agent_deleted, agent_id) |
| 139 | |
| 140 | except HTTPException as http_e: |
| 141 | # * Catch any HTTPException and re-raise it |
| 142 | raise http_e |
| 143 | |
| 144 | except Exception as e: |
| 145 | # * Catch-all for other exceptions |
| 146 | raise HTTPException( |
| 147 | status_code=500, |
| 148 | detail=f"Failed to delete agent {agent_id} from Wazuh Manager: {e}", |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | def handle_agent_upgrade_response(agent_upgraded: dict) -> AgentWazuhUpgradeResponse: |
| 153 | """ |
| 154 | Handle the response from the agent upgrade request. |
| 155 | |
| 156 | Args: |
| 157 | agent_upgraded (dict): The response from the agent upgrade request. |
| 158 | |
| 159 | Returns: |
| 160 | AgentWazuhUpgradeResponse: The response indicating the status of the agent upgrade. |
| 161 | """ |
| 162 | data = agent_upgraded.get("data", {}).get("data", {}) |
| 163 | total_failed_items = data.get("total_failed_items", 0) |
| 164 | |
| 165 | if total_failed_items == 0: |
| 166 | # Upgrade was successful |
| 167 | return AgentWazuhUpgradeResponse( |
| 168 | success=True, |
| 169 | message=agent_upgraded.get("data", {}).get("message", "Unknown error"), |
| 170 | ) |
| 171 | else: |
| 172 | # Upgrade failed |
| 173 | failed_items = data.get("failed_items", [{}]) |
| 174 | error_message = failed_items[0].get("error", {}).get("message", "Unknown error") |
| 175 | return AgentWazuhUpgradeResponse( |
| 176 | success=False, |
| 177 | message=error_message, |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | async def upgrade_wazuh_agent(agent_id: str) -> AgentWazuhUpgradeResponse: |
| 182 | """Upgrade agent from Wazuh Manager. |
| 183 | |
| 184 | Args: |
| 185 | agent_id (str): The ID of the agent to be upgraded. |
| 186 | |
| 187 | Returns: |
| 188 | AgentWazuhUpgradeResponse: The response indicating the status of the agent upgrade. |
| 189 | |
| 190 | Raises: |
| 191 | HTTPException: If there is an HTTP error during the upgrade process. |
| 192 | """ |
| 193 | logger.info(f"Upgrading agent {agent_id} from Wazuh Manager") |
| 194 | |
| 195 | params = { |
| 196 | "agents_list": [agent_id], |
| 197 | } |
| 198 | |
| 199 | try: |
| 200 | agent_upgraded = await send_put_request(endpoint="/agents/upgrade", data=None, params=params) |
| 201 | logger.info(f"Agent upgrade response: {agent_upgraded}") |
| 202 | return handle_agent_upgrade_response(agent_upgraded) |
| 203 | |
| 204 | except HTTPException as http_e: |
| 205 | # * Catch any HTTPException and re-raise it |
| 206 | raise http_e |
| 207 | |
| 208 | except Exception as e: |
| 209 | # * Catch-all for other exceptions |
| 210 | raise HTTPException( |
| 211 | status_code=500, |
| 212 | detail=f"Failed to upgrade agent {agent_id} from Wazuh Manager: {e}", |
| 213 | ) |