| 1 | import json |
| 2 | from datetime import datetime |
| 3 | |
| 4 | from fastapi import HTTPException |
| 5 | from loguru import logger |
| 6 | |
| 7 | from app.connectors.graylog.services.pipelines import get_pipelines |
| 8 | from app.connectors.graylog.utils.universal import send_delete_request |
| 9 | from app.connectors.graylog.utils.universal import send_get_request |
| 10 | from app.connectors.graylog.utils.universal import send_post_request |
| 11 | from app.connectors.graylog.utils.universal import send_post_request_create_entity |
| 12 | from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse |
| 13 | from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest |
| 14 | from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineResponse |
| 15 | from app.customer_provisioning.schema.graylog import StreamCreationResponse |
| 16 | from app.customer_provisioning.schema.graylog import TimeBasedIndexSet |
| 17 | from app.customer_provisioning.schema.graylog import WazuhEventStream |
| 18 | from app.customer_provisioning.schema.provision import ProvisionNewCustomer |
| 19 | |
| 20 | |
| 21 | ######### ! GRAYLOG PROVISIONING ! ############ |
| 22 | # ! INDEX SETS ! # |
| 23 | def build_index_set_config(request: ProvisionNewCustomer) -> TimeBasedIndexSet: |
| 24 | """ |
| 25 | Build the configuration for a time-based index set. |
| 26 | |
| 27 | Args: |
| 28 | request (ProvisionNewCustomer): The request object containing customer information. |
| 29 | |
| 30 | Returns: |
| 31 | TimeBasedIndexSet: The configured time-based index set. |
| 32 | """ |
| 33 | return TimeBasedIndexSet( |
| 34 | title=f"{request.customer_name} - Wazuh EDR EVENTS", |
| 35 | description=f"{request.customer_name} - Wazuh EDR EVENTS", |
| 36 | index_prefix=request.customer_index_name, |
| 37 | rotation_strategy_class="org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategy", |
| 38 | rotation_strategy={ |
| 39 | "type": "org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategyConfig", |
| 40 | "rotation_period": "P1D", |
| 41 | "rotate_empty_index_set": False, |
| 42 | "max_rotation_period": None, |
| 43 | }, |
| 44 | retention_strategy_class="org.graylog2.indexer.retention.strategies.DeletionRetentionStrategy", |
| 45 | retention_strategy={ |
| 46 | "type": "org.graylog2.indexer.retention.strategies.DeletionRetentionStrategyConfig", |
| 47 | "max_number_of_indices": request.hot_data_retention, |
| 48 | }, |
| 49 | creation_date=datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"), |
| 50 | index_analyzer="standard", |
| 51 | shards=request.index_shards, |
| 52 | replicas=request.index_replicas, |
| 53 | index_optimization_max_num_segments=1, |
| 54 | index_optimization_disabled=False, |
| 55 | writable=True, |
| 56 | field_type_refresh_interval=5000, |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | # Function to send the POST request and handle the response |
| 61 | async def send_index_set_creation_request( |
| 62 | index_set: TimeBasedIndexSet, |
| 63 | ) -> GraylogIndexSetCreationResponse: |
| 64 | """ |
| 65 | Sends a request to create an index set in Graylog. |
| 66 | |
| 67 | Args: |
| 68 | index_set (TimeBasedIndexSet): The index set to be created. |
| 69 | |
| 70 | Returns: |
| 71 | GraylogIndexSetCreationResponse: The response from Graylog after creating the index set. |
| 72 | """ |
| 73 | json_index_set = json.dumps(index_set.model_dump()) |
| 74 | logger.info(f"json_index_set set: {json_index_set}") |
| 75 | response_json = await send_post_request( |
| 76 | endpoint="/api/system/indices/index_sets", |
| 77 | data=index_set.model_dump(), |
| 78 | ) |
| 79 | return GraylogIndexSetCreationResponse(**response_json) |
| 80 | |
| 81 | |
| 82 | # Refactored create_index_set function |
| 83 | async def create_index_set( |
| 84 | request: ProvisionNewCustomer, |
| 85 | ) -> GraylogIndexSetCreationResponse: |
| 86 | """ |
| 87 | Creates an index set for a new customer. |
| 88 | |
| 89 | Args: |
| 90 | request (ProvisionNewCustomer): The request object containing the customer information. |
| 91 | |
| 92 | Returns: |
| 93 | GraylogIndexSetCreationResponse: The response object containing the result of the index set creation. |
| 94 | """ |
| 95 | logger.info(f"Creating index set for customer {request.customer_name}") |
| 96 | index_set_config = build_index_set_config(request) |
| 97 | return await send_index_set_creation_request(index_set_config) |
| 98 | |
| 99 | |
| 100 | # Function to extract index set ID |
| 101 | def extract_index_set_id(response: GraylogIndexSetCreationResponse) -> str: |
| 102 | """ |
| 103 | Extracts the index set ID from the given GraylogIndexSetCreationResponse object. |
| 104 | |
| 105 | Args: |
| 106 | response (GraylogIndexSetCreationResponse): The GraylogIndexSetCreationResponse object. |
| 107 | |
| 108 | Returns: |
| 109 | str: The index set ID extracted from the response. |
| 110 | """ |
| 111 | return response.data.id |
| 112 | |
| 113 | |
| 114 | # ! Event STREAMS ! # |
| 115 | # Function to create event stream configuration |
| 116 | def build_event_stream_config( |
| 117 | request: ProvisionNewCustomer, |
| 118 | index_set_id: str, |
| 119 | ) -> WazuhEventStream: |
| 120 | """ |
| 121 | Build the configuration for a Wazuh event stream. |
| 122 | |
| 123 | Args: |
| 124 | request (ProvisionNewCustomer): The request object containing customer information. |
| 125 | index_set_id (str): The ID of the index set. |
| 126 | |
| 127 | Returns: |
| 128 | WazuhEventStream: The configured Wazuh event stream. |
| 129 | """ |
| 130 | sanitized_name = request.customer_code.replace(" ", "_") |
| 131 | return WazuhEventStream( |
| 132 | title=f"{request.customer_name} - Wazuh EDR EVENTS", |
| 133 | description=f"{request.customer_name} - Wazuh EDR EVENTS", |
| 134 | index_set_id=index_set_id, |
| 135 | rules=[ |
| 136 | { |
| 137 | "field": "agent_labels_customer", |
| 138 | "type": 1, |
| 139 | "inverted": False, |
| 140 | "value": request.customer_code, |
| 141 | }, |
| 142 | {"field": "cluster_node", "type": 1, "inverted": False, "value": f"wazuh.worker.{sanitized_name}.1"}, |
| 143 | ], |
| 144 | matching_type="OR", |
| 145 | remove_matches_from_default_stream=True, |
| 146 | content_pack=None, |
| 147 | ) |
| 148 | |
| 149 | |
| 150 | async def send_event_stream_creation_request( |
| 151 | event_stream: WazuhEventStream, |
| 152 | ) -> StreamCreationResponse: |
| 153 | """ |
| 154 | Sends a request to create an event stream. |
| 155 | |
| 156 | Args: |
| 157 | event_stream (WazuhEventStream): The event stream to be created. |
| 158 | |
| 159 | Returns: |
| 160 | StreamCreationResponse: The response containing the created event stream. |
| 161 | """ |
| 162 | json_event_stream = json.dumps(event_stream.model_dump()) |
| 163 | logger.info(f"json_event_stream set: {json_event_stream}") |
| 164 | response_json = await send_post_request_create_entity( |
| 165 | endpoint="/api/streams", |
| 166 | entity=event_stream.model_dump(), |
| 167 | ) |
| 168 | return StreamCreationResponse(**response_json) |
| 169 | |
| 170 | |
| 171 | async def create_event_stream(request: ProvisionNewCustomer, index_set_id: str): |
| 172 | """ |
| 173 | Creates an event stream for a customer. |
| 174 | |
| 175 | Args: |
| 176 | request (ProvisionNewCustomer): The request object containing customer information. |
| 177 | index_set_id (str): The ID of the index set. |
| 178 | |
| 179 | Returns: |
| 180 | The result of the event stream creation request. |
| 181 | """ |
| 182 | logger.info(f"Creating event stream for customer {request.customer_name}") |
| 183 | event_stream_config = build_event_stream_config(request, index_set_id) |
| 184 | return await send_event_stream_creation_request(event_stream_config) |
| 185 | |
| 186 | |
| 187 | # ! PIPELINES ! # |
| 188 | # Function to get pipeline ID |
| 189 | async def get_pipeline_id(subscription: str) -> str: |
| 190 | """ |
| 191 | Retrieves the pipeline ID for a given subscription. |
| 192 | |
| 193 | Args: |
| 194 | subscription (str): The subscription name. |
| 195 | |
| 196 | Returns: |
| 197 | str: The pipeline ID. |
| 198 | |
| 199 | Raises: |
| 200 | HTTPException: If the pipeline ID cannot be retrieved. |
| 201 | """ |
| 202 | logger.info(f"Getting pipeline ID for subscription {subscription}") |
| 203 | pipelines_response = await get_pipelines() |
| 204 | if pipelines_response.success: |
| 205 | for pipeline in pipelines_response.pipelines: |
| 206 | if subscription.lower() in pipeline.description.lower(): |
| 207 | return [pipeline.id] |
| 208 | logger.error(f"Failed to get pipeline ID for subscription {subscription}") |
| 209 | raise HTTPException( |
| 210 | status_code=500, |
| 211 | detail=( |
| 212 | f"Failed to get pipeline ID for subscription {subscription}. " |
| 213 | "Please ensure you have installed the SOCFortress Wazuh Content Pack. " |
| 214 | "See more at: https://youtu.be/euFrHP0VkD8?si=ajqjNobHvBjrTzAH" |
| 215 | ), |
| 216 | ) |
| 217 | else: |
| 218 | logger.error(f"Failed to get pipelines: {pipelines_response.message}") |
| 219 | raise HTTPException( |
| 220 | status_code=500, |
| 221 | detail=f"Failed to get pipelines: {pipelines_response.message}", |
| 222 | ) |
| 223 | |
| 224 | |
| 225 | async def connect_stream_to_pipeline( |
| 226 | stream_and_pipeline: StreamConnectionToPipelineRequest, |
| 227 | ): |
| 228 | """ |
| 229 | Connects a stream to a pipeline. |
| 230 | |
| 231 | Args: |
| 232 | stream_and_pipeline (StreamConnectionToPipelineRequest): The request object containing the stream ID and pipeline IDs. |
| 233 | |
| 234 | Returns: |
| 235 | StreamConnectionToPipelineResponse: The response object containing the connection details. |
| 236 | """ |
| 237 | logger.info( |
| 238 | f"Connecting stream {stream_and_pipeline.stream_id} to pipeline {stream_and_pipeline.pipeline_ids}", |
| 239 | ) |
| 240 | response_json = await send_post_request( |
| 241 | endpoint="/api/system/pipelines/connections/to_stream", |
| 242 | data=stream_and_pipeline.model_dump(), |
| 243 | ) |
| 244 | logger.info(f"Response: {response_json}") |
| 245 | return StreamConnectionToPipelineResponse(**response_json) |
| 246 | |
| 247 | |
| 248 | ######### ! GRAYLOG DECOMISSIONGING ! ############ |
| 249 | async def delete_stream(stream_id: str): |
| 250 | """ |
| 251 | Deletes a stream. |
| 252 | |
| 253 | Args: |
| 254 | stream_id (str): The ID of the stream to be deleted. |
| 255 | |
| 256 | Returns: |
| 257 | The result of the stream deletion request. |
| 258 | """ |
| 259 | logger.info(f"Deleting stream {stream_id}") |
| 260 | response = await send_delete_request(endpoint=f"/api/streams/{stream_id}") |
| 261 | return response |
| 262 | |
| 263 | |
| 264 | async def delete_index_set(index_set_id: str): |
| 265 | """ |
| 266 | Deletes an index set. |
| 267 | |
| 268 | Args: |
| 269 | index_set_id (str): The ID of the index set to be deleted. |
| 270 | |
| 271 | Returns: |
| 272 | The result of the index set deletion request. |
| 273 | """ |
| 274 | logger.info(f"Deleting index set {index_set_id}") |
| 275 | response = await send_delete_request( |
| 276 | endpoint=f"/api/system/indices/index_sets/{index_set_id}", |
| 277 | ) |
| 278 | return response |
| 279 | |
| 280 | |
| 281 | async def get_content_pack_installation_id(content_pack_id: str): |
| 282 | """ |
| 283 | Retrieves the installation ID of a content pack. |
| 284 | |
| 285 | Args: |
| 286 | content_pack_id (str): The ID of the content pack. |
| 287 | |
| 288 | Returns: |
| 289 | str: The installation ID of the content pack. |
| 290 | """ |
| 291 | logger.info(f"Getting installation ID for content pack {content_pack_id}") |
| 292 | response = await send_get_request( |
| 293 | endpoint=f"/api/system/content_packs/{content_pack_id}/installations", |
| 294 | ) |
| 295 | return response["data"]["installations"][0]["_id"] |
| 296 | |
| 297 | |
| 298 | async def uninstall_content_pack(content_pack_id: str): |
| 299 | """ |
| 300 | Uninstalls a content pack. |
| 301 | |
| 302 | Args: |
| 303 | content_pack_id (str): The ID of the content pack to be deleted. |
| 304 | |
| 305 | Returns: |
| 306 | The result of the content pack deletion request. |
| 307 | """ |
| 308 | logger.info(f"Deleting content pack {content_pack_id}") |
| 309 | installation_id = await get_content_pack_installation_id(content_pack_id) |
| 310 | response = await send_delete_request( |
| 311 | endpoint=f"/api/system/content_packs/{content_pack_id}/installations/{installation_id}", |
| 312 | ) |
| 313 | return response |
| 314 | |
| 315 | |
| 316 | async def delete_content_pack(content_pack_id: str): |
| 317 | """ |
| 318 | Deletes a content pack. |
| 319 | |
| 320 | Args: |
| 321 | content_pack_id (str): The ID of the content pack to be deleted. |
| 322 | |
| 323 | Returns: |
| 324 | The result of the content pack deletion request. |
| 325 | """ |
| 326 | logger.info(f"Deleting content pack {content_pack_id}") |
| 327 | response = await send_delete_request( |
| 328 | endpoint=f"/api/system/content_packs/{content_pack_id}", |
| 329 | ) |
| 330 | return response |