Merge branch 'incident-management' of https://github.com/socfortress/…

…CoPilot into incident-management

backend/app/connectors/velociraptor/routes/artifacts.py

@@ -12,7 +12,7 @@
 from app.connectors.velociraptor.schema.artifacts import ArtifactReccomendationAIRequest
 from app.connectors.velociraptor.schema.artifacts import ArtifactReccomendationRequest
 from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
-from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody, CollectFileBody
 from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
 from app.connectors.velociraptor.schema.artifacts import OSPrefixEnum
 from app.connectors.velociraptor.schema.artifacts import OSPrefixModel
@@ -23,7 +23,7 @@
 from app.connectors.velociraptor.services.artifacts import get_artifacts
 from app.connectors.velociraptor.services.artifacts import post_to_copilot_ai_module
 from app.connectors.velociraptor.services.artifacts import quarantine_host
-from app.connectors.velociraptor.services.artifacts import run_artifact_collection
+from app.connectors.velociraptor.services.artifacts import run_artifact_collection, run_file_collection
 from app.connectors.velociraptor.services.artifacts import run_remote_command
 from app.db.db_session import get_db
 from app.db.universal_models import Agents
@@ -476,3 +476,41 @@ async def get_artifact_recommendation(request: ArtifactReccomendationAIRequest):
             prompt=request.prompt,
         ),
     )
+
+# ! WIP ! #
+@velociraptor_artifacts_router.post(
+    "/collect/file",
+    response_model=CollectArtifactResponse,
+    description="Run an the artifact to collect a file",
+)
+async def collect_file(collect_artifact_body: CollectFileBody, session: AsyncSession = Depends(get_db)):
+    """
+    Collects a file based on the artifact.
+
+    Returns:
+        CollectArtifactResponse: The response containing the collected file.
+    """
+    logger.info(f"Received request to collect artifact {collect_artifact_body}")
+    # result = await get_all_artifacts_for_hostname(
+    #     collect_artifact_body.hostname,
+    #     session,
+    # )
+    # artifact_names = [artifact.name for artifact in result.artifacts]
+
+    # if collect_artifact_body.artifact_name not in artifact_names:
+    #     raise HTTPException(
+    #         status_code=400,
+    #         detail=f"Artifact name {collect_artifact_body.artifact_name} does not apply for hostname {collect_artifact_body.hostname} or does not exist",
+    #     )
+
+    # collect_artifact_body.velociraptor_id = await get_velociraptor_id(
+    #     session,
+    #     collect_artifact_body.hostname,
+    # )
+
+    # collect_artifact_body.velociraptor_org = await get_velociraptor_org(
+    #     session,
+    #     collect_artifact_body.hostname,
+    # )
+    return await run_file_collection(collect_artifact_body)
+

backend/app/connectors/velociraptor/schema/artifacts.py

@@ -5,7 +5,8 @@
 from typing import Optional
 
 from pydantic import BaseModel
-from pydantic import Field
+from pydantic import Field, validator
+from fastapi import HTTPException
 
 
 class Artifacts(BaseModel):
@@ -82,6 +83,20 @@ class CollectArtifactBody(BaseBody):
         description="Name of the artifact for collection or command running",
     )
 
+class CollectFileBody(BaseBody):
+    artifact_name: str = Field(
+        'Generic.Collectors.File',
+        description="Name of the artifact for collection or command running",
+    )
+    file: str = Field('Glob\nUsers\\Administrator\\Documents\\*\n', description="File to collect")
+    root_disk: Optional[str] = Field('C:', description="Root disk to collect from")
+
+    @validator('artifact_name')
+    def validate_artifact_name(cls, value):
+        if value != 'Generic.Collectors.File':
+            raise HTTPException(status_code=400, detail="Invalid artifact name. Name should be 'Generic.Collectors.File'")
+        return value
+
 
 class RunCommandBody(BaseBody):
     command: Optional[str] = Field(None, description="Command to run")

backend/app/connectors/velociraptor/services/artifacts.py

@@ -7,7 +7,7 @@
 from app.connectors.velociraptor.schema.artifacts import Artifacts
 from app.connectors.velociraptor.schema.artifacts import ArtifactsResponse
 from app.connectors.velociraptor.schema.artifacts import CollectArtifactBody
-from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse
+from app.connectors.velociraptor.schema.artifacts import CollectArtifactResponse, CollectFileBody
 from app.connectors.velociraptor.schema.artifacts import QuarantineBody
 from app.connectors.velociraptor.schema.artifacts import QuarantineResponse
 from app.connectors.velociraptor.schema.artifacts import RunCommandBody
@@ -164,6 +164,78 @@ async def run_artifact_collection(
             detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}",
         )
 
+async def run_file_collection(
+    collect_artifact_body: CollectFileBody,
+) -> CollectArtifactResponse:
+    """
+    Run an artifact collection on a client.
+
+    Args:
+        run_analyzer_body (RunAnalyzerBody): The body of the request.
+
+    Returns:
+        RunAnalyzerResponse: A dictionary containing the success status and a message.
+    """
+    velociraptor_service = await UniversalService.create("Velociraptor")
+    try:
+        # ! Can specify org_id with org_id='OL680' ! #
+        query = create_query(
+            (
+                f"SELECT collect_client("
+                f"org_id='{collect_artifact_body.velociraptor_org}', "
+                f"client_id='{collect_artifact_body.velociraptor_id}', "
+                f"artifacts=['{collect_artifact_body.artifact_name}'], "
+                f"specs=[{{"
+                f"    'parameters': {{"
+                f"        'env': ["
+                f"            {{'key': 'collectionSpec', 'value': 'Glob\\n{collect_artifact_body.file}\\n'}},"
+                f"            {{'key': 'Root', 'value': '{collect_artifact_body.root_disk}'}}"
+                f"        ]"
+                f"    }}"
+                f"}}]"
+                f") "
+                f"FROM scope()"
+            ),
+        )
+        logger.info(f"Query: {query}")
+        flow = velociraptor_service.execute_query(query, org_id=collect_artifact_body.velociraptor_org)
+        logger.info(f"Successfully ran artifact collection on {flow}")
+
+        artifact_key = get_artifact_key(analyzer_body=collect_artifact_body)
+
+        flow_id = flow["results"][0][artifact_key]["flow_id"]
+        logger.info(f"Extracted flow_id: {flow_id}")
+
+        completed = velociraptor_service.watch_flow_completion(flow_id, org_id=collect_artifact_body.velociraptor_org)
+        logger.info(f"Successfully watched flow completion on {completed}")
+
+        results = velociraptor_service.read_collection_results(
+            client_id=collect_artifact_body.velociraptor_id,
+            flow_id=flow_id,
+            org_id=collect_artifact_body.velociraptor_org,
+            artifact=collect_artifact_body.artifact_name,
+        )
+
+        logger.info(f"Successfully read collection results on {results}")
+
+        return CollectArtifactResponse(
+            success=results["success"],
+            message=results["message"],
+            results=results["results"],
+        )
+    except HTTPException as he:  # Catch HTTPException separately to propagate the original message
+        logger.error(
+            f"HTTPException while running artifact collection on {collect_artifact_body}: {he.detail}",
+        )
+        raise he
+    except Exception as err:
+        logger.error(
+            f"Failed to run artifact collection on {collect_artifact_body}: {err}",
+        )
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to run artifact collection on {collect_artifact_body}: {err}",
+        )
 
 async def run_remote_command(run_command_body: RunCommandBody) -> RunCommandResponse:
     """