| 1 | import hashlib |
| 2 | import os |
| 3 | import tempfile |
| 4 | from typing import Optional |
| 5 | |
| 6 | import aiohttp |
| 7 | from fastapi import HTTPException |
| 8 | from fastapi import UploadFile |
| 9 | from loguru import logger |
| 10 | |
| 11 | from app.data_store.data_store_schema import CaseDataStoreCreation |
| 12 | from app.data_store.data_store_schema import CaseReportTemplateDataStoreCreation |
| 13 | from app.data_store.data_store_session import create_session |
| 14 | |
| 15 | |
| 16 | async def create_bucket_if_not_exists(bucket_name: str) -> None: |
| 17 | client = await create_session() |
| 18 | if not await client.bucket_exists(bucket_name): |
| 19 | await client.make_bucket(bucket_name) |
| 20 | logger.info(f"Created bucket {bucket_name}") |
| 21 | else: |
| 22 | logger.info(f"Bucket {bucket_name} already exists") |
| 23 | |
| 24 | |
| 25 | async def upload_case_data_store(data: CaseDataStoreCreation, file: UploadFile) -> None: |
| 26 | client = await create_session() |
| 27 | logger.info(f"Uploading file {file.filename} to bucket {data.bucket_name}") |
| 28 | |
| 29 | # Save the file to a secure temporary location |
| 30 | with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as tmp: |
| 31 | content = await file.read() |
| 32 | tmp.write(content) |
| 33 | temp_file_path = tmp.name |
| 34 | |
| 35 | try: |
| 36 | await client.fput_object( |
| 37 | bucket_name=data.bucket_name, |
| 38 | object_name=f"{data.case_id}/{file.filename}", |
| 39 | file_path=temp_file_path, |
| 40 | content_type=data.content_type, |
| 41 | ) |
| 42 | finally: |
| 43 | if os.path.exists(temp_file_path): |
| 44 | os.remove(temp_file_path) |
| 45 | |
| 46 | |
| 47 | async def upload_case_report_template_data_store(data: CaseReportTemplateDataStoreCreation, file: UploadFile) -> None: |
| 48 | client = await create_session() |
| 49 | logger.info(f"Uploading file {file.filename} to bucket {data.bucket_name}") |
| 50 | |
| 51 | # Save the file to a secure temporary location |
| 52 | with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as tmp: |
| 53 | content = await file.read() |
| 54 | tmp.write(content) |
| 55 | temp_file_path = tmp.name |
| 56 | |
| 57 | try: |
| 58 | await client.fput_object( |
| 59 | bucket_name=data.bucket_name, |
| 60 | object_name=f"{file.filename}", |
| 61 | file_path=temp_file_path, |
| 62 | content_type=data.content_type, |
| 63 | ) |
| 64 | finally: |
| 65 | if os.path.exists(temp_file_path): |
| 66 | os.remove(temp_file_path) |
| 67 | |
| 68 | |
| 69 | async def download_data_store(bucket_name: str, object_name: str) -> bytes: |
| 70 | client = await create_session() |
| 71 | logger.info(f"Downloading file {object_name} from bucket {bucket_name}") |
| 72 | try: |
| 73 | # Check if the file exists |
| 74 | await client.stat_object(bucket_name, object_name) |
| 75 | |
| 76 | # If no exception is raised, the file exists, proceed to download |
| 77 | async with aiohttp.ClientSession() as session: |
| 78 | response = await client.get_object(bucket_name, object_name, session) |
| 79 | if response is None: |
| 80 | raise Exception("Received None response from get_object") |
| 81 | if not isinstance(response, aiohttp.ClientResponse): |
| 82 | raise Exception("Response is not an instance of aiohttp.ClientResponse") |
| 83 | data = await response.read() # Ensure to read the data |
| 84 | response.close() # Close the response to release resources |
| 85 | logger.info(f"Downloaded file {object_name} from bucket {bucket_name} and returning data") |
| 86 | return data |
| 87 | except Exception as e: |
| 88 | # If an exception is raised, the file does not exist |
| 89 | logger.info(f"Error: {e}") |
| 90 | # List all objects in the bucket |
| 91 | objects = client.list_objects(bucket_name, recursive=True) |
| 92 | objects_list = [obj.object_name async for obj in objects] |
| 93 | logger.info(f"Objects in bucket {bucket_name}: {objects_list}") |
| 94 | raise HTTPException(status_code=404, detail=f"File {object_name} not found in bucket {bucket_name}") |
| 95 | |
| 96 | |
| 97 | async def list_case_data_store_files(bucket_name: str, case_id: int) -> list: |
| 98 | client = await create_session() |
| 99 | objects = await client.list_objects(bucket_name, prefix=f"{case_id}/") |
| 100 | return objects |
| 101 | |
| 102 | |
| 103 | async def list_case_report_template_data_store_files(bucket_name: Optional[str] = "copilot-case-report-templates") -> list: |
| 104 | client = await create_session() |
| 105 | objects_list = [] |
| 106 | logger.info(f"Listing objects in bucket {bucket_name}") |
| 107 | objects = await client.list_objects(bucket_name) |
| 108 | for obj in objects: |
| 109 | logger.info(f"Object: {obj.object_name}") |
| 110 | objects_list.append(obj.object_name) |
| 111 | return objects_list |
| 112 | |
| 113 | |
| 114 | async def upload_file_to_datastore( |
| 115 | file: UploadFile, |
| 116 | bucket_name: str, |
| 117 | object_name: str, |
| 118 | ) -> dict: |
| 119 | """ |
| 120 | Upload a file to MinIO data store. |
| 121 | |
| 122 | Args: |
| 123 | file: The file to upload |
| 124 | bucket_name: The name of the bucket |
| 125 | object_name: The object path/name within the bucket |
| 126 | |
| 127 | Returns: |
| 128 | dict: Upload details including object_key, file_size, file_hash, and content_type |
| 129 | """ |
| 130 | client = await create_session() |
| 131 | |
| 132 | # Create bucket if it doesn't exist |
| 133 | await create_bucket_if_not_exists(bucket_name) |
| 134 | |
| 135 | logger.info(f"Uploading file {file.filename} to bucket {bucket_name} as {object_name}") |
| 136 | |
| 137 | # Save the file to a secure temporary location and calculate hash |
| 138 | sha256_hash = hashlib.sha256() |
| 139 | with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as tmp: |
| 140 | content = await file.read() |
| 141 | tmp.write(content) |
| 142 | sha256_hash.update(content) |
| 143 | temp_file_path = tmp.name |
| 144 | |
| 145 | file_hash = sha256_hash.hexdigest() |
| 146 | file_size = os.path.getsize(temp_file_path) |
| 147 | |
| 148 | # Determine content type |
| 149 | content_type = file.content_type or "application/octet-stream" |
| 150 | |
| 151 | try: |
| 152 | # Upload the file to MinIO |
| 153 | await client.fput_object( |
| 154 | bucket_name=bucket_name, |
| 155 | object_name=object_name, |
| 156 | file_path=temp_file_path, |
| 157 | content_type=content_type, |
| 158 | ) |
| 159 | finally: |
| 160 | if os.path.exists(temp_file_path): |
| 161 | os.remove(temp_file_path) |
| 162 | |
| 163 | logger.info(f"Successfully uploaded {file.filename} ({file_size} bytes) to {bucket_name}/{object_name}") |
| 164 | |
| 165 | return { |
| 166 | "bucket_name": bucket_name, |
| 167 | "object_key": object_name, |
| 168 | "file_name": file.filename, |
| 169 | "file_size": file_size, |
| 170 | "file_hash": file_hash, |
| 171 | "content_type": content_type, |
| 172 | } |
| 173 | |
| 174 | |
| 175 | async def create_buckets() -> None: |
| 176 | await create_bucket_if_not_exists("copilot-cases") |
| 177 | |
| 178 | |
| 179 | async def delete_file(bucket_name: str, object_name: str) -> None: |
| 180 | client = await create_session() |
| 181 | try: |
| 182 | # Check if the file exists |
| 183 | await client.stat_object(bucket_name, object_name) |
| 184 | # If no exception is raised, the file exists, proceed to delete |
| 185 | await client.remove_object(bucket_name, object_name) |
| 186 | logger.info(f"Deleted file {object_name} from bucket {bucket_name}") |
| 187 | except Exception as e: |
| 188 | # If an exception is raised, the file does not exist |
| 189 | logger.info(f"Error: {e}") |
| 190 | raise HTTPException(status_code=404, detail=f"File {object_name} not found in bucket {bucket_name}") |
| 191 | |
| 192 | |
| 193 | # ! Handle Sysmon Config Files ! # |
| 194 | |
| 195 | |
| 196 | async def upload_sysmon_config(customer_code: str, file: UploadFile) -> None: |
| 197 | """ |
| 198 | Upload a sysmon config XML file to the specified customer folder in the sysmon-configs bucket. |
| 199 | Creates the customer folder if it doesn't exist. |
| 200 | |
| 201 | Args: |
| 202 | customer_code: The customer code used as folder name |
| 203 | file: The uploaded XML file |
| 204 | """ |
| 205 | bucket_name = "sysmon-configs" |
| 206 | client = await create_session() |
| 207 | |
| 208 | # Create bucket if it doesn't exist |
| 209 | await create_bucket_if_not_exists(bucket_name) |
| 210 | |
| 211 | logger.info(f"Uploading sysmon config for customer {customer_code}") |
| 212 | |
| 213 | # Save the file to a secure temporary location |
| 214 | with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as tmp: |
| 215 | content = await file.read() |
| 216 | tmp.write(content) |
| 217 | temp_file_path = tmp.name |
| 218 | |
| 219 | # Upload the file to MinIO with customer folder structure |
| 220 | object_name = f"{customer_code}/sysmon_config.xml" |
| 221 | |
| 222 | try: |
| 223 | await client.fput_object(bucket_name=bucket_name, object_name=object_name, file_path=temp_file_path, content_type="application/xml") |
| 224 | finally: |
| 225 | if os.path.exists(temp_file_path): |
| 226 | os.remove(temp_file_path) |
| 227 | |
| 228 | logger.info(f"Successfully uploaded sysmon config for customer {customer_code}") |
| 229 | |
| 230 | |
| 231 | async def download_sysmon_config(customer_code: str) -> bytes: |
| 232 | """ |
| 233 | Download the sysmon config XML file for the specified customer. |
| 234 | |
| 235 | Args: |
| 236 | customer_code: The customer code |
| 237 | |
| 238 | Returns: |
| 239 | bytes: The content of the sysmon config file |
| 240 | """ |
| 241 | bucket_name = "sysmon-configs" |
| 242 | object_name = f"{customer_code}/sysmon_config.xml" |
| 243 | |
| 244 | return await download_data_store(bucket_name, object_name) |
| 245 | |
| 246 | |
| 247 | async def list_sysmon_configs() -> list: |
| 248 | """ |
| 249 | List all customer sysmon configs available in the sysmon-configs bucket. |
| 250 | |
| 251 | Returns: |
| 252 | list: List of customer codes with sysmon configs |
| 253 | """ |
| 254 | bucket_name = "sysmon-configs" |
| 255 | client = await create_session() |
| 256 | |
| 257 | try: |
| 258 | # Make sure bucket exists |
| 259 | if not await client.bucket_exists(bucket_name): |
| 260 | await client.make_bucket(bucket_name) |
| 261 | return [] |
| 262 | |
| 263 | # Get all objects |
| 264 | objects = client.list_objects(bucket_name, recursive=True) |
| 265 | customers = set() |
| 266 | |
| 267 | # Extract customer codes from paths |
| 268 | async for obj in objects: |
| 269 | if obj.object_name.endswith("sysmon_config.xml"): |
| 270 | customer_code = obj.object_name.split("/")[0] |
| 271 | customers.add(customer_code) |
| 272 | |
| 273 | return list(customers) |
| 274 | except Exception as e: |
| 275 | logger.error(f"Error listing sysmon configs: returning empty list - {e}") |
| 276 | return [] |
| 277 | |
| 278 | |
| 279 | # ! Agent Data Store Operations ! # |
| 280 | async def upload_agent_artifact_file( |
| 281 | agent_id: str, |
| 282 | flow_id: str, |
| 283 | file_path: str, |
| 284 | file_name: str, |
| 285 | ) -> dict: |
| 286 | """ |
| 287 | Upload a Velociraptor artifact collection file to MinIO. |
| 288 | |
| 289 | Args: |
| 290 | agent_id: The agent ID |
| 291 | flow_id: The flow ID |
| 292 | file_path: Local path to the file to upload |
| 293 | file_name: Name of the file |
| 294 | |
| 295 | Returns: |
| 296 | dict: Upload details including object_key, file_size, and file_hash |
| 297 | """ |
| 298 | bucket_name = "velociraptor-artifacts" |
| 299 | client = await create_session() |
| 300 | |
| 301 | # Create bucket if it doesn't exist |
| 302 | await create_bucket_if_not_exists(bucket_name) |
| 303 | |
| 304 | # Calculate file hash |
| 305 | sha256_hash = hashlib.sha256() |
| 306 | with open(file_path, "rb") as f: |
| 307 | for byte_block in iter(lambda: f.read(4096), b""): |
| 308 | sha256_hash.update(byte_block) |
| 309 | file_hash = sha256_hash.hexdigest() |
| 310 | |
| 311 | # Get file size |
| 312 | file_size = os.path.getsize(file_path) |
| 313 | |
| 314 | # Construct object key: agent_id/flow_id/filename.zip |
| 315 | object_key = f"{agent_id}/{flow_id}/{file_name}" |
| 316 | |
| 317 | logger.info(f"Uploading artifact file to {bucket_name}/{object_key}") |
| 318 | |
| 319 | # Upload to MinIO |
| 320 | await client.fput_object( |
| 321 | bucket_name=bucket_name, |
| 322 | object_name=object_key, |
| 323 | file_path=file_path, |
| 324 | content_type="application/zip", |
| 325 | ) |
| 326 | |
| 327 | logger.info(f"Successfully uploaded {file_name} ({file_size} bytes) to MinIO") |
| 328 | |
| 329 | return { |
| 330 | "bucket_name": bucket_name, |
| 331 | "object_key": object_key, |
| 332 | "file_name": file_name, |
| 333 | "file_size": file_size, |
| 334 | "file_hash": file_hash, |
| 335 | } |
| 336 | |
| 337 | |
| 338 | async def download_agent_artifact_file(agent_id: str, flow_id: str, file_name: str) -> bytes: |
| 339 | """ |
| 340 | Download a Velociraptor artifact file from MinIO. |
| 341 | |
| 342 | Args: |
| 343 | agent_id: The agent ID |
| 344 | flow_id: The flow ID |
| 345 | file_name: Name of the file |
| 346 | |
| 347 | Returns: |
| 348 | bytes: File content |
| 349 | """ |
| 350 | bucket_name = "velociraptor-artifacts" |
| 351 | object_key = f"{agent_id}/{flow_id}/{file_name}" |
| 352 | |
| 353 | return await download_data_store(bucket_name, object_key) |
| 354 | |
| 355 | |
| 356 | async def list_agent_artifact_files(agent_id: str, flow_id: Optional[str] = None) -> list: |
| 357 | """ |
| 358 | List all artifact files for an agent, optionally filtered by flow_id. |
| 359 | |
| 360 | Args: |
| 361 | agent_id: The agent ID |
| 362 | flow_id: Optional flow ID to filter by |
| 363 | |
| 364 | Returns: |
| 365 | list: List of object names |
| 366 | """ |
| 367 | bucket_name = "velociraptor-artifacts" |
| 368 | client = await create_session() |
| 369 | |
| 370 | prefix = f"{agent_id}/" |
| 371 | if flow_id: |
| 372 | prefix = f"{agent_id}/{flow_id}/" |
| 373 | |
| 374 | objects = await client.list_objects(bucket_name, prefix=prefix, recursive=True) |
| 375 | return [obj.object_name for obj in objects] |
| 376 | |
| 377 | |
| 378 | async def delete_agent_artifact_file(agent_id: str, flow_id: str, file_name: str) -> None: |
| 379 | """ |
| 380 | Delete a Velociraptor artifact file from MinIO. |
| 381 | |
| 382 | Args: |
| 383 | agent_id: The agent ID |
| 384 | flow_id: The flow ID |
| 385 | file_name: Name of the file |
| 386 | """ |
| 387 | bucket_name = "velociraptor-artifacts" |
| 388 | object_key = f"{agent_id}/{flow_id}/{file_name}" |
| 389 | |
| 390 | await delete_file(bucket_name, object_key) |
| 391 | |
| 392 | |
| 393 | async def store_file_in_minio( |
| 394 | file_content: bytes, |
| 395 | bucket_name: str, |
| 396 | object_key: str, |
| 397 | content_type: str = "application/octet-stream", |
| 398 | ) -> dict: |
| 399 | """ |
| 400 | Store file content directly in MinIO without requiring an UploadFile object. |
| 401 | Useful for programmatically generated files like CSV reports. |
| 402 | |
| 403 | Args: |
| 404 | file_content: The file content as bytes |
| 405 | bucket_name: The name of the bucket |
| 406 | object_key: The object path/name within the bucket |
| 407 | content_type: MIME type of the file |
| 408 | |
| 409 | Returns: |
| 410 | dict: Upload details including success status, object_key, file_size, and file_hash |
| 411 | """ |
| 412 | client = await create_session() |
| 413 | |
| 414 | # Create bucket if it doesn't exist |
| 415 | await create_bucket_if_not_exists(bucket_name) |
| 416 | |
| 417 | logger.info(f"Storing file to bucket {bucket_name} as {object_key}") |
| 418 | |
| 419 | try: |
| 420 | # Calculate file hash |
| 421 | sha256_hash = hashlib.sha256() |
| 422 | sha256_hash.update(file_content) |
| 423 | file_hash = sha256_hash.hexdigest() |
| 424 | file_size = len(file_content) |
| 425 | |
| 426 | # Create a temporary file to upload |
| 427 | with tempfile.NamedTemporaryFile(delete=False, dir="/tmp") as temp_file: |
| 428 | temp_file.write(file_content) |
| 429 | temp_file_path = temp_file.name |
| 430 | |
| 431 | try: |
| 432 | # Upload the file to MinIO |
| 433 | await client.fput_object( |
| 434 | bucket_name=bucket_name, |
| 435 | object_name=object_key, |
| 436 | file_path=temp_file_path, |
| 437 | content_type=content_type, |
| 438 | ) |
| 439 | finally: |
| 440 | if os.path.exists(temp_file_path): |
| 441 | os.remove(temp_file_path) |
| 442 | |
| 443 | logger.info(f"Successfully stored file ({file_size} bytes) to {bucket_name}/{object_key}") |
| 444 | |
| 445 | return { |
| 446 | "success": True, |
| 447 | "bucket_name": bucket_name, |
| 448 | "object_key": object_key, |
| 449 | "file_size": file_size, |
| 450 | "file_hash": file_hash, |
| 451 | "content_type": content_type, |
| 452 | } |
| 453 | |
| 454 | except Exception as e: |
| 455 | logger.error(f"Error storing file in MinIO: {e}") |
| 456 | return { |
| 457 | "success": False, |
| 458 | "error": str(e), |
| 459 | } |
| 460 | |
| 461 | |
| 462 | async def retrieve_file_from_minio( |
| 463 | bucket_name: str, |
| 464 | object_key: str, |
| 465 | ) -> dict: |
| 466 | """ |
| 467 | Retrieve file content from MinIO. |
| 468 | |
| 469 | Args: |
| 470 | bucket_name: The name of the bucket |
| 471 | object_key: The object path/name within the bucket |
| 472 | |
| 473 | Returns: |
| 474 | dict: Contains success status and file_content (bytes) if successful, or error message |
| 475 | """ |
| 476 | client = await create_session() |
| 477 | |
| 478 | logger.info(f"Retrieving file {object_key} from bucket {bucket_name}") |
| 479 | |
| 480 | try: |
| 481 | # Check if the file exists |
| 482 | await client.stat_object(bucket_name, object_key) |
| 483 | |
| 484 | # Download the file |
| 485 | async with aiohttp.ClientSession() as session: |
| 486 | response = await client.get_object(bucket_name, object_key, session) |
| 487 | if response is None: |
| 488 | raise Exception("Received None response from get_object") |
| 489 | if not isinstance(response, aiohttp.ClientResponse): |
| 490 | raise Exception("Response is not an instance of aiohttp.ClientResponse") |
| 491 | |
| 492 | file_content = await response.read() |
| 493 | response.close() |
| 494 | |
| 495 | logger.info(f"Successfully retrieved file {object_key} from bucket {bucket_name}") |
| 496 | |
| 497 | return { |
| 498 | "success": True, |
| 499 | "file_content": file_content, |
| 500 | } |
| 501 | |
| 502 | except Exception as e: |
| 503 | logger.error(f"Error retrieving file from MinIO: {e}") |
| 504 | return { |
| 505 | "success": False, |
| 506 | "error": str(e), |
| 507 | } |
| 508 | |
| 509 | |
| 510 | async def delete_file_from_minio( |
| 511 | bucket_name: str, |
| 512 | object_key: str, |
| 513 | ) -> dict: |
| 514 | """ |
| 515 | Delete a file from MinIO. |
| 516 | |
| 517 | Args: |
| 518 | bucket_name: The name of the bucket |
| 519 | object_key: The object path/name within the bucket |
| 520 | |
| 521 | Returns: |
| 522 | dict: Contains success status and message |
| 523 | """ |
| 524 | client = await create_session() |
| 525 | |
| 526 | logger.info(f"Deleting file {object_key} from bucket {bucket_name}") |
| 527 | |
| 528 | try: |
| 529 | # Check if the file exists |
| 530 | await client.stat_object(bucket_name, object_key) |
| 531 | |
| 532 | # Delete the file |
| 533 | await client.remove_object(bucket_name, object_key) |
| 534 | |
| 535 | logger.info(f"Successfully deleted file {object_key} from bucket {bucket_name}") |
| 536 | |
| 537 | return { |
| 538 | "success": True, |
| 539 | "message": f"File {object_key} deleted successfully", |
| 540 | } |
| 541 | |
| 542 | except Exception as e: |
| 543 | logger.error(f"Error deleting file from MinIO: {e}") |
| 544 | return { |
| 545 | "success": False, |
| 546 | "error": str(e), |
| 547 | } |
| 548 | |
| 549 | |
| 550 | async def list_files_in_bucket( |
| 551 | bucket_name: str, |
| 552 | prefix: Optional[str] = None, |
| 553 | ) -> dict: |
| 554 | """ |
| 555 | List all files in a bucket with optional prefix filtering. |
| 556 | |
| 557 | Args: |
| 558 | bucket_name: The name of the bucket |
| 559 | prefix: Optional prefix to filter objects |
| 560 | |
| 561 | Returns: |
| 562 | dict: Contains success status and list of object names |
| 563 | """ |
| 564 | client = await create_session() |
| 565 | |
| 566 | logger.info(f"Listing files in bucket {bucket_name}" + (f" with prefix {prefix}" if prefix else "")) |
| 567 | |
| 568 | try: |
| 569 | # Check if bucket exists |
| 570 | if not await client.bucket_exists(bucket_name): |
| 571 | return { |
| 572 | "success": False, |
| 573 | "error": f"Bucket {bucket_name} does not exist", |
| 574 | "objects": [], |
| 575 | } |
| 576 | |
| 577 | # List objects |
| 578 | objects = client.list_objects(bucket_name, prefix=prefix or "", recursive=True) |
| 579 | object_list = [] |
| 580 | |
| 581 | async for obj in objects: |
| 582 | object_list.append( |
| 583 | { |
| 584 | "object_name": obj.object_name, |
| 585 | "size": obj.size, |
| 586 | "last_modified": obj.last_modified, |
| 587 | }, |
| 588 | ) |
| 589 | |
| 590 | logger.info(f"Found {len(object_list)} files in bucket {bucket_name}") |
| 591 | |
| 592 | return { |
| 593 | "success": True, |
| 594 | "objects": object_list, |
| 595 | "count": len(object_list), |
| 596 | } |
| 597 | |
| 598 | except Exception as e: |
| 599 | logger.error(f"Error listing files in bucket: {e}") |
| 600 | return { |
| 601 | "success": False, |
| 602 | "error": str(e), |
| 603 | "objects": [], |
| 604 | } |