| 1 | from enum import Enum |
| 2 | from typing import Any |
| 3 | from typing import Dict |
| 4 | from typing import List |
| 5 | from typing import Optional |
| 6 | from typing import Union |
| 7 | |
| 8 | from fastapi import HTTPException |
| 9 | from pydantic import BaseModel |
| 10 | from pydantic import ConfigDict |
| 11 | from pydantic import Field |
| 12 | from pydantic import field_validator |
| 13 | |
| 14 | |
| 15 | class ArtifactParameter(BaseModel): |
| 16 | """Represents a parameter definition from Velociraptor artifact.""" |
| 17 | |
| 18 | name: str = Field(..., description="Parameter name") |
| 19 | description: Optional[str] = Field(None, description="Parameter description") |
| 20 | type: Optional[str] = Field(None, description="Parameter type (e.g., 'bool', 'string')") |
| 21 | default: Optional[Union[str, bool]] = Field(None, description="Default value for the parameter") |
| 22 | |
| 23 | |
| 24 | class Artifacts(BaseModel): |
| 25 | description: str = Field(..., description="Description of the artifact.") |
| 26 | name: str = Field(..., description="Name of the artifact.") |
| 27 | author: Optional[str] = Field(None, description="Author of the artifact.") |
| 28 | precondition: Optional[str] = Field(None, description="Precondition for running the artifact.") |
| 29 | parameters: Optional[List[ArtifactParameter]] = Field( |
| 30 | None, |
| 31 | description="List of parameters that can be configured for this artifact.", |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | class ArtifactsResponse(BaseModel): |
| 36 | message: str = Field(...) |
| 37 | # make artifacts optional |
| 38 | artifacts: Optional[List[Artifacts]] = None |
| 39 | success: bool = Field(...) |
| 40 | |
| 41 | |
| 42 | class ArtifactParametersResponse(BaseModel): |
| 43 | """Response containing filtered artifact parameters.""" |
| 44 | |
| 45 | success: bool = Field(..., description="Whether the request was successful") |
| 46 | message: str = Field(..., description="Response message") |
| 47 | artifact_name: str = Field(..., description="Name of the artifact") |
| 48 | parameter_prefix: str = Field(..., description="The prefix used for filtering") |
| 49 | matching_parameters: List[ArtifactParameter] = Field(default_factory=list, description="List of parameters that match the prefix") |
| 50 | total_matches: int = Field(..., description="Total number of matching parameters") |
| 51 | model_config = ConfigDict( |
| 52 | json_schema_extra={ |
| 53 | "example": { |
| 54 | "success": True, |
| 55 | "message": "Found 2 parameters matching prefix 'T1552.001'", |
| 56 | "artifact_name": "Windows.AttackSimulation.AtomicRedTeam", |
| 57 | "parameter_prefix": "T1552.001", |
| 58 | "matching_parameters": [ |
| 59 | {"name": "T1552.001 - 3", "description": "Credentials In Files - Extracting passwords with findstr", "type": "bool"}, |
| 60 | {"name": "T1552.001 - 4", "description": "Credentials In Files - Access unattend.xml", "type": "bool"}, |
| 61 | ], |
| 62 | "total_matches": 2, |
| 63 | }, |
| 64 | }, |
| 65 | ) |
| 66 | |
| 67 | |
| 68 | class OSPrefixEnum(Enum): |
| 69 | LINUX = "Linux." |
| 70 | WINDOWS = "Windows." |
| 71 | MACOS = "MacOS." |
| 72 | |
| 73 | |
| 74 | class OSPrefixModel(BaseModel): |
| 75 | os_name: Optional[str] = None |
| 76 | os_prefix_mapping: Dict[str, str] = { |
| 77 | "windows": "Windows", |
| 78 | "linux": "Linux", |
| 79 | "mac": "MacOS", |
| 80 | "ubuntu": "Linux", # Add more mappings as needed |
| 81 | } |
| 82 | |
| 83 | def get_os_prefix(self) -> Optional[str]: |
| 84 | if self.os_name is None: |
| 85 | return None |
| 86 | return self._map_os_name_to_prefix() |
| 87 | |
| 88 | def _map_os_name_to_prefix(self) -> Optional[str]: |
| 89 | os_name_lower = self.os_name.lower() |
| 90 | for keyword, prefix in self.os_prefix_mapping.items(): |
| 91 | if keyword in os_name_lower: |
| 92 | return prefix |
| 93 | return None |
| 94 | |
| 95 | |
| 96 | class OperationEnum(str, Enum): |
| 97 | collect_artifact = "collect_artifact" |
| 98 | run_command = "run_command" |
| 99 | quarantine = "quarantine" |
| 100 | |
| 101 | |
| 102 | class ActionEnum(str, Enum): |
| 103 | quarantine = "quarantine" |
| 104 | remove_quarantine = "remove_quarantine" |
| 105 | |
| 106 | |
| 107 | class CommandArtifactsEnum(str, Enum): |
| 108 | windows_powershell = "Windows.System.PowerShell" |
| 109 | windows_cmd = "Windows.System.CmdShell" |
| 110 | linux_bash = "Linux.Sys.BashShell" |
| 111 | |
| 112 | |
| 113 | class QuarantineArtifactsEnum(str, Enum): |
| 114 | windows_quarantine = "Windows.Remediation.Quarantine" |
| 115 | linux_quarantine = "Linux.Remediation.Quarantine" |
| 116 | |
| 117 | |
| 118 | class ParameterKeyValue(BaseModel): |
| 119 | """Represents a key-value pair for artifact parameters.""" |
| 120 | |
| 121 | key: str = Field(..., description="Parameter key/name") |
| 122 | value: str = Field(..., description="Parameter value") |
| 123 | |
| 124 | |
| 125 | class BaseBody(BaseModel): |
| 126 | hostname: str = Field(..., description="Name of the client") |
| 127 | velociraptor_id: Optional[str] = Field(None, description="Client ID of the client") |
| 128 | velociraptor_org: Optional[str] = Field(None, description="Organization of the client") |
| 129 | |
| 130 | |
| 131 | class CollectArtifactBody(BaseBody): |
| 132 | """Request body for collecting artifacts with optional parameters.""" |
| 133 | |
| 134 | artifact_name: Optional[str] = Field( |
| 135 | None, |
| 136 | description="Name of the artifact for collection or command running", |
| 137 | ) |
| 138 | parameters: Optional[Dict[str, Union[str, List[ParameterKeyValue]]]] = Field( |
| 139 | None, |
| 140 | description="Optional parameters for the artifact, such as environment variables", |
| 141 | ) |
| 142 | data_store_only: Optional[bool] = Field( |
| 143 | False, |
| 144 | description="If true, only store the collected data in the datastore without sending it back immediately", |
| 145 | ) |
| 146 | model_config = ConfigDict( |
| 147 | json_schema_extra={ |
| 148 | "example": { |
| 149 | "hostname": "WIN-HFOU106TD7K", |
| 150 | "velociraptor_id": "C.475df76785008b04", |
| 151 | "velociraptor_org": "root", |
| 152 | "artifact_name": "Windows.AttackSimulation.AtomicRedTeam", |
| 153 | "parameters": {"env": [{"key": "InstallART", "value": "N"}, {"key": "T1552.001 - 3", "value": "Y"}]}, |
| 154 | "data_store_only": False, |
| 155 | }, |
| 156 | }, |
| 157 | ) |
| 158 | |
| 159 | |
| 160 | class InvokeCopilotActionBody(BaseModel): |
| 161 | """Request body for invoking a Copilot action.""" |
| 162 | |
| 163 | copilot_action_name: str = Field(..., description="Name of the action to invoke") |
| 164 | agent_names: List[str] = Field(..., description="List of agent names to invoke the action on") # Changed from agent_name to agent_names |
| 165 | artifact_name: Optional[str] = Field(None, description="Name of the artifact to use") |
| 166 | parameters: Optional[Dict[str, Union[str, List[ParameterKeyValue]]]] = Field( |
| 167 | None, |
| 168 | description="Optional parameters for the action", |
| 169 | ) |
| 170 | |
| 171 | |
| 172 | class CollectFileBody(BaseBody): |
| 173 | artifact_name: str = Field( |
| 174 | "Generic.Collectors.File", |
| 175 | description="Name of the artifact for collection or command running", |
| 176 | ) |
| 177 | file: str = Field("Glob\nUsers\\Administrator\\Documents\\*\n", description="File to collect") |
| 178 | root_disk: Optional[str] = Field("C:", description="Root disk to collect from") |
| 179 | |
| 180 | @field_validator("artifact_name") |
| 181 | @classmethod |
| 182 | def validate_artifact_name(cls, value): |
| 183 | if value != "Generic.Collectors.File": |
| 184 | raise HTTPException(status_code=400, detail="Invalid artifact name. Name should be 'Generic.Collectors.File'") |
| 185 | return value |
| 186 | |
| 187 | |
| 188 | class FileCollectionBody(BaseModel): |
| 189 | file: str = Field(..., description="File to collect") |
| 190 | root_disk: str = Field("C:", description="Root disk to collect from") |
| 191 | |
| 192 | |
| 193 | # ! Windows Example ! # |
| 194 | # { |
| 195 | # "hostname": "WIN-HFOU106TD7K", |
| 196 | # "velociraptor_id": "C.475df76785008b04", |
| 197 | # "velociraptor_org": "root", |
| 198 | # "artifact_name": "Generic.Collectors.File", |
| 199 | # "file": "Glob\nUsers\\\\Administrator\\\\Downloads\\\\LICENSE.txt\n", |
| 200 | # "root_disk": "C:" |
| 201 | # } |
| 202 | |
| 203 | # ! Linux Example ! # |
| 204 | # { |
| 205 | # "hostname": "ASHWZHMA", |
| 206 | # "velociraptor_id": "C.c4e8798fbab1d9c6", |
| 207 | # "velociraptor_org": "root", |
| 208 | # "artifact_name": "Generic.Collectors.File", |
| 209 | # "file": "Glob\n/tmp/dir/test.txt\n", |
| 210 | # "root_disk": "/" # Always use / for linux |
| 211 | # } |
| 212 | |
| 213 | |
| 214 | class RunCommandBody(BaseBody): |
| 215 | command: Optional[str] = Field(None, description="Command to run") |
| 216 | artifact_name: CommandArtifactsEnum = Field( |
| 217 | None, |
| 218 | description="Name of the artifact for command running", |
| 219 | ) |
| 220 | |
| 221 | |
| 222 | class QuarantineBody(BaseBody): |
| 223 | action: ActionEnum = Field(..., description="Action to perform") |
| 224 | artifact_name: QuarantineArtifactsEnum = Field( |
| 225 | None, |
| 226 | description="Name of the artifact for quarantine or removal of quarantine", |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | class BaseResponse(BaseModel): |
| 231 | message: str = Field(...) |
| 232 | success: bool = Field(...) # Changed from str to bool based on your sample data |
| 233 | results: Optional[List[Dict[str, Any]]] = Field( |
| 234 | None, |
| 235 | description="Results of the operation", |
| 236 | ) |
| 237 | |
| 238 | |
| 239 | class CollectArtifactResponse(BaseResponse): |
| 240 | file_info: Optional[Dict] = Field( |
| 241 | None, |
| 242 | description="Information about collected files", |
| 243 | ) |
| 244 | pass # If you have additional fields, you can define them here |
| 245 | |
| 246 | |
| 247 | class RunCommandResponse(BaseResponse): |
| 248 | pass # If you have additional fields, you can define them here |
| 249 | |
| 250 | |
| 251 | class QuarantineResponse(BaseResponse): |
| 252 | pass # If you have additional fields, you can define them here |
| 253 | |
| 254 | |
| 255 | payload = { |
| 256 | "data_win_system_eventRecordID": "521098", |
| 257 | "data_win_eventdata_user": "WIN-HFOU106TD7K\\Administrator", |
| 258 | "agent_id": "111", |
| 259 | "agent_name": "WIN-HFO106TD7K", |
| 260 | "gl2_remote_ip": "10.255.255.13", |
| 261 | "data_win_system_eventID": "22", |
| 262 | "agent_labels_customer": "00002", |
| 263 | "source": "10.255.255.13", |
| 264 | "gl2_source_input": "660320f176ca320e8393f030", |
| 265 | "rule_level": 3, |
| 266 | "data_win_system_task": "22", |
| 267 | "timestamp_utc": "2024-04-17T15:06:54.742Z", |
| 268 | "syslog_type": "wazuh", |
| 269 | "data_win_system_threadID": "2888", |
| 270 | "rule_description": "Sysmon - Event 22: DNS Request by C:\\Windows\\system32\\PING.EXE", |
| 271 | "gl2_source_node": "3b68efa4-3319-4885-a38f-c944f0fcf191", |
| 272 | "id": "1713366415.56188571", |
| 273 | "rule_mitre_tactic": "Command and Control", |
| 274 | "process_image": "C:\\Windows\\system32\\PING.EXE", |
| 275 | "data_win_eventdata_utcTime": "2024-04-17 15:06:28.457", |
| 276 | "streams": ["661555f676ca320e837b14cc", "660320f176ca320e8393f057"], |
| 277 | "rule_mitre_id": "T1071", |
| 278 | "gl2_message_id": "01HVP9HG8YE31EQH1878V50H89", |
| 279 | "data_win_system_computer": "WIN-HFOU106TD7K", |
| 280 | "agent_ip": "192.168.200.3", |
| 281 | "data_win_eventdata_image": "C:\\Windows\\system32\\PING.EXE", |
| 282 | "threat_intel_value": "evil.socfortress.co", |
| 283 | "data_win_eventdata_queryName": "evil.socfortress.co", |
| 284 | "rule_groups": "windows, sysmon, sysmon_event_22", |
| 285 | "data_win_system_keywords": "0x8000000000000000", |
| 286 | "data_win_system_level": "4", |
| 287 | "process_id": "6072", |
| 288 | "data_win_eventdata_queryStatus": "0", |
| 289 | "data_win_system_severityValue": "INFORMATION", |
| 290 | "dns_response_code": "0", |
| 291 | "dns_query": "evil.socfortress.co", |
| 292 | "data_win_eventdata_processGuid": "{691ff406-e58c-661f-b401-000000002300}", |
| 293 | "rule_mitre_technique": "Application Layer Protocol", |
| 294 | "rule_firedtimes": 2, |
| 295 | "data_win_system_systemTime": "2024-04-17T15:06:54.742696000Z", |
| 296 | "decoder_name": "windows_eventchannel", |
| 297 | "data_win_system_processID": "2180", |
| 298 | "data_win_system_channel": "Microsoft-Windows-Sysmon/Operational", |
| 299 | "syslog_level": "ALERT", |
| 300 | "threat_intel_comment": "This is a test IoC", |
| 301 | "data_win_system_providerName": "Microsoft-Windows-Sysmon", |
| 302 | "data_win_eventdata_processId": "6072", |
| 303 | "data_win_system_version": "5", |
| 304 | "data_win_system_providerGuid": "{5770385f-c22a-43e0-bf4c-06f5698ffbd9}", |
| 305 | "timestamp": "2024-04-17 15:06:57.694", |
| 306 | "threat_intel_ioc_source": "test", |
| 307 | "rule_group1": "windows", |
| 308 | "data_win_system_opcode": "0", |
| 309 | } |
| 310 | |
| 311 | |
| 312 | class OS(str, Enum): |
| 313 | Windows = "Windows" |
| 314 | Linux = "Linux" |
| 315 | MacOS = "MacOS" |
| 316 | |
| 317 | |
| 318 | class ArtifactReccomendationAIRequest(BaseModel): |
| 319 | os: OS = Field(..., description="Operating system of the client") |
| 320 | prompt: dict = Field(..., examples=[payload]) |
| 321 | |
| 322 | |
| 323 | class ArtifactReccomendationRequest(BaseModel): |
| 324 | artifacts: List[Artifacts] = Field(..., description="List of artifacts to be recommended") |
| 325 | os: str = Field(..., description="Operating system of the client") |
| 326 | prompt: dict = Field(..., examples=[payload]) |
| 327 | |
| 328 | |
| 329 | class VelociraptorArtifactRecommendation(BaseModel): |
| 330 | name: str = Field(..., description="The name of the artifact.") |
| 331 | description: str = Field(..., description="A description of the artifact.") |
| 332 | explanation: str = Field(..., description="A detailed explanation of the purpose and why the artifact was selected.") |
| 333 | |
| 334 | |
| 335 | class ArtifactReccomendationResponse(BaseModel): |
| 336 | message: str = Field(...) |
| 337 | success: bool = Field(...) |
| 338 | recommendations: list[VelociraptorArtifactRecommendation] |