Agent sync changes (#130)
* Handle exception when fetching Velociraptor Agent * Add verification after saving file for velociraptor * Add Velociraptor connector and agent processing * precommit fixes * precommit styling
taylor_socfortress committed
Feb 4, 2024 at 10:47 UTC
366c36ed37ac7bd0077dfe4235febae8f1230c4c
2 files changed
+66
-1
backend/app/agents/services/sync.py
+65
-1
@@ -11,6 +11,7 @@ from app.agents.schema.agents import SyncedAgentsResponse
11
from app.agents.velociraptor.schema.agents import VelociraptorAgent
12
from app.agents.wazuh.schema.agents import WazuhAgent
13
from app.agents.wazuh.schema.agents import WazuhAgentsList
14
+from app.connectors.models import Connectors
15
from app.db.universal_models import Agents
16
17
@@ -103,6 +104,65 @@ def extract_customer_code(customer_code: str):
104
return parts[1] if len(parts) > 1 else None
105
106
107
+async def get_velociraptor_connector(session):
108
+ """
109
+ Retrieves the Velociraptor connector from the database.
110
+
111
+ Args:
112
+ session: The database session.
113
+
114
+ Returns:
115
+ The first result of the query as a scalar value.
116
+ """
117
+ connector_query = select(Connectors).filter(Connectors.connector_name == "Velociraptor")
118
+ result = await session.execute(connector_query)
119
+ return result.scalars().first()
120
+
121
+
122
+async def get_velociraptor_agent(agent_name):
123
+ """
124
+ Retrieves a Velociraptor agent with the specified name.
125
+
126
+ Args:
127
+ agent_name (str): The name of the agent to retrieve.
128
+
129
+ Returns:
130
+ VelociraptorAgent: The retrieved Velociraptor agent, or None if retrieval fails.
131
+ """
132
+ try:
133
+ return await fetch_velociraptor_agent(agent_name)
134
+ except Exception as e:
135
+ logger.error(f"Failed to collect Velociraptor Agent for {agent_name}: {e}")
136
+ return None
137
+
138
+
139
+async def process_velociraptor_agent(session, wazuh_agent):
140
+ """
141
+ Process the Velociraptor agent for a given Wazuh agent.
142
+
143
+ Args:
144
+ session (object): The session object for the connection.
145
+ wazuh_agent (object): The Wazuh agent object.
146
+
147
+ Returns:
148
+ object: The Velociraptor agent object if successful, None otherwise.
149
+ """
150
+ try:
151
+ velociraptor_connector = await get_velociraptor_connector(session)
152
+ if velociraptor_connector.connector_verified:
153
+ velociraptor_agent = await get_velociraptor_agent(wazuh_agent.agent_name)
154
+ else:
155
+ velociraptor_agent = VelociraptorAgent(
156
+ client_id="Unknown",
157
+ client_last_seen="1970-01-01T00:00:00+00:00",
158
+ client_version="Unknown",
159
+ )
160
+ return velociraptor_agent
161
+ except Exception as e:
162
+ logger.error(f"Failed to process agent {wazuh_agent.agent_name}: {e}")
163
+ return None
164
+
165
+
166
async def sync_agents(session: AsyncSession) -> SyncedAgentsResponse:
167
"""
168
Synchronize agents from Wazuh and Velociraptor services.
@@ -124,7 +184,11 @@ async def sync_agents(session: AsyncSession) -> SyncedAgentsResponse:
184
for wazuh_agent in wazuh_agents_list.agents:
185
logger.info(f"Collecting Velociraptor Agent for {wazuh_agent.agent_name}")
186
127
- velociraptor_agent = await fetch_velociraptor_agent(wazuh_agent.agent_name)
187
+ try:
188
+ velociraptor_agent = await process_velociraptor_agent(session, wazuh_agent)
189
+ except Exception as e:
190
+ logger.error(f"Failed to collect Velociraptor Agent for {wazuh_agent.agent_name}: {e}")
191
+ continue
192
193
customer_code = extract_customer_code(wazuh_agent.agent_label)
194
backend/app/connectors/routes.py
+1
@@ -171,6 +171,7 @@ async def upload_yaml_file(connector_id: int, file: UploadFile = File(...), sess
171
try:
172
save_file_result = await ConnectorServices.save_file(file, session=session)
173
if save_file_result:
174
+ await ConnectorServices.verify_connector_by_id(connector_id, session=session)
175
return {"success": True, "message": "File uploaded successfully"}
176
else:
177
raise HTTPException(status_code=500, detail="Failed to upload file")