Create agent.py
taylor_socfortress committed
Jul 10, 2023 at 16:45 UTC
089cf9e36c67149474e4b1e6485b0f68476991cc
1 file changed
+137
backend/app/services/WazuhManager/agent.py
new
+137
@@ -0,0 +1,137 @@
1
+from typing import Dict, Optional, List, Any
2
+from loguru import logger
3
+from app.services.WazuhManager.universal import UniversalService
4
+import requests
5
+
6
+class WazuhHttpRequests:
7
+ """
8
+ Class to handle HTTP requests to the Wazuh API.
9
+ """
10
+ def __init__(self, connector_url: str, wazuh_auth_token: str) -> None:
11
+ """
12
+ Args:
13
+ connector_url (str): The URL of the Wazuh Manager.
14
+ wazuh_auth_token (str): The Wazuh API authentication token.
15
+ """
16
+ self.connector_url = connector_url
17
+ self.wazuh_auth_token = wazuh_auth_token
18
+ self.headers = {"Authorization": f"Bearer {wazuh_auth_token}"}
19
+
20
+ def delete_request(self, endpoint: str, params: Optional[Dict[str, str]] = None) -> Dict[str, bool]:
21
+ """
22
+ Function to handle DELETE requests.
23
+
24
+ Args:
25
+ endpoint (str): The endpoint to make a DELETE request to.
26
+ params (Optional[Dict[str, str]]): Any parameters to pass in the DELETE request.
27
+
28
+ Returns:
29
+ Dict[str, bool]: A dictionary indicating the success of the operation.
30
+ """
31
+ try:
32
+ response = requests.delete(
33
+ f"{self.connector_url}/{endpoint}",
34
+ headers=self.headers,
35
+ params=params,
36
+ verify=False,
37
+ )
38
+ response.raise_for_status()
39
+ logger.info(f"Successfully deleted {endpoint}")
40
+ return {"agentDeleted": True}
41
+
42
+ except Exception as e:
43
+ logger.error(f"Failed to delete {endpoint}: {e}")
44
+ return {"agentDeleted": False}
45
+
46
+class WazuhManagerAgentService:
47
+ """
48
+ A service class that encapsulates the logic for handling agent related operations in Wazuh Manager.
49
+ """
50
+ def __init__(self, universal_service: UniversalService) -> None:
51
+ """
52
+ Args:
53
+ universal_service (UniversalService): The UniversalService instance to use.
54
+ """
55
+ self.universal_service = universal_service
56
+ self.auth_token = universal_service.get_auth_token()
57
+ self.wazuh_http_requests = WazuhHttpRequests(self.universal_service.connector_url, self.auth_token)
58
+
59
+ def collect_agents(self) -> Optional[List[Dict[str, str]]]:
60
+ """
61
+ Collect all agents from Wazuh Manager.
62
+
63
+ Returns:
64
+ Optional[List[Dict[str, str]]]: A list of dictionaries containing agent data, or None on failure.
65
+ """
66
+ logger.info("Collecting Wazuh Agents")
67
+ try:
68
+ agent_data = self._get_agent_data()
69
+ if agent_data is None:
70
+ return None
71
+
72
+ wazuh_agents_list = self._build_agent_list(agent_data)
73
+ return wazuh_agents_list
74
+ except Exception as e:
75
+ logger.error(f"Failed to collect Wazuh Agents: {e}")
76
+ return None
77
+
78
+ def _get_agent_data(self) -> Optional[Dict[str, Any]]:
79
+ """
80
+ Get agent data from Wazuh Manager.
81
+
82
+ Returns:
83
+ Optional[Dict[str, Any]]: A dictionary containing agent data, or None on failure.
84
+ """
85
+ headers = {"Authorization": f"Bearer {self.auth_token}"}
86
+ limit = 1000
87
+ response = requests.get(
88
+ f"{self.universal_service.connector_url}/agents?limit={limit}", headers=headers, verify=False
89
+ )
90
+ if response.status_code == 200:
91
+ return response.json()["data"]["affected_items"]
92
+ else:
93
+ return None
94
+
95
+ def _build_agent_list(self, agent_data: Dict[str, Any]) -> List[Dict[str, str]]:
96
+ """
97
+ Build a list of agent data dictionaries.
98
+
99
+ Args:
100
+ agent_data (Dict[str, Any]): The raw agent data.
101
+
102
+ Returns:
103
+ List[Dict[str, str]]: A list of dictionaries containing agent data.
104
+ """
105
+ wazuh_agents_list = []
106
+ for agent in agent_data:
107
+ os_name = agent.get("os", {}).get("name", "Unknown")
108
+ last_keep_alive = agent.get("lastKeepAlive", "Unknown")
109
+ wazuh_agents_list.append(
110
+ {
111
+ "agent_id": agent["id"],
112
+ "agent_name": agent["name"],
113
+ "agent_ip": agent["ip"],
114
+ "agent_os": os_name,
115
+ "agent_last_seen": last_keep_alive,
116
+ },
117
+ )
118
+ logger.info(f"Collected Wazuh Agent: {agent['name']}")
119
+ return wazuh_agents_list
120
+
121
+ def delete_agent(self, agent_id: str) -> Dict[str, bool]:
122
+ """
123
+ Delete an agent from Wazuh Manager.
124
+
125
+ Args:
126
+ agent_id (str): The id of the agent to be deleted.
127
+
128
+ Returns:
129
+ Dict[str, bool]: A dictionary indicating the success of the operation.
130
+ """
131
+ params = {
132
+ "purge": True,
133
+ "agents_list": [agent_id],
134
+ "status": "all",
135
+ "older_than": "0s",
136
+ }
137
+ return self.wazuh_http_requests.delete_request("agents", params)