update `agent_metadata` table when agents are synced (#31)
taylor_socfortress committed
Jul 15, 2023 at 10:19 UTC
163bf4698549b14fcd4ff3a82e4d3ea6f8e49ac3
1 file changed
+44
backend/app/services/agents/agents.py
+44
@@ -141,6 +141,49 @@ class AgentService:
141
logger.error(f"Failed to create agent: {e}")
142
return None
143
144
+ def update_agent(self, agent: Dict[str, str]) -> Optional[AgentMetadata]:
145
+ """
146
+ Updates an agent in the database.
147
+
148
+ Args:
149
+ agent (dict): A dictionary containing the information of an agent.
150
+
151
+ Returns:
152
+ The agent object if the agent was successfully updated, None otherwise.
153
+ """
154
+ try:
155
+ agent_last_seen = datetime.strptime(
156
+ agent["agent_last_seen"],
157
+ "%Y-%m-%dT%H:%M:%S+00:00",
158
+ ) # Convert to datetime
159
+ except ValueError:
160
+ logger.info(
161
+ f"Invalid format for agent_last_seen: {agent['agent_last_seen']}. Fixing...",
162
+ )
163
+
164
+ agent_last_seen = datetime.strptime(
165
+ "1970-01-01T00:00:00+00:00",
166
+ "%Y-%m-%dT%H:%M:%S+00:00",
167
+ )
168
+
169
+ agent_metadata = db.session.query(AgentMetadata).filter_by(agent_id=agent["agent_id"]).first()
170
+ if agent_metadata is None:
171
+ return None
172
+
173
+ agent_metadata.hostname = agent["agent_name"]
174
+ agent_metadata.ip_address = agent["agent_ip"]
175
+ agent_metadata.os = agent["agent_os"]
176
+ agent_metadata.last_seen = agent_last_seen # Use the datetime object
177
+ agent_metadata.critical_asset = False
178
+ logger.info(f"Agent metadata: {agent_metadata}")
179
+
180
+ try:
181
+ db.session.commit()
182
+ return agent_metadata
183
+ except Exception as e:
184
+ logger.error(f"Failed to update agent: {e}")
185
+ return None
186
+
187
def delete_agent_db(self, agent_id: str) -> Dict[str, Union[str, bool]]:
188
"""
189
Deletes a specific agent from the database using its ID.
@@ -268,6 +311,7 @@ class AgentSyncService:
311
for agent in wazuh_agents_list:
312
agent_info = self.agent_service.get_agent(agent["agent_id"])
313
logger.info(f"Agent info: {agent_info}")
314
+ self.agent_service.update_agent(agent)
315
if "message" in agent_info:
316
self.agent_service.create_agent(agent)
317
agents_added_list.append(agent)