| 1 | from typing import List |
| 2 | from typing import Optional |
| 3 | from typing import Union |
| 4 | |
| 5 | from fastapi import HTTPException |
| 6 | from pydantic import BaseModel |
| 7 | from pydantic import ConfigDict |
| 8 | from pydantic import Field |
| 9 | from pydantic import field_validator |
| 10 | |
| 11 | |
| 12 | class RuleDisable(BaseModel): |
| 13 | rule_id: str |
| 14 | reason_for_disabling: str |
| 15 | length_of_time: str |
| 16 | |
| 17 | |
| 18 | class RuleDisableResponse(BaseModel): |
| 19 | previous_level: Optional[str] = None |
| 20 | message: str |
| 21 | success: bool |
| 22 | |
| 23 | |
| 24 | class RuleEnable(BaseModel): |
| 25 | rule_id: str |
| 26 | reason_for_enabling: str |
| 27 | |
| 28 | |
| 29 | class RuleEnableResponse(BaseModel): |
| 30 | new_level: Optional[str] = None |
| 31 | message: str |
| 32 | success: bool |
| 33 | |
| 34 | |
| 35 | class AllDisabledRule(BaseModel): |
| 36 | rule_id: str |
| 37 | previous_level: str |
| 38 | new_level: str |
| 39 | reason_for_disabling: str |
| 40 | length_of_time: str |
| 41 | disabled_by: str |
| 42 | |
| 43 | |
| 44 | class AllDisabledRuleResponse(BaseModel): |
| 45 | disabled_rules: List[AllDisabledRule] |
| 46 | success: bool |
| 47 | message: str |
| 48 | |
| 49 | |
| 50 | # Simplified and more efficient rule models |
| 51 | class WazuhRule(BaseModel): |
| 52 | """Represents a single Wazuh rule from the API.""" |
| 53 | |
| 54 | filename: str |
| 55 | relative_dirname: str |
| 56 | id: int |
| 57 | level: int |
| 58 | status: str |
| 59 | description: str |
| 60 | groups: List[str] = [] |
| 61 | pci_dss: List[str] = [] |
| 62 | gpg13: List[str] = [] |
| 63 | gdpr: List[str] = [] |
| 64 | hipaa: List[str] = [] |
| 65 | nist_800_53: List[str] = Field(default=[], alias="nist-800-53") |
| 66 | tsc: List[str] = [] |
| 67 | mitre: List[str] = [] |
| 68 | details: Optional[dict] = None |
| 69 | model_config = ConfigDict(populate_by_name=True, extra="ignore") |
| 70 | |
| 71 | |
| 72 | class WazuhRulesResponse(BaseModel): |
| 73 | """Response model for Wazuh rules listing.""" |
| 74 | |
| 75 | success: bool |
| 76 | message: str |
| 77 | results: List[WazuhRule] = [] |
| 78 | total_items: Optional[int] = None |
| 79 | |
| 80 | |
| 81 | class WazuhRuleFile(BaseModel): |
| 82 | """Represents a single Wazuh rule file from the API.""" |
| 83 | |
| 84 | filename: str = Field(..., description="Rule file name") |
| 85 | relative_dirname: str = Field(..., description="Relative directory path") |
| 86 | status: str = Field(..., description="File status (enabled/disabled)") |
| 87 | model_config = ConfigDict(extra="ignore") |
| 88 | |
| 89 | |
| 90 | class WazuhRuleFilesResponse(BaseModel): |
| 91 | """Response model for Wazuh rule files listing.""" |
| 92 | |
| 93 | success: bool = Field(..., description="Whether the request was successful") |
| 94 | message: str = Field(..., description="Response message") |
| 95 | results: List[WazuhRuleFile] = Field(default=[], description="List of rule files") |
| 96 | total_items: Optional[int] = Field(None, description="Total number of files") |
| 97 | |
| 98 | |
| 99 | class WazuhRuleFileContentResponse(BaseModel): |
| 100 | """Response model for Wazuh rule file content.""" |
| 101 | |
| 102 | success: bool = Field(..., description="Whether the request was successful") |
| 103 | message: str = Field(..., description="Response message") |
| 104 | filename: str = Field(..., description="The requested filename") |
| 105 | content: Union[dict, str] = Field(..., description="File content (structured or raw)") |
| 106 | is_raw: bool = Field(False, description="Whether the content is raw text") |
| 107 | total_items: Optional[int] = Field(None, description="Total affected items from API") |
| 108 | |
| 109 | |
| 110 | class WazuhRuleFileUploadResponse(BaseModel): |
| 111 | """Response model for Wazuh rule file upload/update.""" |
| 112 | |
| 113 | success: bool = Field(..., description="Whether the request was successful") |
| 114 | message: str = Field(..., description="Response message") |
| 115 | filename: str = Field(..., description="The uploaded/updated filename") |
| 116 | details: Optional[dict] = Field(None, description="Additional response details from API") |
| 117 | total_items: Optional[int] = Field(None, description="Total affected items from API") |
| 118 | |
| 119 | |
| 120 | payload = { |
| 121 | "data_win_system_eventRecordID": "521098", |
| 122 | "data_win_eventdata_user": "WIN-HFOU106TD7K\\Administrator", |
| 123 | "agent_id": "111", |
| 124 | "agent_name": "WIN-HFO106TD7K", |
| 125 | "gl2_remote_ip": "10.255.255.13", |
| 126 | "data_win_system_eventID": "22", |
| 127 | "agent_labels_customer": "00002", |
| 128 | "source": "10.255.255.13", |
| 129 | "gl2_source_input": "660320f176ca320e8393f030", |
| 130 | "rule_level": 3, |
| 131 | "data_win_system_task": "22", |
| 132 | "timestamp_utc": "2024-04-17T15:06:54.742Z", |
| 133 | "syslog_type": "wazuh", |
| 134 | "data_win_system_threadID": "2888", |
| 135 | "rule_description": "Sysmon - Event 22: DNS Request by C:\\Windows\\system32\\PING.EXE", |
| 136 | "gl2_source_node": "3b68efa4-3319-4885-a38f-c944f0fcf191", |
| 137 | "id": "1713366415.56188571", |
| 138 | "rule_mitre_tactic": "Command and Control", |
| 139 | "process_image": "C:\\Windows\\system32\\PING.EXE", |
| 140 | "data_win_eventdata_utcTime": "2024-04-17 15:06:28.457", |
| 141 | "streams": ["661555f676ca320e837b14cc", "660320f176ca320e8393f057"], |
| 142 | "rule_mitre_id": "T1071", |
| 143 | "gl2_message_id": "01HVP9HG8YE31EQH1878V50H89", |
| 144 | "data_win_system_computer": "WIN-HFOU106TD7K", |
| 145 | "agent_ip": "192.168.200.3", |
| 146 | "data_win_eventdata_image": "C:\\Windows\\system32\\PING.EXE", |
| 147 | "threat_intel_value": "evil.socfortress.co", |
| 148 | "data_win_eventdata_queryName": "evil.socfortress.co", |
| 149 | "rule_groups": "windows, sysmon, sysmon_event_22", |
| 150 | "data_win_system_keywords": "0x8000000000000000", |
| 151 | "data_win_system_level": "4", |
| 152 | "process_id": "6072", |
| 153 | "data_win_eventdata_queryStatus": "0", |
| 154 | "data_win_system_severityValue": "INFORMATION", |
| 155 | "dns_response_code": "0", |
| 156 | "dns_query": "evil.socfortress.co", |
| 157 | "data_win_eventdata_processGuid": "{691ff406-e58c-661f-b401-000000002300}", |
| 158 | "rule_mitre_technique": "Application Layer Protocol", |
| 159 | "rule_firedtimes": 2, |
| 160 | "data_win_system_systemTime": "2024-04-17T15:06:54.742696000Z", |
| 161 | "decoder_name": "windows_eventchannel", |
| 162 | "data_win_system_processID": "2180", |
| 163 | "data_win_system_channel": "Microsoft-Windows-Sysmon/Operational", |
| 164 | "syslog_level": "ALERT", |
| 165 | "threat_intel_comment": "This is a test IoC", |
| 166 | "data_win_system_providerName": "Microsoft-Windows-Sysmon", |
| 167 | "data_win_eventdata_processId": "6072", |
| 168 | "data_win_system_version": "5", |
| 169 | "data_win_system_providerGuid": "{5770385f-c22a-43e0-bf4c-06f5698ffbd9}", |
| 170 | "timestamp": "2024-04-17 15:06:57.694", |
| 171 | "threat_intel_ioc_source": "test", |
| 172 | "rule_group1": "windows", |
| 173 | "data_win_system_opcode": "0", |
| 174 | } |
| 175 | |
| 176 | |
| 177 | class RuleExcludeRequest(BaseModel): |
| 178 | integration: str = Field(..., examples=["wazuh-rule-exclusion"]) |
| 179 | prompt: dict = Field(..., examples=[payload]) |
| 180 | |
| 181 | @field_validator("integration") |
| 182 | @classmethod |
| 183 | def check_integration(cls, v): |
| 184 | if v != "wazuh-rule-exclusion": |
| 185 | raise HTTPException( |
| 186 | status_code=400, |
| 187 | detail="Invalid integration. Only 'wazuh-rule-exclusion' is supported.", |
| 188 | ) |
| 189 | return v |
| 190 | |
| 191 | @field_validator("prompt") |
| 192 | @classmethod |
| 193 | def check_rule_group(cls, v): |
| 194 | if "rule_group3" in v: |
| 195 | if "rule_group1" not in v and "rule_group3" not in v: |
| 196 | raise HTTPException( |
| 197 | status_code=400, |
| 198 | detail="Missing 'rule_group1' or 'rule_group3' in prompt.", |
| 199 | ) |
| 200 | if ("rule_group1" in v and v["rule_group1"] != "windows") and ("rule_group3" in v and v["rule_group3"] != "windows"): |
| 201 | raise HTTPException( |
| 202 | status_code=400, |
| 203 | detail="Invalid 'rule_group1' or 'rule_group3'. At least one must be 'windows'.", |
| 204 | ) |
| 205 | else: |
| 206 | if "rule_group1" not in v: |
| 207 | raise HTTPException( |
| 208 | status_code=400, |
| 209 | detail="Missing 'rule_group1' in prompt.", |
| 210 | ) |
| 211 | if v["rule_group1"] != "windows": |
| 212 | raise HTTPException( |
| 213 | status_code=400, |
| 214 | detail="Invalid 'rule_group1'. Only 'windows' is supported.", |
| 215 | ) |
| 216 | return v |
| 217 | |
| 218 | |
| 219 | class RuleExcludeResponse(BaseModel): |
| 220 | wazuh_rule: str |
| 221 | explanation: str |
| 222 | message: str |
| 223 | success: bool |