| 1 | import io |
| 2 | from typing import Optional |
| 3 | |
| 4 | from fastapi import APIRouter |
| 5 | from fastapi import Depends |
| 6 | from fastapi import File |
| 7 | from fastapi import Form |
| 8 | from fastapi import HTTPException |
| 9 | from fastapi import UploadFile |
| 10 | from fastapi import status |
| 11 | from fastapi.responses import StreamingResponse |
| 12 | from loguru import logger |
| 13 | from sqlalchemy import desc |
| 14 | from sqlalchemy import select |
| 15 | from sqlalchemy.ext.asyncio import AsyncSession |
| 16 | from sqlalchemy.orm import selectinload |
| 17 | |
| 18 | from app.auth.utils import AuthHandler |
| 19 | from app.data_store.data_store_operations import delete_agent_artifact_file |
| 20 | from app.data_store.data_store_operations import download_agent_artifact_file |
| 21 | from app.data_store.data_store_operations import upload_file_to_datastore |
| 22 | from app.data_store.data_store_schema import AgentDataStoreData |
| 23 | from app.data_store.data_store_schema import AgentDataStoreListResponse |
| 24 | from app.data_store.data_store_schema import AgentDataStoreResponse |
| 25 | from app.data_store.data_store_schema import FileUploadResponse |
| 26 | from app.db.db_session import get_db |
| 27 | from app.db.universal_models import AgentDataStore |
| 28 | |
| 29 | agent_data_store_router = APIRouter() |
| 30 | |
| 31 | |
| 32 | @agent_data_store_router.post( |
| 33 | "/upload", |
| 34 | response_model=FileUploadResponse, |
| 35 | description="Upload a file to the data store", |
| 36 | dependencies=[Depends(AuthHandler().require_any_scope("admin", "analyst", "customer_user"))], |
| 37 | ) |
| 38 | async def upload_file( |
| 39 | file: UploadFile = File(...), |
| 40 | bucket_name: str = Form(...), |
| 41 | object_name: str = Form(...), |
| 42 | ) -> FileUploadResponse: |
| 43 | """ |
| 44 | Upload a file to the specified bucket and object path in the data store. |
| 45 | |
| 46 | Args: |
| 47 | file: The file to upload |
| 48 | bucket_name: The name of the bucket to upload to |
| 49 | object_name: The object path/name within the bucket (e.g., "folder/subfolder/file.txt") |
| 50 | |
| 51 | Returns: |
| 52 | FileUploadResponse with upload details |
| 53 | """ |
| 54 | try: |
| 55 | logger.info(f"Uploading file {file.filename} to bucket {bucket_name} as {object_name}") |
| 56 | |
| 57 | if not file.filename: |
| 58 | raise HTTPException( |
| 59 | status_code=status.HTTP_400_BAD_REQUEST, |
| 60 | detail="File name is required", |
| 61 | ) |
| 62 | |
| 63 | # Upload the file to MinIO |
| 64 | upload_result = await upload_file_to_datastore( |
| 65 | file=file, |
| 66 | bucket_name=bucket_name, |
| 67 | object_name=object_name, |
| 68 | ) |
| 69 | |
| 70 | return FileUploadResponse( |
| 71 | success=True, |
| 72 | message=f"File {file.filename} uploaded successfully", |
| 73 | bucket_name=upload_result["bucket_name"], |
| 74 | object_key=upload_result["object_key"], |
| 75 | file_name=upload_result["file_name"], |
| 76 | file_size=upload_result["file_size"], |
| 77 | file_hash=upload_result["file_hash"], |
| 78 | content_type=upload_result.get("content_type"), |
| 79 | ) |
| 80 | |
| 81 | except Exception as e: |
| 82 | logger.error(f"Failed to upload file: {e}") |
| 83 | raise HTTPException( |
| 84 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 85 | detail=f"Failed to upload file: {str(e)}", |
| 86 | ) |
| 87 | |
| 88 | |
| 89 | @agent_data_store_router.get( |
| 90 | "/agent/{agent_id}/artifacts", |
| 91 | response_model=AgentDataStoreListResponse, |
| 92 | description="List all artifact files for an agent", |
| 93 | dependencies=[Depends(AuthHandler().require_any_scope("admin", "analyst"))], |
| 94 | ) |
| 95 | async def list_agent_artifacts( |
| 96 | agent_id: str, |
| 97 | flow_id: Optional[str] = None, |
| 98 | session: AsyncSession = Depends(get_db), |
| 99 | ) -> AgentDataStoreListResponse: |
| 100 | """ |
| 101 | List all artifact collection files for a specific agent. |
| 102 | Optionally filter by flow_id. |
| 103 | """ |
| 104 | try: |
| 105 | query = select(AgentDataStore).where(AgentDataStore.agent_id == agent_id) |
| 106 | query = query.options(selectinload(AgentDataStore.agent)) # Eager load agent |
| 107 | |
| 108 | if flow_id: |
| 109 | query = query.where(AgentDataStore.flow_id == flow_id) |
| 110 | |
| 111 | query = query.order_by(desc(AgentDataStore.collection_time)) |
| 112 | |
| 113 | result = await session.execute(query) |
| 114 | artifacts = result.scalars().all() |
| 115 | |
| 116 | # Convert to response data, accessing customer_code via relationship |
| 117 | artifact_data = [] |
| 118 | for artifact in artifacts: |
| 119 | # Create a dict from the artifact |
| 120 | artifact_dict = { |
| 121 | "id": artifact.id, |
| 122 | "agent_id": artifact.agent_id, |
| 123 | "velociraptor_id": artifact.velociraptor_id, |
| 124 | "customer_code": artifact.agent.customer_code if artifact.agent else None, |
| 125 | "artifact_name": artifact.artifact_name, |
| 126 | "flow_id": artifact.flow_id, |
| 127 | "bucket_name": artifact.bucket_name, |
| 128 | "object_key": artifact.object_key, |
| 129 | "file_name": artifact.file_name, |
| 130 | "content_type": artifact.content_type, |
| 131 | "file_size": artifact.file_size, |
| 132 | "file_hash": artifact.file_hash, |
| 133 | "collection_time": artifact.collection_time, |
| 134 | "uploaded_by": artifact.uploaded_by, |
| 135 | "notes": artifact.notes, |
| 136 | "status": artifact.status, |
| 137 | } |
| 138 | artifact_data.append(AgentDataStoreData(**artifact_dict)) |
| 139 | |
| 140 | return AgentDataStoreListResponse( |
| 141 | success=True, |
| 142 | message=f"Found {len(artifact_data)} artifacts for agent {agent_id}", |
| 143 | data=artifact_data, |
| 144 | total=len(artifact_data), |
| 145 | ) |
| 146 | |
| 147 | except Exception as e: |
| 148 | logger.error(f"Failed to list agent artifacts: {e}") |
| 149 | raise HTTPException( |
| 150 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 151 | detail=f"Failed to list agent artifacts: {str(e)}", |
| 152 | ) |
| 153 | |
| 154 | |
| 155 | @agent_data_store_router.get( |
| 156 | "/agent/{agent_id}/artifacts/{artifact_id}", |
| 157 | response_model=AgentDataStoreResponse, |
| 158 | description="Get details of a specific artifact file", |
| 159 | dependencies=[Depends(AuthHandler().require_any_scope("admin", "analyst"))], |
| 160 | ) |
| 161 | async def get_agent_artifact_details( |
| 162 | agent_id: str, |
| 163 | artifact_id: int, |
| 164 | session: AsyncSession = Depends(get_db), |
| 165 | ) -> AgentDataStoreResponse: |
| 166 | """Get details of a specific artifact collection file.""" |
| 167 | try: |
| 168 | query = select(AgentDataStore).where(AgentDataStore.id == artifact_id, AgentDataStore.agent_id == agent_id) |
| 169 | query = query.options(selectinload(AgentDataStore.agent)) |
| 170 | |
| 171 | result = await session.execute(query) |
| 172 | artifact = result.scalars().first() |
| 173 | |
| 174 | if not artifact: |
| 175 | raise HTTPException( |
| 176 | status_code=status.HTTP_404_NOT_FOUND, |
| 177 | detail=f"Artifact {artifact_id} not found for agent {agent_id}", |
| 178 | ) |
| 179 | |
| 180 | # Create response data with customer_code from relationship |
| 181 | artifact_dict = { |
| 182 | "id": artifact.id, |
| 183 | "agent_id": artifact.agent_id, |
| 184 | "velociraptor_id": artifact.velociraptor_id, |
| 185 | "customer_code": artifact.agent.customer_code if artifact.agent else None, |
| 186 | "artifact_name": artifact.artifact_name, |
| 187 | "flow_id": artifact.flow_id, |
| 188 | "bucket_name": artifact.bucket_name, |
| 189 | "object_key": artifact.object_key, |
| 190 | "file_name": artifact.file_name, |
| 191 | "content_type": artifact.content_type, |
| 192 | "file_size": artifact.file_size, |
| 193 | "file_hash": artifact.file_hash, |
| 194 | "collection_time": artifact.collection_time, |
| 195 | "uploaded_by": artifact.uploaded_by, |
| 196 | "notes": artifact.notes, |
| 197 | "status": artifact.status, |
| 198 | } |
| 199 | |
| 200 | return AgentDataStoreResponse( |
| 201 | success=True, |
| 202 | message="Artifact details retrieved successfully", |
| 203 | data=AgentDataStoreData(**artifact_dict), |
| 204 | ) |
| 205 | |
| 206 | except HTTPException: |
| 207 | raise |
| 208 | except Exception as e: |
| 209 | logger.error(f"Failed to get artifact details: {e}") |
| 210 | raise HTTPException( |
| 211 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 212 | detail=f"Failed to get artifact details: {str(e)}", |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | @agent_data_store_router.get( |
| 217 | "/agent/{agent_id}/artifacts/{artifact_id}/download", |
| 218 | description="Download an artifact file", |
| 219 | dependencies=[Depends(AuthHandler().require_any_scope("admin", "analyst"))], |
| 220 | ) |
| 221 | async def download_agent_artifact( |
| 222 | agent_id: str, |
| 223 | artifact_id: int, |
| 224 | session: AsyncSession = Depends(get_db), |
| 225 | ): |
| 226 | """Download a specific artifact collection file.""" |
| 227 | try: |
| 228 | result = await session.execute(select(AgentDataStore).where(AgentDataStore.id == artifact_id, AgentDataStore.agent_id == agent_id)) |
| 229 | artifact = result.scalars().first() |
| 230 | |
| 231 | if not artifact: |
| 232 | raise HTTPException( |
| 233 | status_code=status.HTTP_404_NOT_FOUND, |
| 234 | detail=f"Artifact {artifact_id} not found for agent {agent_id}", |
| 235 | ) |
| 236 | |
| 237 | # Download from MinIO |
| 238 | file_data = await download_agent_artifact_file( |
| 239 | agent_id=agent_id, |
| 240 | flow_id=artifact.flow_id, |
| 241 | file_name=artifact.file_name, |
| 242 | ) |
| 243 | |
| 244 | return StreamingResponse( |
| 245 | io.BytesIO(file_data), |
| 246 | media_type=artifact.content_type, |
| 247 | headers={ |
| 248 | "Content-Disposition": f"attachment; filename={artifact.file_name}", |
| 249 | "Content-Length": str(artifact.file_size), |
| 250 | }, |
| 251 | ) |
| 252 | |
| 253 | except HTTPException: |
| 254 | raise |
| 255 | except Exception as e: |
| 256 | logger.error(f"Failed to download artifact: {e}") |
| 257 | raise HTTPException( |
| 258 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 259 | detail=f"Failed to download artifact: {str(e)}", |
| 260 | ) |
| 261 | |
| 262 | |
| 263 | @agent_data_store_router.delete( |
| 264 | "/agent/{agent_id}/artifacts/{artifact_id}", |
| 265 | description="Delete an artifact file", |
| 266 | dependencies=[Depends(AuthHandler().require_any_scope("admin"))], |
| 267 | ) |
| 268 | async def delete_agent_artifact( |
| 269 | agent_id: str, |
| 270 | artifact_id: int, |
| 271 | session: AsyncSession = Depends(get_db), |
| 272 | ): |
| 273 | """Delete a specific artifact collection file.""" |
| 274 | try: |
| 275 | result = await session.execute(select(AgentDataStore).where(AgentDataStore.id == artifact_id, AgentDataStore.agent_id == agent_id)) |
| 276 | artifact = result.scalars().first() |
| 277 | |
| 278 | if not artifact: |
| 279 | raise HTTPException( |
| 280 | status_code=status.HTTP_404_NOT_FOUND, |
| 281 | detail=f"Artifact {artifact_id} not found for agent {agent_id}", |
| 282 | ) |
| 283 | |
| 284 | # Delete from MinIO |
| 285 | await delete_agent_artifact_file( |
| 286 | agent_id=agent_id, |
| 287 | flow_id=artifact.flow_id, |
| 288 | file_name=artifact.file_name, |
| 289 | ) |
| 290 | |
| 291 | # Delete from database |
| 292 | await session.delete(artifact) |
| 293 | await session.commit() |
| 294 | |
| 295 | logger.info(f"Deleted artifact {artifact_id} for agent {agent_id}") |
| 296 | |
| 297 | return { |
| 298 | "success": True, |
| 299 | "message": f"Artifact {artifact_id} deleted successfully", |
| 300 | } |
| 301 | |
| 302 | except HTTPException: |
| 303 | raise |
| 304 | except Exception as e: |
| 305 | logger.error(f"Failed to delete artifact: {e}") |
| 306 | await session.rollback() |
| 307 | raise HTTPException( |
| 308 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 309 | detail=f"Failed to delete artifact: {str(e)}", |
| 310 | ) |