make agents service more modular and readable (#35)
taylor_socfortress committed
Jul 15, 2023 at 13:02 UTC
d215606328088469c1e3670b414fe484b9de7b35
2 files changed
+77
-130
backend/app/routes/agents.py
+13
-8
@@ -3,6 +3,7 @@ from typing import Any
3
from flask import Blueprint
4
from flask import jsonify
5
6
+from app.models.agents import agent_metadata_schema
7
from app.services.agents.agents import AgentService
8
from app.services.agents.agents import AgentSyncService
9
from app.services.WazuhManager.agent import WazuhManagerAgentService
@@ -21,7 +22,7 @@ def get_agents() -> Any:
22
"""
23
service = AgentService()
24
agents = service.get_all_agents()
24
- return agents
25
+ return jsonify(agents)
26
27
28
@bp.route("/agents/<agent_id>", methods=["GET"])
@@ -35,7 +36,11 @@ def get_agent(agent_id: str) -> Any:
36
"""
37
service = AgentService()
38
agent = service.get_agent(agent_id=agent_id)
38
- return agent
39
+ if agent is None:
40
+ return jsonify({"message": "Agent not found", "success": False}), 404
41
+ else:
42
+ agent_dict = agent_metadata_schema.dump(agent)
43
+ return jsonify(agent_dict)
44
45
46
@bp.route("/agents/<agent_id>/critical", methods=["POST"])
@@ -48,8 +53,8 @@ def mark_as_critical(agent_id: str) -> Any:
53
json: A JSON response containing the updated agent information after being marked as critical.
54
"""
55
service = AgentService()
51
- result = service.mark_agent_as_critical(agent_id=agent_id)
52
- return result
56
+ result = service.mark_agent_criticality(agent_id=agent_id, critical=True)
57
+ return jsonify(result)
58
59
60
@bp.route("/agents/<agent_id>/noncritical", methods=["POST"])
@@ -62,8 +67,8 @@ def unmark_agent_critical(agent_id: str) -> Any:
67
json: A JSON response containing the updated agent information after being unmarked as critical.
68
"""
69
service = AgentService()
65
- result = service.mark_agent_as_non_critical(agent_id=agent_id)
66
- return result
70
+ result = service.mark_agent_criticality(agent_id=agent_id, critical=False)
71
+ return jsonify(result)
72
73
74
@bp.route("/agents/sync", methods=["POST"])
@@ -94,7 +99,7 @@ def delete_agent(agent_id: str) -> Any:
99
agent_service = WazuhManagerAgentService(universal_service)
100
agent_service.delete_agent(agent_id=agent_id)
101
97
- return result
102
+ return jsonify(result)
103
104
105
@bp.route("/agents/<agent_id>/vulnerabilities", methods=["GET"])
@@ -112,4 +117,4 @@ def get_agent_vulnerabilities(agent_id: str) -> Any:
117
agent_vulnerabilities = vulnerability_service.agent_vulnerabilities(
118
agent_id=agent_id,
119
)
115
- return agent_vulnerabilities
120
+ return jsonify(agent_vulnerabilities)
backend/app/services/agents/agents.py
+64
-122
@@ -1,4 +1,4 @@
1
-# agents.py
1
+# Here is the improved version of the code:
2
from datetime import datetime
3
from typing import Dict
4
from typing import List
@@ -24,6 +24,22 @@ class AgentService:
24
A service class that encapsulates the logic for managing agents.
25
"""
26
27
+ def parse_date(self, date_string: str) -> datetime:
28
+ """
29
+ Parses a date string into a datetime object.
30
+
31
+ Args:
32
+ date_string (str): The date string to parse.
33
+
34
+ Returns:
35
+ datetime: The parsed datetime object.
36
+ """
37
+ try:
38
+ return datetime.strptime(date_string, "%Y-%m-%dT%H:%M:%S+00:00")
39
+ except ValueError:
40
+ logger.info(f"Invalid format for date: {date_string}. Using the epoch time as default.")
41
+ return datetime.strptime("1970-01-01T00:00:00+00:00", "%Y-%m-%dT%H:%M:%S+00:00")
42
+
43
def get_all_agents(self) -> List[Dict[str, Union[str, bool]]]:
44
"""
45
Retrieves all agents from the database.
@@ -34,7 +50,7 @@ class AgentService:
50
agents = db.session.query(AgentMetadata).all()
51
return agent_metadatas_schema.dump(agents)
52
37
- def get_agent(self, agent_id: str) -> Dict[str, Union[str, bool]]:
53
+ def get_agent(self, agent_id: str) -> Optional[AgentMetadata]:
54
"""
55
Retrieves a specific agent from the database using its ID.
56
@@ -42,153 +58,86 @@ class AgentService:
58
agent_id (str): The ID of the agent to retrieve.
59
60
Returns:
45
- dict: A dictionary representing the serialized data of the agent if found, otherwise a message indicating
46
- that the agent was not found.
61
+ AgentMetadata: The agent object if found, otherwise None.
62
"""
48
- agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
49
- if agent is None:
50
- return {"message": f"Agent with ID {agent_id} not found"}
51
- return agent_metadata_schema.dump(agent)
63
+ return db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
64
53
- def mark_agent_as_critical(self, agent_id: str) -> Dict[str, Union[str, bool]]:
65
+ def mark_agent_criticality(self, agent_id: str, critical: bool) -> Dict[str, Union[str, bool]]:
66
"""
55
- Marks a specific agent as critical.
67
+ Marks a specific agent as critical or non-critical.
68
69
Args:
58
- agent_id (str): The ID of the agent to mark as critical.
70
+ agent_id (str): The ID of the agent to mark.
71
+ critical (bool): Whether to mark the agent as critical.
72
73
Returns:
61
- dict: A dictionary representing a success message if the operation was successful, otherwise an error
62
- message.
74
+ dict: A dictionary representing a success message if the operation was successful, otherwise an error message.
75
"""
64
- agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
76
+ agent = self.get_agent(agent_id)
77
78
if agent is None:
79
return {"message": f"Agent {agent_id} not found", "success": False}
80
69
- agent.mark_as_critical()
70
- agent_details = agent_metadata_schema.dump(agent)
71
- if agent_details["critical_asset"] is False:
72
- return {
73
- "message": f"Agent {agent_id} failed to mark agent as critical",
74
- "success": False,
75
- }
76
- return {"message": f"Agent {agent_id} marked as critical", "success": True}
77
-
78
- def mark_agent_as_non_critical(self, agent_id: str) -> Dict[str, Union[str, bool]]:
79
- """
80
- Marks a specific agent as non-critical.
81
-
82
- Args:
83
- agent_id (str): The ID of the agent to mark as non-critical.
84
-
85
- Returns:
86
- dict: A dictionary representing a success message if the operation was successful, otherwise an error
87
- message.
88
- """
89
- agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
90
-
91
- if agent is None:
92
- return {"message": f"Agent {agent_id} not found", "success": False}
81
+ if critical:
82
+ agent.mark_as_critical()
83
+ else:
84
+ agent.mark_as_non_critical()
85
94
- agent.mark_as_non_critical()
86
agent_details = agent_metadata_schema.dump(agent)
96
- if agent_details["critical_asset"] is True:
87
+ if agent_details["critical_asset"] is not critical:
88
return {
98
- "message": f"Agent {agent_id} failed to mark agent as non-critical",
89
+ "message": f"Agent {agent_id} failed to mark agent as {'critical' if critical else 'non-critical'}",
90
"success": False,
91
}
101
- return {"message": f"Agent {agent_id} marked as non-critical", "success": True}
92
+ return {"message": f"Agent {agent_id} marked as {'critical' if critical else 'non-critical'}", "success": True}
93
103
- def create_agent(self, agent: Dict[str, str]) -> Optional[AgentMetadata]:
94
+ def create_or_update_agent(self, agent: Dict[str, str]) -> Optional[AgentMetadata]:
95
"""
105
- Creates a new agent in the database.
96
+ Creates or updates an agent in the database.
97
98
Args:
99
agent (dict): A dictionary containing the information of an agent.
100
101
Returns:
111
- The agent object if the agent was successfully created, None otherwise.
102
+ The agent object if the agent was successfully created or updated, None otherwise.
103
"""
113
- try:
114
- agent_last_seen = datetime.strptime(
115
- agent["agent_last_seen"],
116
- "%Y-%m-%dT%H:%M:%S+00:00",
117
- ) # Convert to datetime
118
- except ValueError:
119
- logger.info(
120
- f"Invalid format for agent_last_seen: {agent['agent_last_seen']}. Fixing...",
121
- )
122
- agent_last_seen = datetime.strptime(
123
- "1970-01-01T00:00:00+00:00",
124
- "%Y-%m-%dT%H:%M:%S+00:00",
125
- ) # Use the epoch time as default
104
+ agent_last_seen = self.parse_date(agent["agent_last_seen"])
105
+
106
+ existing_agent = self.get_agent(agent["agent_id"])
107
+ if existing_agent is not None:
108
+ existing_agent.hostname = agent["agent_name"]
109
+ existing_agent.ip_address = agent["agent_ip"]
110
+ existing_agent.os = agent["agent_os"]
111
+ existing_agent.last_seen = agent_last_seen
112
+ existing_agent.client_id = agent["client_id"]
113
+ existing_agent.client_last_seen = agent["client_last_seen"]
114
+ try:
115
+ db.session.commit()
116
+ return existing_agent
117
+ except Exception as e:
118
+ logger.error(f"Failed to update agent: {e}")
119
+ return None
120
127
- agent_metadata = AgentMetadata(
121
+ new_agent = AgentMetadata(
122
agent_id=agent["agent_id"],
123
hostname=agent["agent_name"],
124
ip_address=agent["agent_ip"],
125
os=agent["agent_os"],
132
- last_seen=agent_last_seen, # Use the datetime object
126
+ last_seen=agent_last_seen,
127
critical_asset=False,
128
client_id=agent["client_id"],
129
client_last_seen=agent["client_last_seen"],
130
)
137
- logger.info(f"Agent metadata: {agent_metadata}")
131
+ logger.info(f"Agent metadata: {new_agent}")
132
133
try:
140
- db.session.add(agent_metadata)
134
+ db.session.add(new_agent)
135
db.session.commit()
142
- return agent_metadata
136
+ return new_agent
137
except Exception as e:
138
logger.error(f"Failed to create agent: {e}")
139
return None
140
147
- def update_agent(self, agent: Dict[str, str]) -> Optional[AgentMetadata]:
148
- """
149
- Updates an agent in the database.
150
-
151
- Args:
152
- agent (dict): A dictionary containing the information of an agent.
153
-
154
- Returns:
155
- The agent object if the agent was successfully updated, None otherwise.
156
- """
157
- try:
158
- agent_last_seen = datetime.strptime(
159
- agent["agent_last_seen"],
160
- "%Y-%m-%dT%H:%M:%S+00:00",
161
- ) # Convert to datetime
162
- except ValueError:
163
- logger.info(
164
- f"Invalid format for agent_last_seen: {agent['agent_last_seen']}. Fixing...",
165
- )
166
-
167
- agent_last_seen = datetime.strptime(
168
- "1970-01-01T00:00:00+00:00",
169
- "%Y-%m-%dT%H:%M:%S+00:00",
170
- )
171
-
172
- agent_metadata = db.session.query(AgentMetadata).filter_by(agent_id=agent["agent_id"]).first()
173
- if agent_metadata is None:
174
- return None
175
-
176
- agent_metadata.hostname = agent["agent_name"]
177
- agent_metadata.ip_address = agent["agent_ip"]
178
- agent_metadata.os = agent["agent_os"]
179
- agent_metadata.last_seen = agent_last_seen # Use the datetime object
180
- agent_metadata.critical_asset = False
181
- agent_metadata.client_id = agent["client_id"]
182
- agent_metadata.client_last_seen = agent["client_last_seen"]
183
- logger.info(f"Agent metadata: {agent_metadata}")
184
-
185
- try:
186
- db.session.commit()
187
- return agent_metadata
188
- except Exception as e:
189
- logger.error(f"Failed to update agent: {e}")
190
- return None
191
-
141
def delete_agent_db(self, agent_id: str) -> Dict[str, Union[str, bool]]:
142
"""
143
Deletes a specific agent from the database using its ID.
@@ -197,10 +146,9 @@ class AgentService:
146
agent_id (str): The ID of the agent to delete.
147
148
Returns:
200
- dict: A dictionary representing a success message if the operation was successful, otherwise an error
201
- message.
149
+ dict: A dictionary representing a success message if the operation was successful, otherwise an error message.
150
"""
203
- agent = db.session.query(AgentMetadata).filter_by(agent_id=agent_id).first()
151
+ agent = self.get_agent(agent_id)
152
if agent is None:
153
return {"message": f"Agent with ID {agent_id} not found", "success": False}
154
try:
@@ -235,10 +183,7 @@ class AgentService:
183
return client_id, client_last_seen
184
except Exception as e:
185
logger.error(f"Failed to get last seen at from Velociraptor. Setting to default time. Error: {e}")
238
- client_last_seen = datetime.strptime(
239
- "1970-01-01T00:00:00+00:00",
240
- "%Y-%m-%dT%H:%M:%S+00:00",
241
- )
186
+ client_last_seen = self.parse_date("1970-01-01T00:00:00+00:00")
187
return client_id, client_last_seen
188
189
@@ -337,19 +282,16 @@ class AgentSyncService:
282
"success": False,
283
}
284
340
- logger.info(f"Collected {wazuh_agents_list} Wazuh Agents")
285
+ logger.info(f"Collected {len(wazuh_agents_list)} Wazuh Agents")
286
287
agents_added_list = []
288
for agent in wazuh_agents_list:
344
- agent_info = self.agent_service.get_agent(agent["agent_id"])
345
- logger.info(f"Agent info: {agent_info}")
289
client_id, client_last_seen = self.agent_service.get_velo_metadata(agent["agent_name"])
290
agent["client_id"] = client_id
291
agent["client_last_seen"] = client_last_seen
349
- self.agent_service.update_agent(agent)
350
- if "message" in agent_info:
351
- self.agent_service.create_agent(agent)
352
- agents_added_list.append(agent)
292
+ agent_obj = self.agent_service.create_or_update_agent(agent)
293
+ if agent_obj is not None:
294
+ agents_added_list.append(agent_metadata_schema.dump(agent_obj))
295
296
return {
297
"message": "Successfully synced agents.",