Velo dif org artifact collect (#257)
* refactor: Collect Velociraptor artifacts and flows per organization * precommit fixes
taylor_socfortress committed
Jun 27, 2024 at 14:04 UTC
3f6633ac695be0d1325c3e294c9788c03f51b958
4 files changed
+96
-19
backend/app/connectors/velociraptor/routes/flows.py
+75
-2
@@ -56,6 +56,74 @@ async def get_velociraptor_id(session: AsyncSession, hostname: str) -> str:
56
return agent.velociraptor_id
57
58
59
+async def get_velociraptor_org(session: AsyncSession, hostname: str) -> str:
60
+ """
61
+ Retrieves the velociraptor_org associated with the given hostname.
62
+
63
+ Args:
64
+ session (AsyncSession): The database session.
65
+ hostname (str): The hostname of the agent.
66
+
67
+ Returns:
68
+ str: The velociraptor_org associated with the hostname.
69
+
70
+ Raises:
71
+ HTTPException: If the agent with the given hostname is not found or if the velociraptor_org is not available.
72
+ """
73
+ logger.info(f"Getting velociraptor_org from hostname {hostname}")
74
+ result = await session.execute(select(Agents).filter(Agents.hostname == hostname))
75
+ agent = result.scalars().first()
76
+
77
+ if not agent:
78
+ raise HTTPException(
79
+ status_code=404,
80
+ detail=f"Agent with hostname {hostname} not found",
81
+ )
82
+
83
+ if agent.velociraptor_org is None:
84
+ raise HTTPException(
85
+ status_code=404,
86
+ detail=f"Velociraptor ORG for hostname {hostname} is not available",
87
+ )
88
+
89
+ logger.info(f"velociraptor_org for hostname {hostname} is {agent.velociraptor_org}")
90
+ return agent.velociraptor_org
91
+
92
+
93
+async def get_velociraptor_org_via_client_id(session: AsyncSession, client_id: str) -> str:
94
+ """
95
+ Retrieves the velociraptor_org associated with the given hostname.
96
+
97
+ Args:
98
+ session (AsyncSession): The database session.
99
+ hostname (str): The hostname of the agent.
100
+
101
+ Returns:
102
+ str: The velociraptor_org associated with the hostname.
103
+
104
+ Raises:
105
+ HTTPException: If the agent with the given hostname is not found or if the velociraptor_org is not available.
106
+ """
107
+ logger.info(f"Getting velociraptor_org from client id {client_id}")
108
+ result = await session.execute(select(Agents).filter(Agents.velociraptor_id == client_id))
109
+ agent = result.scalars().first()
110
+
111
+ if not agent:
112
+ raise HTTPException(
113
+ status_code=404,
114
+ detail=f"Agent with client id {client_id} not found",
115
+ )
116
+
117
+ if agent.velociraptor_org is None:
118
+ raise HTTPException(
119
+ status_code=404,
120
+ detail=f"Velociraptor ORG for hostname {client_id} is not available",
121
+ )
122
+
123
+ logger.info(f"velociraptor_org for hostname {client_id} is {agent.velociraptor_org}")
124
+ return agent.velociraptor_org
125
+
126
+
127
@velociraptor_flows_router.get(
128
"/{hostname}",
129
response_model=FlowResponse,
@@ -79,8 +147,12 @@ async def get_all_flows_for_hostname(
147
logger.info(f"Fetching all flows for hostname {hostname}")
148
149
velociraptor_id = await get_velociraptor_id(session, hostname)
150
+ velociraptor_org = await get_velociraptor_org(
151
+ session,
152
+ hostname,
153
+ )
154
logger.info(f"velociraptor_id for hostname {hostname} is {velociraptor_id}")
83
- return await get_flows(velociraptor_id)
155
+ return await get_flows(velociraptor_id, velociraptor_org)
156
157
158
@velociraptor_flows_router.post(
@@ -91,6 +163,7 @@ async def get_all_flows_for_hostname(
163
)
164
async def retrieve_flow(
165
retrieve_flow_request: RetrieveFlowRequest,
166
+ session: AsyncSession = Depends(get_db),
167
) -> CollectArtifactResponse:
168
"""
169
Retrieve ran flows for a specific host.
@@ -103,4 +176,4 @@ async def retrieve_flow(
176
CollectArtifactResponse: The response containing the retrieved flows.
177
"""
178
logger.info(f"Fetching flow for flow_id {retrieve_flow_request.session_id}")
106
- return await get_flow(retrieve_flow_request)
179
+ return await get_flow(retrieve_flow_request, await get_velociraptor_org_via_client_id(session, retrieve_flow_request.client_id))
backend/app/connectors/velociraptor/services/artifacts.py
+9
-6
@@ -125,7 +125,7 @@ async def run_artifact_collection(
125
f"FROM scope()"
126
),
127
)
128
- flow = velociraptor_service.execute_query(query)
128
+ flow = velociraptor_service.execute_query(query, org_id=collect_artifact_body.velociraptor_org)
129
logger.info(f"Successfully ran artifact collection on {flow}")
130
131
artifact_key = get_artifact_key(analyzer_body=collect_artifact_body)
@@ -133,12 +133,13 @@ async def run_artifact_collection(
133
flow_id = flow["results"][0][artifact_key]["flow_id"]
134
logger.info(f"Extracted flow_id: {flow_id}")
135
136
- completed = velociraptor_service.watch_flow_completion(flow_id)
136
+ completed = velociraptor_service.watch_flow_completion(flow_id, org_id=collect_artifact_body.velociraptor_org)
137
logger.info(f"Successfully watched flow completion on {completed}")
138
139
results = velociraptor_service.read_collection_results(
140
client_id=collect_artifact_body.velociraptor_id,
141
flow_id=flow_id,
142
+ org_id=collect_artifact_body.velociraptor_org,
143
artifact=collect_artifact_body.artifact_name,
144
)
145
@@ -186,7 +187,7 @@ async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResp
187
"FROM scope()"
188
),
189
)
189
- flow = velociraptor_service.execute_query(query)
190
+ flow = velociraptor_service.execute_query(query, org_id=run_command_body.velociraptor_org)
191
logger.info(f"Successfully ran artifact collection on {flow}")
192
193
artifact_key = get_artifact_key(analyzer_body=run_command_body)
@@ -194,12 +195,13 @@ async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResp
195
flow_id = flow["results"][0][artifact_key]["flow_id"]
196
logger.info(f"Extracted flow_id: {flow_id}")
197
197
- completed = velociraptor_service.watch_flow_completion(flow_id)
198
+ completed = velociraptor_service.watch_flow_completion(flow_id, org_id=run_command_body.velociraptor_org)
199
logger.info(f"Successfully watched flow completion on {completed}")
200
201
results = velociraptor_service.read_collection_results(
202
client_id=run_command_body.velociraptor_id,
203
flow_id=flow_id,
204
+ org_id=run_command_body.velociraptor_org,
205
artifact=run_command_body.artifact_name,
206
)
207
@@ -250,7 +252,7 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse
252
"FROM scope()"
253
),
254
)
253
- flow = velociraptor_service.execute_query(query)
255
+ flow = velociraptor_service.execute_query(query, org_id=quarantine_body.velociraptor_org)
256
logger.info(f"Successfully ran artifact collection on {flow}")
257
258
artifact_key = get_artifact_key(analyzer_body=quarantine_body)
@@ -258,12 +260,13 @@ async def quarantine_host(quarantine_body: QuarantineBody) -> QuarantineResponse
260
flow_id = flow["results"][0][artifact_key]["flow_id"]
261
logger.info(f"Extracted flow_id: {flow_id}")
262
261
- completed = velociraptor_service.watch_flow_completion(flow_id)
263
+ completed = velociraptor_service.watch_flow_completion(flow_id, org_id=quarantine_body.velociraptor_org)
264
logger.info(f"Successfully watched flow completion on {completed}")
265
266
results = velociraptor_service.read_collection_results(
267
client_id=quarantine_body.velociraptor_id,
268
flow_id=flow_id,
269
+ org_id=quarantine_body.velociraptor_org,
270
artifact=quarantine_body.artifact_name,
271
)
272
backend/app/connectors/velociraptor/services/flows.py
+4
-4
@@ -21,7 +21,7 @@ def create_query(query: str) -> str:
21
return query
22
23
24
-async def get_flows(velociraptor_id: str) -> FlowResponse:
24
+async def get_flows(velociraptor_id: str, velociraptor_org: str = "root") -> FlowResponse:
25
"""
26
Get all artifacts from Velociraptor.
27
@@ -33,7 +33,7 @@ async def get_flows(velociraptor_id: str) -> FlowResponse:
33
query = create_query(
34
f"SELECT * FROM flows(client_id='{velociraptor_id}')",
35
)
36
- all_flows = velociraptor_service.execute_query(query)
36
+ all_flows = velociraptor_service.execute_query(query, org_id=velociraptor_org)
37
logger.info(f"all_flows: {all_flows}")
38
flows = [FlowClientSession(**flow) for flow in all_flows["results"]]
39
logger.info(f"flows: {flows}")
@@ -59,7 +59,7 @@ async def get_flows(velociraptor_id: str) -> FlowResponse:
59
)
60
61
62
-async def get_flow(retrieve_flow_request: RetrieveFlowRequest):
62
+async def get_flow(retrieve_flow_request: RetrieveFlowRequest, velociraptor_org: str = "root"):
63
"""
64
Get all artifacts from Velociraptor.
65
@@ -71,7 +71,7 @@ async def get_flow(retrieve_flow_request: RetrieveFlowRequest):
71
query = create_query(
72
f"SELECT * FROM flow_results(client_id='{retrieve_flow_request.client_id}', flow_id='{retrieve_flow_request.session_id}')",
73
)
74
- flow_results = velociraptor_service.execute_query(query)
74
+ flow_results = velociraptor_service.execute_query(query, org_id=velociraptor_org)
75
logger.info(f"flow_results: {flow_results}")
76
try:
77
if flow_results["success"]:
backend/app/connectors/velociraptor/utils/universal.py
+8
-7
@@ -141,7 +141,7 @@ class UniversalService:
141
142
# ! Modify this to use AsyncSessionLocal End
143
144
- def create_vql_request(self, vql: str):
144
+ def create_vql_request(self, vql: str, org_id: str = "root"):
145
"""
146
Creates a VQLCollectorArgs object with given VQL query.
147
@@ -153,6 +153,7 @@ class UniversalService:
153
"""
154
return api_pb2.VQLCollectorArgs(
155
max_wait=1,
156
+ org_id=org_id,
157
Query=[
158
api_pb2.VQLRequest(
159
Name="VQLRequest",
@@ -161,7 +162,7 @@ class UniversalService:
162
],
163
)
164
164
- def execute_query(self, vql: str):
165
+ def execute_query(self, vql: str, org_id: str = "root"):
166
"""
167
Executes a VQL query and returns the results.
168
@@ -173,7 +174,7 @@ class UniversalService:
174
"""
175
logger.info(f"Executing query: {vql}")
176
176
- client_request = self.create_vql_request(vql)
177
+ client_request = self.create_vql_request(vql, org_id)
178
179
try:
180
results = []
@@ -202,7 +203,7 @@ class UniversalService:
203
logger.error(f"Failed to execute query: {e}")
204
raise HTTPException(status_code=500, detail=f"Failed to execute query: {e}")
205
205
- def watch_flow_completion(self, flow_id: str):
206
+ def watch_flow_completion(self, flow_id: str, org_id: str = "root"):
207
"""
208
Watch for the completion of a flow.
209
@@ -213,14 +214,14 @@ class UniversalService:
214
dict: A dictionary with the success status and a message.
215
"""
216
vql = f"SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1"
216
- # vql = f"SELECT * FROM query(org_id='OL680', query='SELECT * FROM watch_monitoring(artifact='System.Flow.Completion') WHERE FlowId='{flow_id}' LIMIT 1')"
217
logger.info(f"Watching flow {flow_id} for completion")
218
- return self.execute_query(vql)
218
+ return self.execute_query(vql, org_id)
219
220
def read_collection_results(
221
self,
222
client_id: str,
223
flow_id: str,
224
+ org_id: str = "root",
225
artifact: str = "Generic.Client.Info/BasicInformation",
226
):
227
"""
@@ -235,7 +236,7 @@ class UniversalService:
236
dict: A dictionary with the success status, a message, and potentially the results.
237
"""
238
vql = f"SELECT * FROM source(client_id='{client_id}', flow_id='{flow_id}', artifact='{artifact}')"
238
- return self.execute_query(vql)
239
+ return self.execute_query(vql, org_id)
240
241
async def get_client_id(self, client_name: str):
242
"""