| 1 | from typing import List |
| 2 | |
| 3 | from fastapi import APIRouter |
| 4 | from fastapi import Depends |
| 5 | from fastapi import HTTPException |
| 6 | from fastapi import Security |
| 7 | from loguru import logger |
| 8 | from sqlalchemy.ext.asyncio import AsyncSession |
| 9 | from sqlalchemy.future import select |
| 10 | |
| 11 | from app.auth.utils import AuthHandler |
| 12 | from app.connectors.velociraptor.schema.artifacts import ArtifactParametersResponse |
| 13 | from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse |
| 14 | from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody |
| 15 | from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse |
| 16 | from app.connectors.velociraptor.schema.artifacts import CollectFileBody |
| 17 | from app.connectors.velociraptor.schema.artifacts import FileCollectionBody |
| 18 | from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum |
| 19 | from app.connectors.velociraptor.schema.artifacts import OSPrefixModel |
| 20 | from app.connectors.velociraptor.schema.artifacts import QuarantineBody |
| 21 | from app.connectors.velociraptor.schema.artifacts import QuarantineResponse |
| 22 | from app.connectors.velociraptor.schema.artifacts import RunCommandBody |
| 23 | from app.connectors.velociraptor.schema.artifacts import RunCommandResponse |
| 24 | from app.connectors.velociraptor.services.artifacts import get_artifact_by_name |
| 25 | from app.connectors.velociraptor.services.artifacts import ( |
| 26 | get_artifact_parameters_by_prefix_service, |
| 27 | ) |
| 28 | from app.connectors.velociraptor.services.artifacts import get_artifacts |
| 29 | from app.connectors.velociraptor.services.artifacts import quarantine_host |
| 30 | from app.connectors.velociraptor.services.artifacts import run_artifact_collection |
| 31 | from app.connectors.velociraptor.services.artifacts import run_file_collection |
| 32 | from app.connectors.velociraptor.services.artifacts import run_remote_command |
| 33 | from app.connectors.velociraptor.services.artifacts import validate_artifact_parameters |
| 34 | from app.db.db_session import get_db |
| 35 | from app.db.universal_models import Agents |
| 36 | |
| 37 | # App specific imports |
| 38 | |
| 39 | |
| 40 | velociraptor_artifacts_router = APIRouter() |
| 41 | |
| 42 | |
| 43 | # Get all valid OS prefixes |
| 44 | |
| 45 | |
| 46 | def get_valid_os_prefixes() -> List[str]: |
| 47 | """ |
| 48 | Returns a list of valid operating system prefixes. |
| 49 | |
| 50 | Returns: |
| 51 | List[str]: A list of valid operating system prefixes. |
| 52 | """ |
| 53 | return [prefix.name.lower() for prefix in OSPrefixEnum] |
| 54 | |
| 55 | |
| 56 | # Verify the OS prefix exists and return the appropriate Enum value |
| 57 | def verify_os_prefix_exists(os_prefix: str) -> str: |
| 58 | """ |
| 59 | Verify if the given OS prefix exists. |
| 60 | |
| 61 | Args: |
| 62 | os_prefix (str): The OS prefix to be verified. |
| 63 | |
| 64 | Returns: |
| 65 | str: The value of the OS prefix. |
| 66 | |
| 67 | Raises: |
| 68 | HTTPException: If the OS prefix does not exist. |
| 69 | """ |
| 70 | os_prefix_lower = os_prefix.lower() |
| 71 | os_prefix_upper = os_prefix.upper() # Convert to uppercase for Enum matching |
| 72 | valid_os_prefixes = get_valid_os_prefixes() |
| 73 | |
| 74 | if os_prefix_lower not in valid_os_prefixes: |
| 75 | raise HTTPException( |
| 76 | status_code=400, |
| 77 | detail=f"OS prefix {os_prefix} does not exist.", |
| 78 | ) |
| 79 | |
| 80 | return OSPrefixEnum[os_prefix_upper].value # Use the uppercase version for Enum matching |
| 81 | |
| 82 | |
| 83 | def get_os_prefix_from_os_name(os_name: str) -> str: |
| 84 | """ |
| 85 | Get the OS prefix from the OS name. |
| 86 | |
| 87 | Args: |
| 88 | os_name (str): The name of the operating system. |
| 89 | |
| 90 | Returns: |
| 91 | str: The OS prefix corresponding to the OS name. |
| 92 | """ |
| 93 | # Use the OSPrefixModel to get the OS prefix from the OS name |
| 94 | logger.info(f"Getting OS prefix from OS name {os_name}") |
| 95 | os_prefix_model = OSPrefixModel(os_name=os_name) |
| 96 | result = os_prefix_model.get_os_prefix() |
| 97 | logger.info(f"OS prefix for OS name {os_name} is {result}") |
| 98 | return result |
| 99 | |
| 100 | |
| 101 | async def get_velociraptor_id(session: AsyncSession, hostname: str) -> str: |
| 102 | """ |
| 103 | Retrieves the velociraptor_id associated with the given hostname. |
| 104 | |
| 105 | Args: |
| 106 | session (AsyncSession): The database session. |
| 107 | hostname (str): The hostname of the agent. |
| 108 | |
| 109 | Returns: |
| 110 | str: The velociraptor_id associated with the hostname. |
| 111 | |
| 112 | Raises: |
| 113 | HTTPException: If the agent with the given hostname is not found or if the velociraptor_id is not available. |
| 114 | """ |
| 115 | logger.info(f"Getting velociraptor_id from hostname {hostname}") |
| 116 | result = await session.execute(select(Agents).filter(Agents.hostname == hostname)) |
| 117 | agent = result.scalars().first() |
| 118 | |
| 119 | if not agent: |
| 120 | raise HTTPException( |
| 121 | status_code=404, |
| 122 | detail=f"Agent with hostname {hostname} not found", |
| 123 | ) |
| 124 | |
| 125 | if agent.velociraptor_id == "n/a": |
| 126 | raise HTTPException( |
| 127 | status_code=404, |
| 128 | detail=f"Velociraptor ID for hostname {hostname} is not available", |
| 129 | ) |
| 130 | |
| 131 | logger.info(f"velociraptor_id for hostname {hostname} is {agent.velociraptor_id}") |
| 132 | return agent.velociraptor_id |
| 133 | |
| 134 | |
| 135 | async def get_velociraptor_org(session: AsyncSession, hostname: str) -> str: |
| 136 | """ |
| 137 | Retrieves the velociraptor_org associated with the given hostname. |
| 138 | |
| 139 | Args: |
| 140 | session (AsyncSession): The database session. |
| 141 | hostname (str): The hostname of the agent. |
| 142 | |
| 143 | Returns: |
| 144 | str: The velociraptor_org associated with the hostname. |
| 145 | |
| 146 | Raises: |
| 147 | HTTPException: If the agent with the given hostname is not found or if the velociraptor_org is not available. |
| 148 | """ |
| 149 | logger.info(f"Getting velociraptor_org from hostname {hostname}") |
| 150 | result = await session.execute(select(Agents).filter(Agents.hostname == hostname)) |
| 151 | agent = result.scalars().first() |
| 152 | |
| 153 | if not agent: |
| 154 | raise HTTPException( |
| 155 | status_code=404, |
| 156 | detail=f"Agent with hostname {hostname} not found", |
| 157 | ) |
| 158 | |
| 159 | if agent.velociraptor_org is None: |
| 160 | raise HTTPException( |
| 161 | status_code=404, |
| 162 | detail=f"Velociraptor ORG for hostname {hostname} is not available", |
| 163 | ) |
| 164 | |
| 165 | logger.info(f"velociraptor_org for hostname {hostname} is {agent.velociraptor_org}") |
| 166 | return agent.velociraptor_org |
| 167 | |
| 168 | |
| 169 | def format_file_path_for_os(file_path: str, os_prefix: str) -> str: |
| 170 | """ |
| 171 | Format the file path based on the operating system prefix. |
| 172 | |
| 173 | For Windows OS: |
| 174 | - Prepends "Glob\n" to the file path |
| 175 | - Escapes backslashes (single \ becomes \\\\) |
| 176 | - Appends "\n" at the end |
| 177 | |
| 178 | For Linux/MacOS: |
| 179 | - Prepends "Glob\n" to the file path |
| 180 | - Appends "\n" at the end |
| 181 | - Forward slashes are kept as-is |
| 182 | |
| 183 | Args: |
| 184 | file_path (str): The original file path to format. |
| 185 | os_prefix (str): The OS prefix (e.g., "Windows", "Linux", "MacOS", "Generic.Client"). |
| 186 | |
| 187 | Returns: |
| 188 | str: The formatted file path. |
| 189 | |
| 190 | Examples: |
| 191 | >>> format_file_path_for_os("Users\\Administrator\\Downloads\\LICENSE.txt", "Windows") |
| 192 | 'Glob\\nUsers\\\\\\\\Administrator\\\\\\\\Downloads\\\\\\\\LICENSE.txt\\n' |
| 193 | |
| 194 | >>> format_file_path_for_os("/home/user/document.txt", "Linux") |
| 195 | 'Glob\\n/home/user/document.txt\\n' |
| 196 | |
| 197 | >>> format_file_path_for_os("/Users/user/document.txt", "MacOS") |
| 198 | 'Glob\\n/Users/user/document.txt\\n' |
| 199 | """ |
| 200 | logger.info(f"Formatting file path '{file_path}' for OS prefix '{os_prefix}'") |
| 201 | |
| 202 | # Check if the OS is Windows |
| 203 | if os_prefix.lower() in ["windows", "windows."]: |
| 204 | # Escape backslashes: single \ becomes \\\\ |
| 205 | # This is because we need double escaping: once for Python string, once for Velociraptor |
| 206 | escaped_path = file_path.replace("\\", "\\\\") |
| 207 | |
| 208 | # Format with Glob prefix and newline suffix |
| 209 | formatted_path = f"Glob\n{escaped_path}\n" |
| 210 | |
| 211 | logger.info(f"Formatted Windows path: '{formatted_path}'") |
| 212 | return formatted_path |
| 213 | |
| 214 | # For Linux, MacOS, or other Unix-based systems |
| 215 | # Also add Glob prefix and newline suffix, but keep forward slashes as-is |
| 216 | formatted_path = f"Glob\n{file_path}\n" |
| 217 | |
| 218 | logger.info(f"Formatted {os_prefix} path: '{formatted_path}'") |
| 219 | return formatted_path |
| 220 | |
| 221 | |
| 222 | async def update_agent_quarantine_status( |
| 223 | session: AsyncSession, |
| 224 | quarantine_body: QuarantineBody, |
| 225 | quarantine_response: QuarantineResponse, |
| 226 | ): |
| 227 | """ |
| 228 | Updates the quarantine status of an agent. |
| 229 | |
| 230 | Args: |
| 231 | session (AsyncSession): The database session. |
| 232 | quarantine_body (QuarantineBody): The body of the quarantine request. |
| 233 | quarantine_response (QuarantineResponse): The response of the quarantine request. |
| 234 | |
| 235 | Raises: |
| 236 | HTTPException: If the agent with the specified hostname is not found or if the quarantine action fails. |
| 237 | |
| 238 | Returns: |
| 239 | None |
| 240 | """ |
| 241 | logger.info( |
| 242 | f"Updating agent quarantine status for hostname {quarantine_body.hostname}", |
| 243 | ) |
| 244 | result = await session.execute( |
| 245 | select(Agents).filter(Agents.hostname == quarantine_body.hostname), |
| 246 | ) |
| 247 | agent = result.scalars().first() |
| 248 | |
| 249 | if not agent: |
| 250 | raise HTTPException( |
| 251 | status_code=404, |
| 252 | detail=f"Agent with hostname {quarantine_body.hostname} not found", |
| 253 | ) |
| 254 | |
| 255 | if quarantine_body.action == "quarantine": |
| 256 | if quarantine_response.success: |
| 257 | agent.quarantined = True |
| 258 | else: |
| 259 | raise HTTPException( |
| 260 | status_code=500, |
| 261 | detail=f"Failed to quarantine hostname {quarantine_body.hostname}", |
| 262 | ) |
| 263 | elif quarantine_body.action == "remove_quarantine": |
| 264 | if quarantine_response.success: |
| 265 | agent.quarantined = False |
| 266 | else: |
| 267 | raise HTTPException( |
| 268 | status_code=500, |
| 269 | detail=f"Failed to remove quarantine for hostname {quarantine_body.hostname}", |
| 270 | ) |
| 271 | |
| 272 | await session.commit() |
| 273 | |
| 274 | logger.info( |
| 275 | f"Agent quarantine status for hostname {quarantine_body.hostname} updated to {agent.quarantined}", |
| 276 | ) |
| 277 | |
| 278 | return None |
| 279 | |
| 280 | |
| 281 | @velociraptor_artifacts_router.get( |
| 282 | "", |
| 283 | response_model=ArtifactsResponse, |
| 284 | description="Get all artifacts", |
| 285 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 286 | ) |
| 287 | async def get_all_artifacts() -> ArtifactsResponse: |
| 288 | """ |
| 289 | Retrieve all artifacts. |
| 290 | |
| 291 | Returns: |
| 292 | ArtifactsResponse: The response containing all artifacts. |
| 293 | """ |
| 294 | logger.info("Fetching all artifacts") |
| 295 | return await get_artifacts() |
| 296 | |
| 297 | |
| 298 | @velociraptor_artifacts_router.get( |
| 299 | "/{os_prefix}", |
| 300 | response_model=ArtifactsResponse, |
| 301 | description="Get all artifacts for a specific OS prefix", |
| 302 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 303 | ) |
| 304 | async def get_all_artifacts_for_os_prefix( |
| 305 | os_prefix: str = Depends(verify_os_prefix_exists), |
| 306 | ) -> ArtifactsResponse: |
| 307 | """ |
| 308 | Fetches all artifacts for a specific OS prefix. |
| 309 | |
| 310 | Args: |
| 311 | os_prefix (str): The OS prefix to filter the artifacts. |
| 312 | |
| 313 | Returns: |
| 314 | ArtifactsResponse: The response containing the success status, message, and artifacts. |
| 315 | """ |
| 316 | logger.info(f"Fetching all artifacts for OS prefix {os_prefix}") |
| 317 | # Get all the artifacts names that begin with the OS prefix |
| 318 | artifacts = await get_artifacts() |
| 319 | artifacts = artifacts.artifacts |
| 320 | # Match artifacts that start with the OS prefix OR start with "Custom." followed by the OS prefix |
| 321 | artifacts_for_os_prefix = [ |
| 322 | artifact for artifact in artifacts if artifact.name.startswith(os_prefix) or artifact.name.startswith(f"Custom.{os_prefix}") |
| 323 | ] |
| 324 | return ArtifactsResponse( |
| 325 | success=True, |
| 326 | message=f"All artifacts for OS prefix {os_prefix} retrieved", |
| 327 | artifacts=artifacts_for_os_prefix, |
| 328 | ) |
| 329 | |
| 330 | |
| 331 | @velociraptor_artifacts_router.get( |
| 332 | "/artifact/{artifact_name}", |
| 333 | response_model=ArtifactsResponse, |
| 334 | description="Get a specific artifact by name", |
| 335 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 336 | ) |
| 337 | async def get_artifact_by_name_route(artifact_name: str) -> ArtifactsResponse: |
| 338 | """ |
| 339 | Retrieve a specific artifact by its name. |
| 340 | |
| 341 | Args: |
| 342 | artifact_name (str): The name of the artifact to retrieve. |
| 343 | |
| 344 | Returns: |
| 345 | ArtifactsResponse: The response containing the specific artifact. |
| 346 | """ |
| 347 | logger.info(f"Fetching artifact by name: {artifact_name}") |
| 348 | return await get_artifact_by_name(artifact_name) |
| 349 | |
| 350 | |
| 351 | @velociraptor_artifacts_router.get( |
| 352 | "/artifact/{artifact_name}/parameters/{parameter_prefix}", |
| 353 | response_model=ArtifactParametersResponse, |
| 354 | description="Get parameters from an artifact that match a specific prefix", |
| 355 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 356 | ) |
| 357 | async def get_artifact_parameters_by_prefix(artifact_name: str, parameter_prefix: str) -> ArtifactParametersResponse: |
| 358 | """ |
| 359 | Retrieve parameters from a specific artifact that start with the given prefix. |
| 360 | |
| 361 | Args: |
| 362 | artifact_name (str): The name of the artifact to retrieve parameters from. |
| 363 | parameter_prefix (str): The prefix to filter parameters by (e.g., "T1552.001"). |
| 364 | |
| 365 | Returns: |
| 366 | ArtifactParametersResponse: The response containing matching parameters. |
| 367 | """ |
| 368 | logger.info(f"Fetching parameters with prefix '{parameter_prefix}' from artifact '{artifact_name}'") |
| 369 | return await get_artifact_parameters_by_prefix_service(artifact_name, parameter_prefix) |
| 370 | |
| 371 | |
| 372 | @velociraptor_artifacts_router.get( |
| 373 | "/hostname/{hostname}", |
| 374 | response_model=ArtifactsResponse, |
| 375 | description="Get all artifacts for a specific host's OS prefix", |
| 376 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 377 | ) |
| 378 | async def get_all_artifacts_for_hostname( |
| 379 | hostname: str, |
| 380 | session: AsyncSession = Depends(get_db), |
| 381 | ) -> ArtifactsResponse: |
| 382 | """ |
| 383 | Retrieve all artifacts for a specific host's OS prefix. |
| 384 | |
| 385 | Args: |
| 386 | hostname (str): The hostname of the host. |
| 387 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 388 | |
| 389 | Returns: |
| 390 | ArtifactsResponse: The response containing the retrieved artifacts. |
| 391 | """ |
| 392 | logger.info(f"Fetching all artifacts for hostname {hostname}") |
| 393 | |
| 394 | # Asynchronous query to find the agent |
| 395 | agent_result = await session.execute( |
| 396 | select(Agents).filter(Agents.hostname == hostname), |
| 397 | ) |
| 398 | agent = agent_result.scalars().first() |
| 399 | |
| 400 | if not agent: |
| 401 | raise HTTPException( |
| 402 | status_code=404, |
| 403 | detail=f"Agent with hostname {hostname} not found", |
| 404 | ) |
| 405 | |
| 406 | os_prefix = get_os_prefix_from_os_name(os_name=agent.os.lower()) |
| 407 | if not os_prefix: |
| 408 | raise HTTPException( |
| 409 | status_code=404, |
| 410 | detail=f"OS prefix of {agent.os.lower()} for hostname {hostname} not found", |
| 411 | ) |
| 412 | |
| 413 | # Assuming get_all_artifacts_for_os_prefix is an async function |
| 414 | result = await get_all_artifacts_for_os_prefix(os_prefix) |
| 415 | |
| 416 | return ArtifactsResponse( |
| 417 | success=True, |
| 418 | message=f"All available artifacts that can be ran for hostname {hostname} retrieved", |
| 419 | artifacts=result.artifacts, |
| 420 | ) |
| 421 | |
| 422 | |
| 423 | @velociraptor_artifacts_router.post( |
| 424 | "/collect", |
| 425 | response_model=CollectArtifactResponse, |
| 426 | description="Run an analyzer", |
| 427 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 428 | ) |
| 429 | async def collect_artifact( |
| 430 | collect_artifact_body: CollectArtifactBody, |
| 431 | session: AsyncSession = Depends(get_db), |
| 432 | ) -> CollectArtifactResponse: |
| 433 | """ |
| 434 | Collects an artifact for a given hostname. |
| 435 | |
| 436 | Args: |
| 437 | collect_artifact_body (CollectArtifactBody): The request body containing the hostname and artifact name. |
| 438 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 439 | |
| 440 | Returns: |
| 441 | CollectArtifactResponse: The response containing the collected artifact. |
| 442 | """ |
| 443 | logger.info(f"Received request to collect artifact {collect_artifact_body}") |
| 444 | result = await get_all_artifacts_for_hostname( |
| 445 | collect_artifact_body.hostname, |
| 446 | session, |
| 447 | ) |
| 448 | artifact_names = [artifact.name for artifact in result.artifacts] |
| 449 | |
| 450 | if collect_artifact_body.artifact_name not in artifact_names: |
| 451 | raise HTTPException( |
| 452 | status_code=400, |
| 453 | detail=f"Artifact name {collect_artifact_body.artifact_name} does not apply for hostname {collect_artifact_body.hostname} or does not exist", |
| 454 | ) |
| 455 | |
| 456 | # Validate parameters if provided |
| 457 | await validate_artifact_parameters( |
| 458 | collect_artifact_body.artifact_name, |
| 459 | collect_artifact_body.parameters, |
| 460 | ) |
| 461 | |
| 462 | collect_artifact_body.velociraptor_id = await get_velociraptor_id( |
| 463 | session, |
| 464 | collect_artifact_body.hostname, |
| 465 | ) |
| 466 | |
| 467 | collect_artifact_body.velociraptor_org = await get_velociraptor_org( |
| 468 | session, |
| 469 | collect_artifact_body.hostname, |
| 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, session) |
| 474 | |
| 475 | |
| 476 | @velociraptor_artifacts_router.post( |
| 477 | "/command", |
| 478 | response_model=RunCommandResponse, |
| 479 | description="Run a remote command", |
| 480 | dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])], |
| 481 | ) |
| 482 | async def run_command( |
| 483 | run_command_body: RunCommandBody, |
| 484 | session: AsyncSession = Depends(get_db), |
| 485 | ) -> RunCommandResponse: |
| 486 | """ |
| 487 | Run a remote command. |
| 488 | |
| 489 | Args: |
| 490 | run_command_body (RunCommandBody): The request body containing the command details. |
| 491 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 492 | |
| 493 | Returns: |
| 494 | RunCommandResponse: The response containing the result of the command execution. |
| 495 | """ |
| 496 | logger.info(f"Received request to run command {run_command_body}") |
| 497 | result = await get_all_artifacts_for_hostname(run_command_body.hostname, session) |
| 498 | artifact_names = [artifact.name for artifact in result.artifacts] |
| 499 | if run_command_body.artifact_name not in artifact_names: |
| 500 | raise HTTPException( |
| 501 | status_code=400, |
| 502 | detail=f"Artifact name {run_command_body.artifact_name.value} does not apply for hostname {run_command_body.hostname} or does not exist", |
| 503 | ) |
| 504 | # Add the velociraptor_id to the run_command_body object |
| 505 | run_command_body.velociraptor_id = await get_velociraptor_id( |
| 506 | session, |
| 507 | run_command_body.hostname, |
| 508 | ) |
| 509 | |
| 510 | run_command_body.velociraptor_org = await get_velociraptor_org( |
| 511 | session, |
| 512 | run_command_body.hostname, |
| 513 | ) |
| 514 | # Run the command |
| 515 | return await run_remote_command(run_command_body) |
| 516 | |
| 517 | |
| 518 | @velociraptor_artifacts_router.post( |
| 519 | "/quarantine", |
| 520 | response_model=QuarantineResponse, |
| 521 | description="Quarantine a host", |
| 522 | dependencies=[Security(AuthHandler().get_current_user, scopes=["admin"])], |
| 523 | ) |
| 524 | async def quarantine( |
| 525 | quarantine_body: QuarantineBody, |
| 526 | session: AsyncSession = Depends(get_db), |
| 527 | ) -> QuarantineResponse: |
| 528 | """ |
| 529 | Quarantine a host. |
| 530 | |
| 531 | Args: |
| 532 | quarantine_body (QuarantineBody): The body of the request containing the hostname and artifact name. |
| 533 | session (AsyncSession, optional): The database session. Defaults to Depends(get_db). |
| 534 | |
| 535 | Returns: |
| 536 | QuarantineResponse: The response containing the result of the quarantine operation. |
| 537 | """ |
| 538 | logger.info(f"Received request to quarantine host {quarantine_body}") |
| 539 | result = await get_all_artifacts_for_hostname(quarantine_body.hostname, session) |
| 540 | artifact_names = [artifact.name for artifact in result.artifacts] |
| 541 | if quarantine_body.artifact_name not in artifact_names: |
| 542 | raise HTTPException( |
| 543 | status_code=400, |
| 544 | detail=f"Artifact name {quarantine_body.artifact_name.value} does not apply for hostname {quarantine_body.hostname} or does not exist", |
| 545 | ) |
| 546 | # Add the velociraptor_id to the run_command_body object |
| 547 | # Add the velociraptor_id to the quarantine_body object |
| 548 | quarantine_body.velociraptor_id = await get_velociraptor_id( |
| 549 | session, |
| 550 | quarantine_body.hostname, |
| 551 | ) |
| 552 | |
| 553 | quarantine_body.velociraptor_org = await get_velociraptor_org( |
| 554 | session, |
| 555 | quarantine_body.hostname, |
| 556 | ) |
| 557 | |
| 558 | # Quarantine the host |
| 559 | quarantine_response = await quarantine_host(quarantine_body) |
| 560 | |
| 561 | # If the host was successfully quarantined, update the database |
| 562 | await update_agent_quarantine_status(session, quarantine_body, quarantine_response) |
| 563 | |
| 564 | return quarantine_response |
| 565 | |
| 566 | |
| 567 | # Add this new route after the existing collect_file route |
| 568 | |
| 569 | |
| 570 | @velociraptor_artifacts_router.post( |
| 571 | "/collect/file/agent/{agent_id}", |
| 572 | response_model=CollectArtifactResponse, |
| 573 | description="Collect a file from an agent using agent ID", |
| 574 | dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))], |
| 575 | ) |
| 576 | async def collect_file_by_agent_id( |
| 577 | agent_id: str, |
| 578 | collect_file_body: FileCollectionBody, |
| 579 | session: AsyncSession = Depends(get_db), |
| 580 | ) -> CollectArtifactResponse: |
| 581 | """ |
| 582 | Collects a file from an agent using the agent ID. |
| 583 | |
| 584 | Args: |
| 585 | agent_id (str): The agent ID to collect the file from. |
| 586 | collect_file_body (FileCollectionBody): The request body containing file path and root disk. |
| 587 | session (AsyncSession): The database session. |
| 588 | |
| 589 | Returns: |
| 590 | CollectArtifactResponse: The response containing the collection status. |
| 591 | |
| 592 | Raises: |
| 593 | HTTPException: If the agent is not found or velociraptor details are not available. |
| 594 | """ |
| 595 | logger.info(f"Received request to collect file for agent ID {agent_id}") |
| 596 | |
| 597 | # Query the agent from the database using agent_id |
| 598 | result = await session.execute(select(Agents).filter(Agents.agent_id == agent_id)) |
| 599 | agent = result.scalars().first() |
| 600 | |
| 601 | if not agent: |
| 602 | raise HTTPException( |
| 603 | status_code=404, |
| 604 | detail=f"Agent with ID {agent_id} not found", |
| 605 | ) |
| 606 | |
| 607 | # Verify hostname exists |
| 608 | if not agent.hostname: |
| 609 | raise HTTPException( |
| 610 | status_code=404, |
| 611 | detail=f"Hostname for agent {agent_id} is not available", |
| 612 | ) |
| 613 | |
| 614 | # Verify velociraptor_id exists and is valid |
| 615 | if not agent.velociraptor_id or agent.velociraptor_id == "n/a": |
| 616 | raise HTTPException( |
| 617 | status_code=404, |
| 618 | detail=f"Velociraptor ID for agent {agent_id} is not available", |
| 619 | ) |
| 620 | |
| 621 | # Verify velociraptor_org exists |
| 622 | if not agent.velociraptor_org: |
| 623 | raise HTTPException( |
| 624 | status_code=404, |
| 625 | detail=f"Velociraptor ORG for agent {agent_id} is not available", |
| 626 | ) |
| 627 | |
| 628 | # Collect the agent os |
| 629 | os_prefix = get_os_prefix_from_os_name(os_name=agent.os.lower()) |
| 630 | |
| 631 | logger.info(f"OS prefix for agent ID {agent_id} is {os_prefix}") |
| 632 | # Format the file path based on the OS |
| 633 | original_file_path = collect_file_body.file |
| 634 | formatted_file_path = format_file_path_for_os(original_file_path, os_prefix) |
| 635 | |
| 636 | logger.info(f"Original file path: '{original_file_path}'") |
| 637 | logger.info(f"Formatted file path: '{formatted_file_path}'") |
| 638 | |
| 639 | # Update the file path in the request body |
| 640 | collect_file_body.file = formatted_file_path |
| 641 | return await run_file_collection( |
| 642 | CollectFileBody( |
| 643 | hostname=agent.hostname, |
| 644 | velociraptor_id=agent.velociraptor_id, |
| 645 | velociraptor_org=agent.velociraptor_org, |
| 646 | artifact_name="Generic.Collectors.File", |
| 647 | file=collect_file_body.file, |
| 648 | root_disk=collect_file_body.root_disk, |
| 649 | ), |
| 650 | session, |
| 651 | ) |