| 1 | from datetime import datetime |
| 2 | from typing import Optional |
| 3 | |
| 4 | from pydantic import BaseModel |
| 5 | from pydantic import ConfigDict |
| 6 | from pydantic import Field |
| 7 | |
| 8 | |
| 9 | class FileUploadResponse(BaseModel): |
| 10 | success: bool |
| 11 | message: str |
| 12 | bucket_name: str |
| 13 | object_key: str |
| 14 | file_name: str |
| 15 | file_size: int |
| 16 | file_hash: str |
| 17 | content_type: Optional[str] = None |
| 18 | |
| 19 | |
| 20 | class CaseDataStoreCreation(BaseModel): |
| 21 | case_id: int |
| 22 | bucket_name: str = Field(max_length=255) |
| 23 | object_key: str = Field(max_length=1024) |
| 24 | file_name: str = Field(max_length=255) |
| 25 | content_type: Optional[str] = Field(max_length=100, default=None) |
| 26 | file_size: Optional[int] = None |
| 27 | upload_time: datetime = Field(default_factory=datetime.utcnow) |
| 28 | file_hash: str = Field(max_length=128) |
| 29 | |
| 30 | |
| 31 | class CaseReportTemplateDataStoreCreation(BaseModel): |
| 32 | report_template_name: str = Field(max_length=255) |
| 33 | bucket_name: str = Field(max_length=255) |
| 34 | object_key: str = Field(max_length=1024) |
| 35 | file_name: str = Field(max_length=255) |
| 36 | content_type: Optional[str] = Field(max_length=100, default=None) |
| 37 | upload_time: datetime = Field(default_factory=datetime.utcnow) |
| 38 | file_hash: str = Field(max_length=128) |
| 39 | |
| 40 | |
| 41 | class AgentDataStoreCreation(BaseModel): |
| 42 | agent_id: str |
| 43 | velociraptor_id: str |
| 44 | artifact_name: str |
| 45 | flow_id: str |
| 46 | bucket_name: str = "velociraptor-artifacts" |
| 47 | file_name: str |
| 48 | content_type: str = "application/zip" |
| 49 | file_size: int |
| 50 | file_hash: str |
| 51 | uploaded_by: Optional[int] = None |
| 52 | notes: Optional[str] = None |
| 53 | status: str = "completed" |
| 54 | |
| 55 | |
| 56 | class AgentDataStoreData(BaseModel): |
| 57 | id: int |
| 58 | agent_id: str |
| 59 | velociraptor_id: str |
| 60 | customer_code: Optional[str] = None # Made optional since it comes from relationship |
| 61 | artifact_name: str |
| 62 | flow_id: str |
| 63 | bucket_name: str |
| 64 | object_key: str |
| 65 | file_name: str |
| 66 | content_type: str |
| 67 | file_size: int |
| 68 | file_hash: str |
| 69 | collection_time: datetime |
| 70 | uploaded_by: Optional[int] = None |
| 71 | notes: Optional[str] = None |
| 72 | status: str |
| 73 | model_config = ConfigDict(from_attributes=True) |
| 74 | |
| 75 | |
| 76 | class AgentDataStoreResponse(BaseModel): |
| 77 | success: bool |
| 78 | message: str |
| 79 | data: Optional[AgentDataStoreData] = None # Removed the string quotes here |
| 80 | |
| 81 | |
| 82 | class AgentDataStoreListResponse(BaseModel): |
| 83 | success: bool |
| 84 | message: str |
| 85 | data: list[AgentDataStoreData] = [] |
| 86 | total: int = 0 |