576 velo art upload (#577)
* feat: enhance artifact collection to include session management and file upload to MinIO * chore: update current version to 0.1.16
taylor_socfortress committed
Dec 29, 2025 at 13:34 UTC
2e46c726d4fe708f6167c84d640dbabbd6a7cc8c
3 files changed
+191
-3
backend/app/connectors/velociraptor/routes/artifacts.py
+1
-1
@@ -470,7 +470,7 @@ async def collect_artifact(
470
)
471
472
# Assuming run_artifact_collection is an async function and takes a session as a parameter
473
- return await run_artifact_collection(collect_artifact_body)
473
+ return await run_artifact_collection(collect_artifact_body, session)
474
475
476
@velociraptor_artifacts_router.post(
backend/app/connectors/velociraptor/services/artifacts.py
+189
-1
@@ -303,20 +303,171 @@ async def get_artifact_parameters_by_prefix_service(
303
)
304
305
306
+# async def run_artifact_collection(
307
+# collect_artifact_body: CollectArtifactBody,
308
+# ) -> CollectArtifactResponse:
309
+# """
310
+# Run an artifact collection on a client with optional parameters.
311
+
312
+# Args:
313
+# collect_artifact_body: The body of the request with optional parameters.
314
+
315
+# Returns:
316
+# CollectArtifactResponse: A dictionary containing the success status, message and results.
317
+# """
318
+# velociraptor_service = await UniversalService.create("Velociraptor")
319
+# try:
320
+# # Build the query dynamically based on whether parameters are provided
321
+# parameters = getattr(collect_artifact_body, "parameters", None)
322
+
323
+# if parameters:
324
+# # Velociraptor expects parameters in a very specific format
325
+# # For the "env" parameter, we need to construct a dict
326
+# if "env" in parameters and isinstance(parameters["env"], list):
327
+# env_dict = {}
328
+# for item in parameters["env"]:
329
+# env_dict[item.key] = item.value
330
+
331
+# # Format the query with proper VQL syntax
332
+# query = create_query(
333
+# f"SELECT collect_client("
334
+# f"org_id='{collect_artifact_body.velociraptor_org}', "
335
+# f"client_id='{collect_artifact_body.velociraptor_id}', "
336
+# f"artifacts=['{collect_artifact_body.artifact_name}'], "
337
+# f"env=dict(",
338
+# )
339
+
340
+# # Add each environment variable as a key-value pair
341
+# env_parts = []
342
+# for key, value in env_dict.items():
343
+# # Escape any single quotes in the values
344
+# escaped_value = value.replace("'", "\\'")
345
+# env_parts.append(f"`{key}`='{escaped_value}'")
346
+
347
+# query += ", ".join(env_parts)
348
+# query += ")) FROM scope()"
349
+# else:
350
+# # Handle other types of parameters
351
+# query = create_query(
352
+# f"SELECT collect_client("
353
+# f"org_id='{collect_artifact_body.velociraptor_org}', "
354
+# f"client_id='{collect_artifact_body.velociraptor_id}', "
355
+# f"artifacts=['{collect_artifact_body.artifact_name}']",
356
+# )
357
+
358
+# # Add other parameters if needed
359
+# for param_key, param_value in parameters.items():
360
+# if isinstance(param_value, str):
361
+# query += f", `{param_key}`='{param_value}'"
362
+
363
+# query += ") FROM scope()"
364
+# else:
365
+# # Original query without parameters
366
+# query = create_query(
367
+# f"SELECT collect_client("
368
+# f"org_id='{collect_artifact_body.velociraptor_org}', "
369
+# f"client_id='{collect_artifact_body.velociraptor_id}', "
370
+# f"artifacts=['{collect_artifact_body.artifact_name}']) "
371
+# f"FROM scope()",
372
+# )
373
+
374
+# logger.info(f"Running artifact collection with query: {query}")
375
+# flow = velociraptor_service.execute_query(query, org_id=collect_artifact_body.velociraptor_org)
376
+# logger.info(f"Successfully ran artifact collection on {flow}")
377
+
378
+# # Check if results are available
379
+# if not flow.get("results") or len(flow["results"]) == 0:
380
+# logger.error("No results returned from query execution")
381
+# raise HTTPException(
382
+# status_code=500,
383
+# detail="Query execution did not return any results",
384
+# )
385
+
386
+# # Instead of relying on get_artifact_key, extract the flow_id directly from results
387
+# # by checking all keys in the first result for a flow_id
388
+# result_dict = flow["results"][0]
389
+# flow_id = None
390
+
391
+# # Look for any key that has a flow_id in its value
392
+# for key, value in result_dict.items():
393
+# if isinstance(value, dict) and "flow_id" in value:
394
+# flow_id = value["flow_id"]
395
+# logger.debug(f"Found flow_id {flow_id} in key: {key}")
396
+# break
397
+
398
+# if not flow_id:
399
+# logger.error(f"Could not find flow_id in results: {result_dict}")
400
+# raise HTTPException(
401
+# status_code=500,
402
+# detail="Failed to extract flow ID from results",
403
+# )
404
+
405
+# logger.info(f"Extracted flow_id: {flow_id}")
406
+
407
+# completed = velociraptor_service.watch_flow_completion(flow_id, org_id=collect_artifact_body.velociraptor_org)
408
+# logger.info(f"Successfully watched flow completion on {completed}")
409
+
410
+# results = velociraptor_service.read_collection_results(
411
+# client_id=collect_artifact_body.velociraptor_id,
412
+# flow_id=flow_id,
413
+# org_id=collect_artifact_body.velociraptor_org,
414
+# artifact=collect_artifact_body.artifact_name,
415
+# )
416
+
417
+# logger.info(f"Successfully read collection results on {results}")
418
+
419
+# return CollectArtifactResponse(
420
+# success=results["success"],
421
+# message=results["message"],
422
+# results=results["results"],
423
+# )
424
+# except HTTPException as he: # Catch HTTPException separately to propagate the original message
425
+# logger.error(
426
+# f"HTTPException while running artifact collection on {collect_artifact_body}: {he.detail}",
427
+# )
428
+# raise he
429
+# except Exception as err:
430
+# logger.error(
431
+# f"Failed to run artifact collection on {collect_artifact_body}: {err}",
432
+# )
433
+# raise HTTPException(
434
+# status_code=500,
435
+# detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}",
436
+# )
437
+
438
async def run_artifact_collection(
439
collect_artifact_body: CollectArtifactBody,
440
+ session: AsyncSession, # Add session parameter
441
) -> CollectArtifactResponse:
442
"""
310
- Run an artifact collection on a client with optional parameters.
443
+ Run an artifact collection on a client with optional parameters and upload results to MinIO.
444
445
Args:
446
collect_artifact_body: The body of the request with optional parameters.
447
+ session: Database session for storing metadata.
448
449
Returns:
450
CollectArtifactResponse: A dictionary containing the success status, message and results.
451
"""
452
+ from sqlalchemy import select
453
+
454
+ from app.db.universal_models import AgentDataStore
455
+ from app.db.universal_models import Agents
456
+
457
velociraptor_service = await UniversalService.create("Velociraptor")
458
try:
459
+ # Get agent details from database
460
+ result = await session.execute(
461
+ select(Agents).where(Agents.velociraptor_id == collect_artifact_body.velociraptor_id),
462
+ )
463
+ agent = result.scalars().first()
464
+
465
+ if not agent:
466
+ raise HTTPException(
467
+ status_code=404,
468
+ detail=f"Agent with velociraptor_id {collect_artifact_body.velociraptor_id} not found",
469
+ )
470
+
471
# Build the query dynamically based on whether parameters are provided
472
parameters = getattr(collect_artifact_body, "parameters", None)
473
@@ -416,10 +567,47 @@ async def run_artifact_collection(
567
568
logger.info(f"Successfully read collection results on {results}")
569
570
+ # Fetch the collected file from filestore and upload to MinIO
571
+ file_data = None
572
+ try:
573
+ logger.info("Fetching collected file from filestore and uploading to MinIO")
574
+ file_data = await fetch_file_from_filestore(
575
+ client_id=collect_artifact_body.velociraptor_id,
576
+ flow_id=flow_id,
577
+ org_id=collect_artifact_body.velociraptor_org,
578
+ agent_id=agent.agent_id,
579
+ )
580
+
581
+ # Save metadata to database if file upload was successful
582
+ if file_data.get("success"):
583
+ agent_data_store = AgentDataStore(
584
+ agent_id=agent.agent_id,
585
+ velociraptor_id=collect_artifact_body.velociraptor_id,
586
+ artifact_name=collect_artifact_body.artifact_name,
587
+ flow_id=flow_id,
588
+ bucket_name=file_data["bucket_name"],
589
+ object_key=file_data["object_key"],
590
+ file_name=file_data["file_name"],
591
+ content_type=file_data.get("content_type", "application/zip"),
592
+ file_size=file_data["file_size"],
593
+ file_hash=file_data["file_hash"],
594
+ collection_time=datetime.utcnow(),
595
+ status="completed",
596
+ )
597
+ session.add(agent_data_store)
598
+ await session.commit()
599
+ await session.refresh(agent_data_store)
600
+
601
+ logger.info(f"Saved artifact metadata to database with ID {agent_data_store.id}")
602
+ except Exception as file_err:
603
+ logger.warning(f"Failed to upload file to MinIO, but artifact collection succeeded: {file_err}")
604
+ # Continue execution even if file upload fails
605
+
606
return CollectArtifactResponse(
607
success=results["success"],
608
message=results["message"],
609
results=results["results"],
610
+ file_info=file_data if file_data and file_data.get("success") else None,
611
)
612
except HTTPException as he: # Catch HTTPException separately to propagate the original message
613
logger.error(
backend/app/version/services/version.py
+1
-1
@@ -7,7 +7,7 @@ from loguru import logger
7
from packaging.version import Version
8
9
# Current version - update this with each release
10
-CURRENT_VERSION = "0.1.15"
10
+CURRENT_VERSION = "0.1.16"
11
VERSION_CHECK_URL = "https://api.github.com/repos/socfortress/CoPilot/releases/latest"
12
13