Create agents.py
taylor_socfortress committed
Jul 10, 2023 at 16:47 UTC
b1f604a2997a05f1b5d6d0e1927adf369136be62
1 file changed
+257
backend/app/services/agents/agents.py
new
+257
@@ -0,0 +1,257 @@
1
+# services.py
2
+from app.models.agents import (
3
+ AgentMetadata,
4
+ agent_metadata_schema,
5
+ agent_metadatas_schema,
6
+)
7
+from app import db
8
+from datetime import datetime
9
+import requests
10
+from loguru import logger
11
+from app.models.connectors import connector_factory, Connector, WazuhManagerConnector
12
+
13
+
14
+class AgentService:
15
+ """
16
+ A service class that encapsulates the logic for managing agents.
17
+ """
18
+
19
+ def get_all_agents(self):
20
+ """
21
+ Retrieves all agents from the database.
22
+
23
+ Returns:
24
+ List[dict]: A list of dictionaries where each dictionary represents the serialized data of an agent.
25
+ """
26
+ agents = db.session.query(AgentMetadata).all()
27
+ return agent_metadatas_schema.dump(agents)
28
+
29
+ def get_agent(self, agent_id):
30
+ """
31
+ Retrieves a specific agent from the database using its ID.
32
+
33
+ Args:
34
+ agent_id (str): The ID of the agent to retrieve.
35
+
36
+ Returns:
37
+ dict: A dictionary representing the serialized data of the agent if found, otherwise a message indicating that the agent was not found.
38
+ """
39
+ agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
40
+ if agent is None:
41
+ return {"message": f"Agent with ID {agent_id} not found"}
42
+ return agent_metadata_schema.dump(agent)
43
+
44
+ def mark_agent_as_critical(self, agent_id):
45
+ """
46
+ Marks a specific agent as critical.
47
+
48
+ Args:
49
+ agent_id (str): The ID of the agent to mark as critical.
50
+
51
+ Returns:
52
+ dict: A dictionary representing a success message if the operation was successful, otherwise an error message.
53
+ """
54
+ agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
55
+
56
+ if agent is None:
57
+ return {"message": f"Agent {agent_id} not found", "success": False}
58
+
59
+ agent.mark_as_critical()
60
+ agent_details = agent_metadata_schema.dump(agent)
61
+ if agent_details["critical_asset"] == False:
62
+ return {
63
+ "message": f"Agent {agent_id} failed to mark agent as critical",
64
+ "success": False,
65
+ }
66
+ return {"message": f"Agent {agent_id} marked as critical", "success": True}
67
+
68
+ def mark_agent_as_non_critical(self, agent_id):
69
+ """
70
+ Marks a specific agent as non-critical.
71
+
72
+ Args:
73
+ agent_id (str): The ID of the agent to mark as non-critical.
74
+
75
+ Returns:
76
+ dict: A dictionary representing a success message if the operation was successful, otherwise an error message.
77
+ """
78
+ agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
79
+
80
+ if agent is None:
81
+ return {"message": f"Agent {agent_id} not found", "success": False}
82
+
83
+ agent.mark_as_non_critical()
84
+ agent_details = agent_metadata_schema.dump(agent)
85
+ if agent_details["critical_asset"] == True:
86
+ return {
87
+ "message": f"Agent {agent_id} failed to mark agent as non-critical",
88
+ "success": False,
89
+ }
90
+ return {"message": f"Agent {agent_id} marked as non-critical", "success": True}
91
+
92
+ def create_agent(self, agent):
93
+ """
94
+ Creates a new agent in the database.
95
+
96
+ Args:
97
+ agent (dict): A dictionary containing the information of an agent.
98
+
99
+ Returns:
100
+ The agent object if the agent was successfully created, None otherwise.
101
+ """
102
+ try:
103
+ agent_last_seen = datetime.strptime(
104
+ agent["agent_last_seen"], "%Y-%m-%dT%H:%M:%S+00:00"
105
+ ) # Convert to datetime
106
+ except ValueError:
107
+ logger.info(
108
+ f"Invalid format for agent_last_seen: {agent['agent_last_seen']}. Fixing..."
109
+ )
110
+ agent_last_seen = datetime.strptime(
111
+ "1970-01-01T00:00:00+00:00", "%Y-%m-%dT%H:%M:%S+00:00"
112
+ ) # Use the epoch time as default
113
+
114
+ agent_metadata = AgentMetadata(
115
+ agent_id=agent["agent_id"],
116
+ hostname=agent["agent_name"],
117
+ ip_address=agent["agent_ip"],
118
+ os=agent["agent_os"],
119
+ last_seen=agent_last_seen, # Use the datetime object
120
+ critical_asset=False,
121
+ )
122
+ logger.info(f"Agent metadata: {agent_metadata}")
123
+
124
+ try:
125
+ db.session.add(agent_metadata)
126
+ db.session.commit()
127
+ return agent_metadata
128
+ except Exception as e:
129
+ logger.error(f"Failed to create agent: {e}")
130
+ return None
131
+
132
+ def delete_agent_db(self, agent_id):
133
+ """
134
+ Deletes a specific agent from the database using its ID.
135
+
136
+ Args:
137
+ agent_id (str): The ID of the agent to delete.
138
+
139
+ Returns:
140
+ dict: A dictionary representing a success message if the operation was successful, otherwise an error message.
141
+ """
142
+ agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
143
+ if agent is None:
144
+ return {"message": f"Agent with ID {agent_id} not found", "success": False}
145
+ try:
146
+ db.session.delete(agent)
147
+ db.session.commit()
148
+ return {"message": f"Agent with ID {agent_id} deleted", "success": True}
149
+ except Exception as e:
150
+ logger.error(f"Failed to delete agent: {e}")
151
+ return {
152
+ "message": f"Failed to delete agent with ID {agent_id}",
153
+ "success": False,
154
+ }
155
+
156
+
157
+class AgentSyncService:
158
+ def __init__(self):
159
+ self.agent_service = AgentService()
160
+
161
+ def collect_wazuh_details(self, connector_name: str):
162
+ """
163
+ Collects the information of all Wazuh API credentials using the WazuhIndexerConnector class details.
164
+
165
+ Returns:
166
+ tuple: A tuple containing the connection URL, username, and password.
167
+ """
168
+ connector_instance = connector_factory.create(connector_name, connector_name)
169
+ connection_successful = connector_instance.verify_connection()
170
+ if connection_successful:
171
+ connection_details = Connector.get_connector_info_from_db(connector_name)
172
+ return (
173
+ connection_details.get("connector_url"),
174
+ connection_details.get("connector_username"),
175
+ connection_details.get("connector_password"),
176
+ )
177
+ else:
178
+ return None, None, None
179
+
180
+ def collect_wazuh_agents(self, connection_url: str, wazuh_auth_token: str):
181
+ """
182
+ Collects the information of all agents from the Wazuh API.
183
+
184
+ Returns:
185
+ list: A list containing the information of all Wazuh agents.
186
+ """
187
+ logger.info("Collecting Wazuh Agents")
188
+ try:
189
+ headers = {"Authorization": f"Bearer {wazuh_auth_token}"}
190
+ limit = 1000
191
+ agents_collected = requests.get(
192
+ f"{connection_url}/agents?limit={limit}", headers=headers, verify=False
193
+ )
194
+ if agents_collected.status_code == 200:
195
+ wazuh_agents_list = []
196
+ for agent in agents_collected.json()["data"]["affected_items"]:
197
+ os_name = agent.get("os", {}).get("name", "Unknown")
198
+ last_keep_alive = agent.get("lastKeepAlive", "Unknown")
199
+ wazuh_agents_list.append(
200
+ {
201
+ "agent_id": agent["id"],
202
+ "agent_name": agent["name"],
203
+ "agent_ip": agent["ip"],
204
+ "agent_os": os_name,
205
+ "agent_last_seen": last_keep_alive,
206
+ },
207
+ )
208
+ logger.info(f"Collected Wazuh Agent: {agent['name']}")
209
+ return wazuh_agents_list
210
+ else:
211
+ return None
212
+ except Exception as e:
213
+ logger.error(f"Failed to collect Wazuh Agents: {e}")
214
+ return None
215
+
216
+ def sync_agents(self):
217
+ (
218
+ connection_url,
219
+ connection_username,
220
+ connection_password,
221
+ ) = self.collect_wazuh_details("Wazuh-Manager")
222
+ if connection_url is None:
223
+ return {
224
+ "message": "Failed to get Wazuh-Manager API details",
225
+ "success": False,
226
+ }
227
+
228
+ wazuh_manager_connector = WazuhManagerConnector("Wazuh-Manager")
229
+ wazuh_auth_token = wazuh_manager_connector.get_auth_token()
230
+ if wazuh_auth_token is None:
231
+ return {
232
+ "message": "Failed to get Wazuh-Manager API Auth Token",
233
+ "success": False,
234
+ }
235
+
236
+ wazuh_agents_list = self.collect_wazuh_agents(connection_url, wazuh_auth_token)
237
+ if wazuh_agents_list is None:
238
+ return {
239
+ "message": "Failed to collect Wazuh-Manager Agents",
240
+ "success": False,
241
+ }
242
+
243
+ logger.info(f"Collected {wazuh_agents_list} Wazuh Agents")
244
+
245
+ agents_added_list = []
246
+ for agent in wazuh_agents_list:
247
+ agent_info = self.agent_service.get_agent(agent["agent_id"])
248
+ logger.info(f"Agent info: {agent_info}")
249
+ if "message" in agent_info:
250
+ self.agent_service.create_agent(agent)
251
+ agents_added_list.append(agent)
252
+
253
+ return {
254
+ "message": "Successfully synced agents.",
255
+ "success": True,
256
+ "agents_added": agents_added_list,
257
+ }