| 1 | from typing import List |
| 2 | from typing import Optional |
| 3 | |
| 4 | from pydantic import BaseModel |
| 5 | from pydantic import ConfigDict |
| 6 | from pydantic import Field |
| 7 | |
| 8 | |
| 9 | # ! INDEX SETS ! # |
| 10 | class TimeBasedRotationStrategyConfig(BaseModel): |
| 11 | type: str |
| 12 | rotation_period: Optional[str] = None |
| 13 | |
| 14 | |
| 15 | class TimeBasedRetentionStrategyConfig(BaseModel): |
| 16 | type: str |
| 17 | max_number_of_indices: Optional[int] = None |
| 18 | |
| 19 | |
| 20 | class TimeBasedIndexSet(BaseModel): |
| 21 | title: str |
| 22 | description: str |
| 23 | index_prefix: str |
| 24 | rotation_strategy_class: str |
| 25 | rotation_strategy: TimeBasedRotationStrategyConfig |
| 26 | retention_strategy_class: str |
| 27 | retention_strategy: TimeBasedRetentionStrategyConfig |
| 28 | creation_date: str |
| 29 | index_analyzer: str |
| 30 | shards: int |
| 31 | replicas: int |
| 32 | index_optimization_max_num_segments: int |
| 33 | index_optimization_disabled: bool |
| 34 | writable: bool |
| 35 | field_type_refresh_interval: int |
| 36 | model_config = ConfigDict( |
| 37 | json_schema_extra={ |
| 38 | "example": { |
| 39 | "title": "Wazuh - Example Company", |
| 40 | "description": "Wazuh - Example Company", |
| 41 | "index_prefix": "wazuh-examplecode", |
| 42 | "rotation_strategy_class": "org.graylog2.indexer.rotation.strategies.SizeBasedRotationStrategy", |
| 43 | "rotation_strategy": { |
| 44 | "type": "org.graylog2.indexer.rotation.strategies.SizeBasedRotationStrategyConfig", |
| 45 | "max_size": 2684354560, |
| 46 | }, |
| 47 | "retention_strategy_class": "org.graylog2.indexer.retention.strategies.DeletionRetentionStrategy", |
| 48 | "retention_strategy": { |
| 49 | "type": "org.graylog2.indexer.retention.strategies.DeletionRetentionStrategyConfig", |
| 50 | "max_number_of_indices": 20, |
| 51 | }, |
| 52 | "creation_date": "2021-01-01T00:00:00.000Z", |
| 53 | "index_analyzer": "standard", |
| 54 | "shards": 1, |
| 55 | "replicas": 0, |
| 56 | "index_optimization_max_num_segments": 1, |
| 57 | "index_optimization_disabled": False, |
| 58 | "writable": True, |
| 59 | "field_type_refresh_interval": 5000, |
| 60 | }, |
| 61 | }, |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | class RotationStrategyConfig(BaseModel): |
| 66 | type: str |
| 67 | rotation_period: str = Field(..., alias="rotation_period") |
| 68 | max_rotation_period: Optional[str] = Field(None, alias="max_rotation_period") |
| 69 | rotate_empty_index_set: Optional[bool] = Field(None, alias="rotate_empty_index_set") |
| 70 | |
| 71 | |
| 72 | class RetentionStrategyConfig(BaseModel): |
| 73 | type: str |
| 74 | max_number_of_indices: int = Field(..., alias="max_number_of_indices") |
| 75 | |
| 76 | |
| 77 | class GraylogIndexSetData(BaseModel): |
| 78 | id: str |
| 79 | title: str |
| 80 | description: str |
| 81 | can_be_default: bool = Field(..., alias="can_be_default") |
| 82 | index_prefix: str = Field(..., alias="index_prefix") |
| 83 | shards: int |
| 84 | replicas: int |
| 85 | rotation_strategy_class: str = Field(..., alias="rotation_strategy_class") |
| 86 | rotation_strategy: RotationStrategyConfig |
| 87 | retention_strategy_class: str = Field(..., alias="retention_strategy_class") |
| 88 | retention_strategy: RetentionStrategyConfig |
| 89 | creation_date: str = Field(..., alias="creation_date") |
| 90 | index_analyzer: str = Field(..., alias="index_analyzer") |
| 91 | index_optimization_max_num_segments: int = Field( |
| 92 | ..., |
| 93 | alias="index_optimization_max_num_segments", |
| 94 | ) |
| 95 | index_optimization_disabled: bool = Field(..., alias="index_optimization_disabled") |
| 96 | field_type_refresh_interval: int = Field(..., alias="field_type_refresh_interval") |
| 97 | index_template_type: Optional[str] = Field(None, alias="index_template_type") |
| 98 | default: bool |
| 99 | writable: bool |
| 100 | |
| 101 | |
| 102 | class GraylogIndexSetCreationResponse(BaseModel): |
| 103 | data: GraylogIndexSetData |
| 104 | success: bool |
| 105 | message: str |
| 106 | |
| 107 | |
| 108 | # ! STREAMS ! # |
| 109 | class StreamRule(BaseModel): |
| 110 | field: str |
| 111 | type: int |
| 112 | inverted: bool |
| 113 | value: str |
| 114 | |
| 115 | |
| 116 | class WazuhEventStream(BaseModel): |
| 117 | title: str = Field(..., description="Title of the stream") |
| 118 | description: str = Field(..., description="Description of the stream") |
| 119 | index_set_id: str = Field(..., description="ID of the associated index set") |
| 120 | rules: List[StreamRule] = Field(..., description="List of rules for the stream") |
| 121 | matching_type: str = Field(..., description="Matching type for the rules") |
| 122 | remove_matches_from_default_stream: bool = Field( |
| 123 | ..., |
| 124 | description="Whether to remove matches from the default stream", |
| 125 | ) |
| 126 | content_pack: Optional[str] = Field( |
| 127 | None, |
| 128 | description="Associated content pack, if any", |
| 129 | ) |
| 130 | model_config = ConfigDict( |
| 131 | json_schema_extra={ |
| 132 | "example": { |
| 133 | "title": "WAZUH EVENTS CUSTOMERS - Example Company", |
| 134 | "description": "WAZUH EVENTS CUSTOMERS - Example Company", |
| 135 | "index_set_id": "12345", |
| 136 | "rules": [ |
| 137 | { |
| 138 | "field": "agent_labels_customer", |
| 139 | "type": 1, |
| 140 | "inverted": False, |
| 141 | "value": "ExampleCode", |
| 142 | }, |
| 143 | ], |
| 144 | "matching_type": "AND", |
| 145 | "remove_matches_from_default_stream": True, |
| 146 | "content_pack": None, |
| 147 | }, |
| 148 | }, |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | class Office365EventStream(BaseModel): |
| 153 | title: str = Field(..., description="Title of the stream") |
| 154 | description: str = Field(..., description="Description of the stream") |
| 155 | index_set_id: str = Field(..., description="ID of the associated index set") |
| 156 | rules: List[StreamRule] = Field(..., description="List of rules for the stream") |
| 157 | matching_type: str = Field(..., description="Matching type for the rules") |
| 158 | remove_matches_from_default_stream: bool = Field( |
| 159 | ..., |
| 160 | description="Whether to remove matches from the default stream", |
| 161 | ) |
| 162 | content_pack: Optional[str] = Field( |
| 163 | None, |
| 164 | description="Associated content pack, if any", |
| 165 | ) |
| 166 | model_config = ConfigDict( |
| 167 | json_schema_extra={ |
| 168 | "example": { |
| 169 | "title": "Office365 EVENTS - Example Company", |
| 170 | "description": "Office365 EVENTS - Example Company", |
| 171 | "index_set_id": "12345", |
| 172 | "rules": [ |
| 173 | { |
| 174 | "field": "agent_labels_customer", |
| 175 | "type": 1, |
| 176 | "inverted": False, |
| 177 | "value": "ExampleCode", |
| 178 | }, |
| 179 | { |
| 180 | "field": "agent_labels_integration", |
| 181 | "type": 1, |
| 182 | "inverted": False, |
| 183 | "value": "Office365", |
| 184 | }, |
| 185 | ], |
| 186 | "matching_type": "AND", |
| 187 | "remove_matches_from_default_stream": True, |
| 188 | "content_pack": None, |
| 189 | }, |
| 190 | }, |
| 191 | ) |
| 192 | |
| 193 | |
| 194 | class StreamData(BaseModel): |
| 195 | stream_id: str = Field(..., description="ID of the created stream") |
| 196 | |
| 197 | |
| 198 | class StreamCreationResponse(BaseModel): |
| 199 | data: StreamData |
| 200 | success: bool = Field(..., description="Indicates if the request was successful") |
| 201 | message: str = Field( |
| 202 | ..., |
| 203 | description="A message detailing the outcome of the request", |
| 204 | ) |
| 205 | |
| 206 | |
| 207 | class StreamAndPipelineData(BaseModel): |
| 208 | stream_id: str = Field(..., description="ID of the stream") |
| 209 | pipeline_ids: List[str] = Field( |
| 210 | ..., |
| 211 | description="List of pipeline IDs connected to the stream", |
| 212 | ) |
| 213 | |
| 214 | |
| 215 | class StreamConnectionToPipelineRequest(BaseModel): |
| 216 | stream_id: str = Field(..., description="ID of the stream to connect") |
| 217 | pipeline_ids: List[str] = Field( |
| 218 | ..., |
| 219 | description="List of pipeline IDs to connect to the stream", |
| 220 | ) |
| 221 | |
| 222 | |
| 223 | class StreamConnectionToPipelineResponse(BaseModel): |
| 224 | data: StreamAndPipelineData |
| 225 | success: bool = Field(..., description="Indicates if the request was successful") |
| 226 | message: str = Field( |
| 227 | ..., |
| 228 | description="A message detailing the outcome of the request", |
| 229 | ) |