| 1 | import json |
| 2 | import os |
| 3 | import xml.etree.ElementTree as ET |
| 4 | from datetime import datetime |
| 5 | from typing import List |
| 6 | from xml.dom import minidom |
| 7 | from xml.dom.minidom import parseString |
| 8 | from xml.etree.ElementTree import SubElement |
| 9 | from xml.etree.ElementTree import parse |
| 10 | from xml.etree.ElementTree import tostring |
| 11 | |
| 12 | import aiofiles |
| 13 | from dotenv import load_dotenv |
| 14 | from fastapi import HTTPException |
| 15 | from loguru import logger |
| 16 | from sqlalchemy import and_ |
| 17 | from sqlalchemy import update |
| 18 | from sqlalchemy.ext.asyncio import AsyncSession |
| 19 | |
| 20 | from app.connectors.grafana.schema.dashboards import DashboardProvisionRequest |
| 21 | from app.connectors.grafana.schema.dashboards import Office365Dashboard |
| 22 | from app.connectors.grafana.services.dashboards import provision_dashboards |
| 23 | from app.connectors.grafana.utils.universal import create_grafana_client |
| 24 | from app.connectors.graylog.schema.pipelines import CreatePipeline |
| 25 | from app.connectors.graylog.schema.pipelines import CreatePipelineRule |
| 26 | from app.connectors.graylog.schema.pipelines import GraylogPipelinesResponse |
| 27 | from app.connectors.graylog.schema.pipelines import PipelineRulesResponse |
| 28 | from app.connectors.graylog.services.management import start_stream |
| 29 | from app.connectors.graylog.services.pipelines import connect_stream_to_pipeline |
| 30 | from app.connectors.graylog.services.pipelines import create_pipeline_graylog |
| 31 | from app.connectors.graylog.services.pipelines import create_pipeline_rule |
| 32 | from app.connectors.graylog.services.pipelines import get_pipeline_id |
| 33 | from app.connectors.graylog.services.pipelines import get_pipeline_rules |
| 34 | from app.connectors.graylog.services.pipelines import get_pipelines |
| 35 | from app.connectors.graylog.utils.universal import send_post_request |
| 36 | from app.connectors.graylog.utils.universal import send_post_request_create_entity |
| 37 | from app.connectors.wazuh_indexer.services.monitoring import ( |
| 38 | output_shard_number_to_be_set_based_on_nodes, |
| 39 | ) |
| 40 | from app.connectors.wazuh_manager.utils.universal import send_get_request |
| 41 | from app.connectors.wazuh_manager.utils.universal import send_put_request |
| 42 | from app.customer_provisioning.schema.grafana import GrafanaDatasource |
| 43 | from app.customer_provisioning.schema.grafana import GrafanaDataSourceCreationResponse |
| 44 | from app.customer_provisioning.schema.graylog import GraylogIndexSetCreationResponse |
| 45 | from app.customer_provisioning.schema.graylog import Office365EventStream |
| 46 | from app.customer_provisioning.schema.graylog import StreamConnectionToPipelineRequest |
| 47 | from app.customer_provisioning.schema.graylog import StreamCreationResponse |
| 48 | from app.customer_provisioning.schema.graylog import TimeBasedIndexSet |
| 49 | from app.customer_provisioning.services.grafana import create_grafana_folder |
| 50 | from app.customer_provisioning.services.grafana import get_opensearch_version |
| 51 | from app.customers.routes.customers import get_customer |
| 52 | from app.customers.routes.customers import get_customer_meta |
| 53 | from app.db.universal_models import CustomersMeta |
| 54 | from app.integrations.models.customer_integration_settings import CustomerIntegrations |
| 55 | from app.integrations.office365.schema.provision import PipelineRuleTitles |
| 56 | from app.integrations.office365.schema.provision import PipelineTitles |
| 57 | from app.integrations.office365.schema.provision import ProvisionOffice365AuthKeys |
| 58 | from app.integrations.office365.schema.provision import ProvisionOffice365Response |
| 59 | from app.integrations.routes import create_integration_meta |
| 60 | from app.integrations.schema import CustomerIntegrationsMetaSchema |
| 61 | from app.utils import get_connector_attribute |
| 62 | from app.utils import get_customer_default_settings_attribute |
| 63 | |
| 64 | load_dotenv() |
| 65 | |
| 66 | |
| 67 | ############ ! WAZUH MANAGER ! ############ |
| 68 | async def get_wazuh_configuration(file_name: str) -> str: |
| 69 | """ |
| 70 | Retrieves the Wazuh configuration from the manager and writes it to a file. |
| 71 | |
| 72 | Args: |
| 73 | file_name (str): The name of the file where the configuration data will be written. |
| 74 | |
| 75 | Returns: |
| 76 | str: The Wazuh configuration data. |
| 77 | """ |
| 78 | endpoint = "/manager/configuration" |
| 79 | params = {"raw": True} |
| 80 | response = await send_get_request(endpoint=endpoint, params=params) |
| 81 | config_data = response["data"] |
| 82 | |
| 83 | # Get the directory of the current module |
| 84 | dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 85 | # Create the full file path |
| 86 | file_path = os.path.join(dir_path, file_name) |
| 87 | |
| 88 | async with aiofiles.open(file_path, "w") as f: |
| 89 | await f.write(config_data) |
| 90 | |
| 91 | return config_data |
| 92 | |
| 93 | |
| 94 | async def office365_template_with_api_type( |
| 95 | customer_code: str, |
| 96 | provision_office365_auth_keys: ProvisionOffice365AuthKeys, |
| 97 | ) -> str: |
| 98 | """ |
| 99 | Returns a configured Office365 template for Wazuh. |
| 100 | |
| 101 | Args: |
| 102 | customer_code (str): The customer code. |
| 103 | provision_office365_auth_keys (ProvisionOffice365AuthKeys): The Office365 auth keys. |
| 104 | |
| 105 | Returns: |
| 106 | str: The Office365 template configured with the given parameters. |
| 107 | """ |
| 108 | |
| 109 | # Get the directory of the current module |
| 110 | dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 111 | # Create the full file path |
| 112 | wazuh_config = os.path.join(dir_path, "wazuh_config.xml") |
| 113 | |
| 114 | # Parse the existing XML file |
| 115 | tree = parse(wazuh_config) |
| 116 | root = tree.getroot() |
| 117 | |
| 118 | # Create the office365 element and add it to ossec_config |
| 119 | office365 = SubElement(root, "office365") |
| 120 | |
| 121 | # Add the child elements to office365 |
| 122 | SubElement(office365, "enabled").text = "yes" |
| 123 | SubElement(office365, "interval").text = "1m" |
| 124 | SubElement(office365, "curl_max_size").text = "5M" |
| 125 | SubElement(office365, "only_future_events").text = "yes" |
| 126 | |
| 127 | # Create the api_auth element and add it to office365 |
| 128 | api_auth = SubElement(office365, "api_auth") |
| 129 | |
| 130 | # Add the customer code as a comment within the api_auth block |
| 131 | # api_auth.insert(0, Comment(f' Customer Code: {customer_code} ')) |
| 132 | |
| 133 | # Add the child elements to api_auth |
| 134 | SubElement(api_auth, "tenant_id").text = provision_office365_auth_keys.TENANT_ID |
| 135 | SubElement(api_auth, "client_id").text = provision_office365_auth_keys.CLIENT_ID |
| 136 | SubElement(api_auth, "client_secret").text = provision_office365_auth_keys.CLIENT_SECRET |
| 137 | SubElement(api_auth, "api_type").text = provision_office365_auth_keys.API_TYPE |
| 138 | |
| 139 | # Create the subscriptions element and add it to office365 |
| 140 | subscriptions = SubElement(office365, "subscriptions") |
| 141 | |
| 142 | # Add the child elements to subscriptions |
| 143 | SubElement(subscriptions, "subscription").text = "Audit.SharePoint" |
| 144 | SubElement(subscriptions, "subscription").text = "Audit.Exchange" |
| 145 | SubElement(subscriptions, "subscription").text = "DLP.ALL" |
| 146 | SubElement(subscriptions, "subscription").text = "Audit.General" |
| 147 | SubElement(subscriptions, "subscription").text = "Audit.AzureActiveDirectory" |
| 148 | |
| 149 | # Convert the office365 element to a string |
| 150 | office365_str = tostring(office365).decode("utf-8") |
| 151 | |
| 152 | # Pretty print the office365 element |
| 153 | dom = parseString(office365_str) |
| 154 | pretty_office365_str = dom.toprettyxml(indent=" ") |
| 155 | |
| 156 | # Remove the XML declaration from the pretty printed office365 string |
| 157 | pretty_office365_str = pretty_office365_str.replace('<?xml version="1.0" ?>', "").strip() |
| 158 | |
| 159 | # Convert the entire XML to a string |
| 160 | xml_str = tostring(root).decode("utf-8") |
| 161 | |
| 162 | # Replace the original office365 string with the pretty printed office365 string |
| 163 | xml_str = xml_str.replace(office365_str, pretty_office365_str) |
| 164 | |
| 165 | # Overwrite the existing XML file with the new contents |
| 166 | async with aiofiles.open(wazuh_config, "w") as f: |
| 167 | await f.write(xml_str) |
| 168 | |
| 169 | return xml_str |
| 170 | |
| 171 | |
| 172 | async def append_office365_template(wazuh_config: str, office365_template: str) -> str: |
| 173 | """ |
| 174 | Appends the Office365 template to the Wazuh configuration. |
| 175 | |
| 176 | Args: |
| 177 | wazuh_config (str): The current Wazuh configuration. |
| 178 | office365_template (str): The Office365 template to append. |
| 179 | |
| 180 | Returns: |
| 181 | str: The Wazuh configuration with the Office365 template appended. |
| 182 | """ |
| 183 | |
| 184 | return wazuh_config + office365_template |
| 185 | |
| 186 | |
| 187 | async def add_api_auth_to_office365_block(customer_code: str, provision_office365_auth_keys: ProvisionOffice365AuthKeys) -> str: |
| 188 | try: |
| 189 | # Get the directory of the current module |
| 190 | dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 191 | # Create the full file path |
| 192 | wazuh_config = os.path.join(dir_path, "wazuh_config.xml") |
| 193 | |
| 194 | # Parse the existing XML file |
| 195 | tree = ET.ElementTree() |
| 196 | tree.parse(wazuh_config) |
| 197 | root = tree.getroot() |
| 198 | |
| 199 | # Find the office365 block |
| 200 | office365_block = root.find("office365") |
| 201 | |
| 202 | # If the office365 block exists |
| 203 | if office365_block is not None: |
| 204 | # Find the index of the subscriptions block |
| 205 | subscriptions_index = list(office365_block).index(office365_block.find("subscriptions")) |
| 206 | |
| 207 | # Create a new api_auth block |
| 208 | api_auth_block = ET.Element("api_auth") |
| 209 | |
| 210 | # Add the tenant_id, client_id, client_secret, and api_type to the api_auth block |
| 211 | ET.SubElement(api_auth_block, "tenant_id").text = provision_office365_auth_keys.TENANT_ID |
| 212 | ET.SubElement(api_auth_block, "client_id").text = provision_office365_auth_keys.CLIENT_ID |
| 213 | ET.SubElement(api_auth_block, "client_secret").text = provision_office365_auth_keys.CLIENT_SECRET |
| 214 | ET.SubElement(api_auth_block, "api_type").text = provision_office365_auth_keys.API_TYPE |
| 215 | |
| 216 | # Pretty print the new api_auth block |
| 217 | pretty_api_auth_block = minidom.parseString(ET.tostring(api_auth_block)).toprettyxml(indent=" ") |
| 218 | |
| 219 | # Parse the pretty printed api_auth block back to an Element |
| 220 | pretty_api_auth_element = ET.fromstring(pretty_api_auth_block) |
| 221 | |
| 222 | # Insert the new pretty printed api_auth block above the subscriptions block |
| 223 | office365_block.insert(subscriptions_index, pretty_api_auth_element) |
| 224 | |
| 225 | # Convert the modified configuration back to string format |
| 226 | modified_config = ET.tostring(root, encoding="utf-8").decode("utf-8") |
| 227 | |
| 228 | # Overwrite the existing XML file with the new contents |
| 229 | with open(wazuh_config, "w") as f: |
| 230 | f.write(modified_config) |
| 231 | |
| 232 | return modified_config |
| 233 | |
| 234 | except Exception as e: |
| 235 | logger.error(f"An error occurred: {e}") |
| 236 | raise HTTPException( |
| 237 | status_code=500, |
| 238 | detail="Error found in ossec.conf. Multiple <ossec_config> blocks found. Remove all additional <ossec_config> blocks and try again.", |
| 239 | ) |
| 240 | |
| 241 | |
| 242 | async def update_wazuh_configuration( |
| 243 | provision_office365_auth_keys: ProvisionOffice365AuthKeys, |
| 244 | ) -> None: |
| 245 | """ |
| 246 | Updates the Wazuh configuration. If it fails, remove the <api_type> tag and retry. |
| 247 | |
| 248 | Args: |
| 249 | provision_office365_auth_keys (ProvisionOffice365AuthKeys): The Office365 authentication keys. |
| 250 | """ |
| 251 | endpoint = "/manager/configuration" |
| 252 | |
| 253 | # Get the directory of the current module |
| 254 | dir_path = os.path.dirname(os.path.realpath(__file__)) |
| 255 | # Create the full file path |
| 256 | wazuh_config_path = os.path.join(dir_path, "wazuh_config.xml") |
| 257 | |
| 258 | # Read the Wazuh configuration from the file |
| 259 | async with aiofiles.open(wazuh_config_path, "r") as f: |
| 260 | wazuh_config = await f.read() |
| 261 | |
| 262 | data = wazuh_config.encode("utf-8") |
| 263 | |
| 264 | try: |
| 265 | # First attempt to update configuration |
| 266 | response = await send_put_request( |
| 267 | endpoint=endpoint, |
| 268 | data=data, |
| 269 | binary_data=True, |
| 270 | ) |
| 271 | if response.get("success") and response["data"].get("error") == 0: |
| 272 | logger.info("Wazuh configuration updated successfully.") |
| 273 | return |
| 274 | else: |
| 275 | logger.error( |
| 276 | "Failed to update Wazuh configuration. Error: {}".format(response), |
| 277 | ) |
| 278 | |
| 279 | except Exception as e: |
| 280 | logger.error(f"Exception occurred during Wazuh configuration update: {e}") |
| 281 | |
| 282 | # Remove <api_type> tag and retry |
| 283 | api_type_tag = f"<api_type>{provision_office365_auth_keys.API_TYPE}</api_type>" |
| 284 | modified_wazuh_config = wazuh_config.replace(api_type_tag, "") |
| 285 | data = modified_wazuh_config.encode("utf-8") |
| 286 | |
| 287 | try: |
| 288 | response = await send_put_request( |
| 289 | endpoint=endpoint, |
| 290 | data=data, |
| 291 | binary_data=True, |
| 292 | ) |
| 293 | if response.get("success") and response["data"].get("error") == 0: |
| 294 | logger.info( |
| 295 | "Wazuh configuration updated successfully after removing <api_type> tag.", |
| 296 | ) |
| 297 | else: |
| 298 | logger.error( |
| 299 | "Failed to update Wazuh configuration after removing <api_type> tag. Error: {}".format( |
| 300 | response, |
| 301 | ), |
| 302 | ) |
| 303 | |
| 304 | except Exception as e: |
| 305 | logger.error( |
| 306 | f"Exception occurred during retry of Wazuh configuration update: {e}", |
| 307 | ) |
| 308 | raise HTTPException( |
| 309 | status_code=500, |
| 310 | detail="Failed to update Wazuh configuration.", |
| 311 | ) |
| 312 | |
| 313 | |
| 314 | async def check_if_office365_is_already_provisioned( |
| 315 | customer_code: str, |
| 316 | wazuh_config: str, |
| 317 | ) -> bool: |
| 318 | """ |
| 319 | If the string "Office365 Integration" is found in the Wazuh configuration, return True. |
| 320 | |
| 321 | Args: |
| 322 | customer_code (str): The customer code. |
| 323 | wazuh_config (str): The Wazuh configuration in string format. |
| 324 | |
| 325 | Returns: |
| 326 | bool: True if the Office365 integration is already provisioned, False otherwise. |
| 327 | """ |
| 328 | if "office365" in wazuh_config: |
| 329 | return True |
| 330 | return False |
| 331 | |
| 332 | |
| 333 | async def check_if_office365_is_already_provisioned_for_customer( |
| 334 | tenant_id: str, |
| 335 | wazuh_config: str, |
| 336 | ) -> bool: |
| 337 | """ |
| 338 | If the string "Office365 Integration" is found in the Wazuh configuration, return True. |
| 339 | |
| 340 | Args: |
| 341 | customer_code (str): The customer code. |
| 342 | wazuh_config (str): The Wazuh configuration in string format. |
| 343 | |
| 344 | Returns: |
| 345 | bool: True if the Office365 integration is already provisioned, False otherwise. |
| 346 | """ |
| 347 | if f"{tenant_id}" in wazuh_config: |
| 348 | raise HTTPException( |
| 349 | status_code=400, |
| 350 | detail=f"Office365 integration is already provisioned for customer {tenant_id}.", |
| 351 | ) |
| 352 | |
| 353 | |
| 354 | async def restart_wazuh_manager() -> None: |
| 355 | """ |
| 356 | Restarts the Wazuh manager service. |
| 357 | """ |
| 358 | logger.info("Restarting Wazuh manager service.") |
| 359 | await send_put_request(endpoint="/manager/restart", data=None) |
| 360 | |
| 361 | |
| 362 | ################## ! GRAYLOG ! ################## |
| 363 | |
| 364 | |
| 365 | async def build_index_set_config( |
| 366 | customer_code: str, |
| 367 | session: AsyncSession, |
| 368 | ) -> TimeBasedIndexSet: |
| 369 | """ |
| 370 | Build the configuration for a time-based index set. |
| 371 | |
| 372 | Args: |
| 373 | request (ProvisionNewCustomer): The request object containing customer information. |
| 374 | |
| 375 | Returns: |
| 376 | TimeBasedIndexSet: The configured time-based index set. |
| 377 | """ |
| 378 | # Lowercase the customer code since Graylog index sets must be lowercase |
| 379 | customer_code = customer_code.lower() |
| 380 | return TimeBasedIndexSet( |
| 381 | title=f"{(await get_customer(customer_code, session)).customer.customer_name} - Office365", |
| 382 | description=f"{customer_code} - Office365", |
| 383 | index_prefix=f"office365-{customer_code}", |
| 384 | rotation_strategy_class="org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategy", |
| 385 | rotation_strategy={ |
| 386 | "type": "org.graylog2.indexer.rotation.strategies.TimeBasedRotationStrategyConfig", |
| 387 | "rotation_period": "P1D", |
| 388 | "rotate_empty_index_set": False, |
| 389 | "max_rotation_period": None, |
| 390 | }, |
| 391 | retention_strategy_class="org.graylog2.indexer.retention.strategies.DeletionRetentionStrategy", |
| 392 | retention_strategy={ |
| 393 | "type": "org.graylog2.indexer.retention.strategies.DeletionRetentionStrategyConfig", |
| 394 | "max_number_of_indices": 30, |
| 395 | }, |
| 396 | creation_date=datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"), |
| 397 | index_analyzer="standard", |
| 398 | shards=await output_shard_number_to_be_set_based_on_nodes(), |
| 399 | replicas=0, |
| 400 | index_optimization_max_num_segments=1, |
| 401 | index_optimization_disabled=False, |
| 402 | writable=True, |
| 403 | field_type_refresh_interval=5000, |
| 404 | ) |
| 405 | |
| 406 | |
| 407 | # Function to send the POST request and handle the response |
| 408 | async def send_index_set_creation_request( |
| 409 | index_set: TimeBasedIndexSet, |
| 410 | ) -> GraylogIndexSetCreationResponse: |
| 411 | """ |
| 412 | Sends a request to create an index set in Graylog. |
| 413 | |
| 414 | Args: |
| 415 | index_set (TimeBasedIndexSet): The index set to be created. |
| 416 | |
| 417 | Returns: |
| 418 | GraylogIndexSetCreationResponse: The response from Graylog after creating the index set. |
| 419 | """ |
| 420 | json_index_set = json.dumps(index_set.model_dump()) |
| 421 | logger.info(f"json_index_set set: {json_index_set}") |
| 422 | response_json = await send_post_request( |
| 423 | endpoint="/api/system/indices/index_sets", |
| 424 | data=index_set.model_dump(), |
| 425 | ) |
| 426 | return GraylogIndexSetCreationResponse(**response_json) |
| 427 | |
| 428 | |
| 429 | # Refactored create_index_set function |
| 430 | async def create_index_set( |
| 431 | customer_code: str, |
| 432 | session: AsyncSession, |
| 433 | ) -> GraylogIndexSetCreationResponse: |
| 434 | """ |
| 435 | Creates an index set for a new customer. |
| 436 | |
| 437 | Args: |
| 438 | request (ProvisionNewCustomer): The request object containing the customer information. |
| 439 | |
| 440 | Returns: |
| 441 | GraylogIndexSetCreationResponse: The response object containing the result of the index set creation. |
| 442 | """ |
| 443 | logger.info(f"Creating index set for customer {customer_code}") |
| 444 | index_set_config = await build_index_set_config(customer_code, session) |
| 445 | return await send_index_set_creation_request(index_set_config) |
| 446 | |
| 447 | |
| 448 | # Function to extract index set ID |
| 449 | def extract_index_set_id(response: GraylogIndexSetCreationResponse) -> str: |
| 450 | """ |
| 451 | Extracts the index set ID from the given GraylogIndexSetCreationResponse object. |
| 452 | |
| 453 | Args: |
| 454 | response (GraylogIndexSetCreationResponse): The GraylogIndexSetCreationResponse object. |
| 455 | |
| 456 | Returns: |
| 457 | str: The index set ID extracted from the response. |
| 458 | """ |
| 459 | return response.data.id |
| 460 | |
| 461 | |
| 462 | # ! Event STREAMS ! # |
| 463 | # Function to create event stream configuration |
| 464 | async def build_event_stream_config( |
| 465 | customer_code: str, |
| 466 | provision_office365_auth_keys: ProvisionOffice365AuthKeys, |
| 467 | index_set_id: str, |
| 468 | session: AsyncSession, |
| 469 | ) -> Office365EventStream: |
| 470 | """ |
| 471 | Build the configuration for a Wazuh event stream. |
| 472 | |
| 473 | Args: |
| 474 | request (ProvisionNewCustomer): The request object containing customer information. |
| 475 | index_set_id (str): The ID of the index set. |
| 476 | |
| 477 | Returns: |
| 478 | Office365EventStream: The configured Wazuh event stream. |
| 479 | """ |
| 480 | return Office365EventStream( |
| 481 | title=f"{(await get_customer(customer_code, session)).customer.customer_name} - Office365", |
| 482 | description=f"{(await get_customer(customer_code, session)).customer.customer_name} - Office365", |
| 483 | index_set_id=index_set_id, |
| 484 | rules=[ |
| 485 | { |
| 486 | "field": "rule_group1", |
| 487 | "type": 1, |
| 488 | "inverted": False, |
| 489 | "value": "office365", |
| 490 | }, |
| 491 | { |
| 492 | "field": "data_office365_OrganizationId", |
| 493 | "type": 1, |
| 494 | "inverted": False, |
| 495 | "value": f"{provision_office365_auth_keys.TENANT_ID}", |
| 496 | }, |
| 497 | ], |
| 498 | matching_type="AND", |
| 499 | remove_matches_from_default_stream=True, |
| 500 | content_pack=None, |
| 501 | ) |
| 502 | |
| 503 | |
| 504 | async def send_event_stream_creation_request( |
| 505 | event_stream: Office365EventStream, |
| 506 | ) -> StreamCreationResponse: |
| 507 | """ |
| 508 | Sends a request to create an event stream. |
| 509 | |
| 510 | Args: |
| 511 | event_stream (WazuhEventStream): The event stream to be created. |
| 512 | |
| 513 | Returns: |
| 514 | StreamCreationResponse: The response containing the created event stream. |
| 515 | """ |
| 516 | json_event_stream = json.dumps(event_stream.model_dump()) |
| 517 | logger.info(f"json_event_stream set: {json_event_stream}") |
| 518 | response_json = await send_post_request_create_entity( |
| 519 | endpoint="/api/streams", |
| 520 | entity=event_stream.model_dump(), |
| 521 | ) |
| 522 | return StreamCreationResponse(**response_json) |
| 523 | |
| 524 | |
| 525 | async def create_event_stream( |
| 526 | customer_code: str, |
| 527 | provision_office365_auth_keys: ProvisionOffice365AuthKeys, |
| 528 | index_set_id: str, |
| 529 | session: AsyncSession, |
| 530 | ) -> StreamCreationResponse: |
| 531 | """ |
| 532 | Creates an event stream for a customer. |
| 533 | |
| 534 | Args: |
| 535 | request (ProvisionNewCustomer): The request object containing customer information. |
| 536 | index_set_id (str): The ID of the index set. |
| 537 | |
| 538 | Returns: |
| 539 | The result of the event stream creation request. |
| 540 | """ |
| 541 | event_stream_config = await build_event_stream_config( |
| 542 | customer_code, |
| 543 | provision_office365_auth_keys, |
| 544 | index_set_id, |
| 545 | session, |
| 546 | ) |
| 547 | return await send_event_stream_creation_request(event_stream_config) |
| 548 | |
| 549 | |
| 550 | ############### ! PIPELINES AND RULES ! ################ |
| 551 | |
| 552 | |
| 553 | # ! PIPELINE RULES ! # |
| 554 | async def check_pipeline_rules() -> None: |
| 555 | """ |
| 556 | Checks if the pipeline rules exist in Graylog. If they don't, create them. |
| 557 | """ |
| 558 | pipeline_rules = await get_pipeline_rules() |
| 559 | non_existing_rules = await pipeline_rules_exists(pipeline_rules) |
| 560 | if non_existing_rules: |
| 561 | logger.info(f"Creating pipeline rules: {non_existing_rules}") |
| 562 | await create_pipeline_rules(non_existing_rules) |
| 563 | |
| 564 | |
| 565 | async def pipeline_rules_exists(pipeline_rules: PipelineRulesResponse) -> List[str]: |
| 566 | """ |
| 567 | Checks if the pipeline rules exist in Graylog and returns a list of non-existing pipeline rules. |
| 568 | """ |
| 569 | return [ |
| 570 | rule_title.value |
| 571 | for rule_title in PipelineRuleTitles |
| 572 | if not any(rule.title == rule_title.value for rule in pipeline_rules.pipeline_rules) |
| 573 | ] |
| 574 | |
| 575 | |
| 576 | async def create_pipeline_rules(non_existing_rules: List[str]) -> None: |
| 577 | """ |
| 578 | Creates the given pipeline rules. |
| 579 | """ |
| 580 | rule_creators = { |
| 581 | "Office365 Timestamp - UTC": create_office365_utc_rule, |
| 582 | "Office365 Syslog Type": create_office365_syslog_type_rule, |
| 583 | "WAZUH CREATE FIELD SYSLOG LEVEL - INFO": create_wazuh_info_rule, |
| 584 | "WAZUH CREATE FIELD SYSLOG LEVEL - WARNING": create_wazuh_warning_rule, |
| 585 | "WAZUH CREATE FIELD SYSLOG LEVEL - NOTICE": create_wazuh_notice_rule, |
| 586 | "WAZUH CREATE FIELD SYSLOG LEVEL - ALERT": create_wazuh_alert_rule, |
| 587 | } |
| 588 | |
| 589 | for rule_title in non_existing_rules: |
| 590 | logger.info(f"Creating pipeline rule {rule_title}.") |
| 591 | await rule_creators[rule_title](rule_title) |
| 592 | |
| 593 | |
| 594 | async def create_office365_utc_rule(rule_title: str) -> None: |
| 595 | """ |
| 596 | Creates the 'Office365 Timestamp - UTC' pipeline rule. |
| 597 | """ |
| 598 | rule_source = ( |
| 599 | f'rule "{rule_title}"\n' |
| 600 | "when\n" |
| 601 | ' has_field("data_office365_CreationTime")\n' |
| 602 | "then\n" |
| 603 | " let creation_time = $message.data_office365_CreationTime;\n" |
| 604 | ' set_field("timestamp_utc", creation_time);\n' |
| 605 | "end" |
| 606 | ) |
| 607 | await create_pipeline_rule( |
| 608 | CreatePipelineRule( |
| 609 | title=rule_title, |
| 610 | description=rule_title, |
| 611 | source=rule_source, |
| 612 | ), |
| 613 | ) |
| 614 | |
| 615 | |
| 616 | async def create_office365_syslog_type_rule(rule_title: str) -> None: |
| 617 | """ |
| 618 | Creates the 'Office365 Syslog Type' pipeline rule. |
| 619 | """ |
| 620 | rule_source = ( |
| 621 | f'rule "{rule_title}"\n' |
| 622 | "when\n" |
| 623 | ' $message.rule_group1 == "office365"\n' |
| 624 | "then\n" |
| 625 | ' set_field("syslog_type", "office365");\n' |
| 626 | "end" |
| 627 | ) |
| 628 | await create_pipeline_rule( |
| 629 | CreatePipelineRule( |
| 630 | title=rule_title, |
| 631 | description=rule_title, |
| 632 | source=rule_source, |
| 633 | ), |
| 634 | ) |
| 635 | |
| 636 | |
| 637 | async def create_wazuh_info_rule(rule_title: str) -> None: |
| 638 | """ |
| 639 | Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - INFO' pipeline rule. |
| 640 | """ |
| 641 | rule_source = ( |
| 642 | f'rule "{rule_title}"\n' |
| 643 | "when\n" |
| 644 | " to_long($message.rule_level) > 0 AND to_long($message.rule_level) < 4\n" |
| 645 | "then\n" |
| 646 | ' set_field("syslog_level", "INFO");\n' |
| 647 | "end" |
| 648 | ) |
| 649 | await create_pipeline_rule( |
| 650 | CreatePipelineRule( |
| 651 | title=rule_title, |
| 652 | description=rule_title, |
| 653 | source=rule_source, |
| 654 | ), |
| 655 | ) |
| 656 | |
| 657 | |
| 658 | async def create_wazuh_warning_rule(rule_title: str) -> None: |
| 659 | """ |
| 660 | Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - WARNING' pipeline rule. |
| 661 | """ |
| 662 | rule_source = ( |
| 663 | f'rule "{rule_title}"\n' |
| 664 | "when\n" |
| 665 | " to_long($message.rule_level) > 7 AND to_long($message.rule_level) < 12\n" |
| 666 | "then\n" |
| 667 | ' set_field("syslog_level", "WARNING");\n' |
| 668 | "end" |
| 669 | ) |
| 670 | await create_pipeline_rule( |
| 671 | CreatePipelineRule( |
| 672 | title=rule_title, |
| 673 | description=rule_title, |
| 674 | source=rule_source, |
| 675 | ), |
| 676 | ) |
| 677 | |
| 678 | |
| 679 | async def create_wazuh_notice_rule(rule_title: str) -> None: |
| 680 | """ |
| 681 | Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - NOTICE' pipeline rule. |
| 682 | """ |
| 683 | rule_source = ( |
| 684 | f'rule "{rule_title}"\n' |
| 685 | "when\n" |
| 686 | " to_long($message.rule_level) > 3 AND to_long($message.rule_level) < 8\n" |
| 687 | "then\n" |
| 688 | ' set_field("syslog_level", "NOTICE");\n' |
| 689 | "end" |
| 690 | ) |
| 691 | await create_pipeline_rule( |
| 692 | CreatePipelineRule( |
| 693 | title=rule_title, |
| 694 | description=rule_title, |
| 695 | source=rule_source, |
| 696 | ), |
| 697 | ) |
| 698 | |
| 699 | |
| 700 | async def create_wazuh_alert_rule(rule_title: str) -> None: |
| 701 | """ |
| 702 | Creates the 'WAZUH CREATE FIELD SYSLOG LEVEL - ALERT' pipeline rule. |
| 703 | """ |
| 704 | rule_source = ( |
| 705 | f'rule "{rule_title}"\n' "when\n" " to_long($message.rule_level) > 11\n" "then\n" ' set_field("syslog_level", "ALERT");\n' "end" |
| 706 | ) |
| 707 | await create_pipeline_rule( |
| 708 | CreatePipelineRule( |
| 709 | title=rule_title, |
| 710 | description=rule_title, |
| 711 | source=rule_source, |
| 712 | ), |
| 713 | ) |
| 714 | |
| 715 | |
| 716 | # ! PIPELINE ! # |
| 717 | async def check_pipeline() -> None: |
| 718 | """ |
| 719 | Checks if the pipeline exists in Graylog. If it doesn't, create it. |
| 720 | """ |
| 721 | pipelines = await get_pipelines() |
| 722 | non_existing_pipelines = await pipeline_exists(pipelines) |
| 723 | if non_existing_pipelines: |
| 724 | logger.info(f"Creating pipelines: {non_existing_pipelines}") |
| 725 | await create_pipeline(non_existing_pipelines) |
| 726 | |
| 727 | |
| 728 | async def pipeline_exists(pipelines: GraylogPipelinesResponse) -> List[str]: |
| 729 | """ |
| 730 | Checks if the pipeline exists in Graylog and returns a list of non-existing pipelines. |
| 731 | """ |
| 732 | return [ |
| 733 | pipeline_title.value |
| 734 | for pipeline_title in PipelineTitles |
| 735 | if not any(pipeline.title == pipeline_title.value for pipeline in pipelines.pipelines) |
| 736 | ] |
| 737 | |
| 738 | |
| 739 | async def create_pipeline(non_existing_pipelines: List[str]) -> None: |
| 740 | """ |
| 741 | Creates the given pipeline. |
| 742 | """ |
| 743 | pipeline_creators = { |
| 744 | "OFFICE365 PROCESSING PIPELINE": create_office365_pipeline, |
| 745 | } |
| 746 | |
| 747 | for pipeline_title in non_existing_pipelines: |
| 748 | logger.info(f"Creating pipeline {pipeline_title}.") |
| 749 | await pipeline_creators[pipeline_title](pipeline_title) |
| 750 | |
| 751 | |
| 752 | async def create_office365_pipeline(pipeline_title: str) -> None: |
| 753 | """ |
| 754 | Creates the 'OFFICE365 PROCESSING PIPELINE' pipeline. |
| 755 | """ |
| 756 | pipeline_description = "OFFICE365 PROCESSING PIPELINE" |
| 757 | pipeline_source = ( |
| 758 | 'pipeline "OFFICE365 PROCESSING PIPELINE"\n' |
| 759 | "stage 0 match either\n" |
| 760 | 'rule "WAZUH CREATE FIELD SYSLOG LEVEL - ALERT"\n' |
| 761 | 'rule "WAZUH CREATE FIELD SYSLOG LEVEL - INFO"\n' |
| 762 | 'rule "WAZUH CREATE FIELD SYSLOG LEVEL - NOTICE"\n' |
| 763 | 'rule "WAZUH CREATE FIELD SYSLOG LEVEL - WARNING"\n' |
| 764 | 'rule "Office365 Timestamp - UTC"\n' |
| 765 | 'rule "SYSLOG TYPE OFFICE365"\n' |
| 766 | "end" |
| 767 | ) |
| 768 | await create_pipeline_graylog( |
| 769 | CreatePipeline( |
| 770 | title=pipeline_title, |
| 771 | description=pipeline_description, |
| 772 | source=pipeline_source, |
| 773 | ), |
| 774 | ) |
| 775 | |
| 776 | |
| 777 | #### ! GRAFANA ! #### |
| 778 | async def create_grafana_datasource( |
| 779 | customer_code: str, |
| 780 | session: AsyncSession, |
| 781 | ) -> GrafanaDataSourceCreationResponse: |
| 782 | """ |
| 783 | Creates a Grafana Wazuh datasource for a new customer using the OpenSearch Data Source. |
| 784 | |
| 785 | Args: |
| 786 | request (ProvisionNewCustomer): The request object containing customer information. |
| 787 | organization_id (int): The ID of the organization to create the datasource for. |
| 788 | session (AsyncSession): The database session. |
| 789 | |
| 790 | Returns: |
| 791 | GrafanaDataSourceCreationResponse: The response object containing the result of the datasource creation. |
| 792 | """ |
| 793 | logger.info("Creating Grafana datasource") |
| 794 | # Lowercase the customer code since Graylog index sets must be lowercase |
| 795 | customer_code = customer_code.lower() |
| 796 | grafana_client = await create_grafana_client("Grafana") |
| 797 | grafana_url = await get_connector_attribute( |
| 798 | connector_id=12, |
| 799 | column_name="connector_url", |
| 800 | session=session, |
| 801 | ) |
| 802 | # Switch to the newly created organization |
| 803 | grafana_client.user.switch_actual_user_organisation( |
| 804 | (await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, |
| 805 | ) |
| 806 | datasource_payload = GrafanaDatasource( |
| 807 | name="O365", |
| 808 | type="grafana-opensearch-datasource", |
| 809 | typeName="OpenSearch", |
| 810 | access="proxy", |
| 811 | url=await get_connector_attribute( |
| 812 | connector_id=1, |
| 813 | column_name="connector_url", |
| 814 | session=session, |
| 815 | ), |
| 816 | database=f"office365-{customer_code}*", |
| 817 | basicAuth=True, |
| 818 | basicAuthUser=await get_connector_attribute( |
| 819 | connector_id=1, |
| 820 | column_name="connector_username", |
| 821 | session=session, |
| 822 | ), |
| 823 | secureJsonData={ |
| 824 | "basicAuthPassword": await get_connector_attribute( |
| 825 | connector_id=1, |
| 826 | column_name="connector_password", |
| 827 | session=session, |
| 828 | ), |
| 829 | }, |
| 830 | isDefault=False, |
| 831 | jsonData={ |
| 832 | "dataLinks": [ |
| 833 | { |
| 834 | "field": "^_id$", |
| 835 | "url": ( |
| 836 | "{}/explore?left=%7B%22datasource%22:%22O365%22,%22queries%22:%5B%7B" |
| 837 | "%22refId%22:%22A%22,%22query%22:%22_id:${{__value.raw}}%22,%22alias%22:%22%22," |
| 838 | "%22metrics%22:%5B%7B%22id%22:%221%22,%22type%22:%22logs%22,%22settings%22:" |
| 839 | "%7B%22limit%22:%22500%22%7D%7D%5D,%22bucketAggs%22:%5B%7B%22type%22:%22date_histogram%22,%22id%22:%221%22,%22settings%22:%7B%22interval%22:%22auto%22%7D%7D%5D,%22timeField%22:" |
| 840 | "%22timestamp%22%7D%5D,%22range%22:%7B%22from%22:%22now-6h%22,%22to%22:%22now%22%7D%7D" |
| 841 | ).format(grafana_url), |
| 842 | }, |
| 843 | ], |
| 844 | "database": f"office365-{customer_code}*", |
| 845 | "flavor": "opensearch", |
| 846 | "includeFrozen": False, |
| 847 | "logLevelField": "syslog_level", |
| 848 | "logMessageField": "rule_description", |
| 849 | "maxConcurrentShardRequests": 5, |
| 850 | "pplEnabled": True, |
| 851 | "timeField": "timestamp", |
| 852 | "tlsSkipVerify": True, |
| 853 | "version": await get_opensearch_version(), |
| 854 | }, |
| 855 | readOnly=True, |
| 856 | ) |
| 857 | results = grafana_client.datasource.create_datasource( |
| 858 | datasource=datasource_payload.model_dump(), |
| 859 | ) |
| 860 | return GrafanaDataSourceCreationResponse(**results) |
| 861 | |
| 862 | |
| 863 | ################## ! MAIN FUNCTION ! ################## |
| 864 | |
| 865 | |
| 866 | async def provision_office365( |
| 867 | customer_code: str, |
| 868 | provision_office365_auth_keys: ProvisionOffice365AuthKeys, |
| 869 | session: AsyncSession, |
| 870 | ) -> ProvisionOffice365Response: |
| 871 | logger.info(f"Provisioning Office365 integration for customer {customer_code}.") |
| 872 | |
| 873 | # Get Wazuh configuration |
| 874 | wazuh_config = await get_wazuh_configuration(file_name="wazuh_config.xml") |
| 875 | |
| 876 | # Check if Office365 is already provisioned |
| 877 | office365_provisioned = await check_if_office365_is_already_provisioned(customer_code, wazuh_config) |
| 878 | |
| 879 | # Create Office365 template |
| 880 | if office365_provisioned: |
| 881 | logger.info("Office365 integration is already provisioned.") |
| 882 | else: |
| 883 | logger.info("Office365 integration is not yet provisioned.") |
| 884 | logger.info("Creating new Office365 block.") |
| 885 | office365_templated = await office365_template_with_api_type( |
| 886 | customer_code, |
| 887 | provision_office365_auth_keys, |
| 888 | ) |
| 889 | |
| 890 | # Check if Office365 is already provisioned for customer |
| 891 | await check_if_office365_is_already_provisioned_for_customer(provision_office365_auth_keys.TENANT_ID, wazuh_config) |
| 892 | |
| 893 | # If Office365 is already provisioned but not for the customer, add the api_auth contents to the office365 block |
| 894 | if office365_provisioned and not await check_if_office365_is_already_provisioned_for_customer( |
| 895 | provision_office365_auth_keys.TENANT_ID, |
| 896 | wazuh_config, |
| 897 | ): |
| 898 | wazuh_config = await add_api_auth_to_office365_block(customer_code, provision_office365_auth_keys) |
| 899 | else: |
| 900 | # Append Office365 template to Wazuh configuration |
| 901 | wazuh_config = await append_office365_template(wazuh_config, office365_templated) |
| 902 | |
| 903 | # Update Wazuh configuration |
| 904 | # await update_wazuh_configuration(wazuh_config, provision_office365_auth_keys) |
| 905 | await update_wazuh_configuration(provision_office365_auth_keys) |
| 906 | |
| 907 | # Restart Wazuh manager |
| 908 | await restart_wazuh_manager() |
| 909 | |
| 910 | # Graylog Deployment |
| 911 | await check_pipeline_rules() |
| 912 | await check_pipeline() |
| 913 | |
| 914 | # Create Index Set |
| 915 | index_set_id = (await create_index_set(customer_code=customer_code, session=session)).data.id |
| 916 | logger.info(f"Index set: {index_set_id}") |
| 917 | # Create event stream |
| 918 | stream_id = ( |
| 919 | await create_event_stream( |
| 920 | customer_code, |
| 921 | provision_office365_auth_keys, |
| 922 | index_set_id, |
| 923 | session, |
| 924 | ) |
| 925 | ).data.stream_id |
| 926 | pipeline_id = await get_pipeline_id(subscription="OFFICE365") |
| 927 | # Combine stream and pipeline IDs |
| 928 | stream_and_pipeline = StreamConnectionToPipelineRequest( |
| 929 | stream_id=stream_id, |
| 930 | pipeline_ids=pipeline_id, |
| 931 | ) |
| 932 | # Connect stream to pipeline |
| 933 | logger.info(f"Stream and pipeline: {stream_and_pipeline}") |
| 934 | await connect_stream_to_pipeline(stream_and_pipeline) |
| 935 | # Start stream |
| 936 | await start_stream(stream_id=stream_id) |
| 937 | |
| 938 | # Grafana Deployment |
| 939 | office365_datasource_uid = (await create_grafana_datasource(customer_code=customer_code, session=session)).datasource.uid |
| 940 | grafana_o365_folder_id = ( |
| 941 | await create_grafana_folder( |
| 942 | organization_id=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, |
| 943 | folder_title="OFFICE 365", |
| 944 | ) |
| 945 | ).id |
| 946 | await provision_dashboards( |
| 947 | DashboardProvisionRequest( |
| 948 | dashboards=[dashboard.name for dashboard in Office365Dashboard], |
| 949 | organizationId=(await get_customer_meta(customer_code, session)).customer_meta.customer_meta_grafana_org_id, |
| 950 | folderId=grafana_o365_folder_id, |
| 951 | datasourceUid=office365_datasource_uid, |
| 952 | grafana_url=(await get_customer_default_settings_attribute(column_name="grafana_url", session=session)) |
| 953 | or "grafana.company.local", |
| 954 | ), |
| 955 | ) |
| 956 | |
| 957 | await update_customer_integration_table(customer_code, session) |
| 958 | await update_customermeta_table(customer_code, session, provision_office365_auth_keys.TENANT_ID) |
| 959 | |
| 960 | await create_integration_meta_entry( |
| 961 | CustomerIntegrationsMetaSchema( |
| 962 | customer_code=customer_code, |
| 963 | integration_name="Office365", |
| 964 | graylog_input_id=None, |
| 965 | graylog_index_id=index_set_id, |
| 966 | graylog_stream_id=stream_id, |
| 967 | grafana_org_id=( |
| 968 | await get_customer_meta( |
| 969 | customer_code, |
| 970 | session, |
| 971 | ) |
| 972 | ).customer_meta.customer_meta_grafana_org_id, |
| 973 | grafana_dashboard_folder_id=grafana_o365_folder_id, |
| 974 | grafana_datasource_uid=office365_datasource_uid, |
| 975 | ), |
| 976 | session, |
| 977 | ) |
| 978 | |
| 979 | return ProvisionOffice365Response( |
| 980 | success=True, |
| 981 | message=f"Successfully provisioned Office365 integration for customer {customer_code}.", |
| 982 | ) |
| 983 | |
| 984 | |
| 985 | ######### ! Update Database ! ############ |
| 986 | async def update_customer_integration_table( |
| 987 | customer_code: str, |
| 988 | session: AsyncSession, |
| 989 | ) -> None: |
| 990 | """ |
| 991 | Updates the `customer_integrations` table to set the `deployed` column to True where the `customer_code` |
| 992 | matches the given customer code and the `integration_service_name` is "Office365". |
| 993 | |
| 994 | Args: |
| 995 | customer_code (str): The customer code. |
| 996 | session (AsyncSession): The async session object for making HTTP requests. |
| 997 | """ |
| 998 | await session.execute( |
| 999 | update(CustomerIntegrations) |
| 1000 | .where( |
| 1001 | and_( |
| 1002 | CustomerIntegrations.customer_code == customer_code, |
| 1003 | CustomerIntegrations.integration_service_name == "Office365", |
| 1004 | ), |
| 1005 | ) |
| 1006 | .values(deployed=True), |
| 1007 | ) |
| 1008 | await session.commit() |
| 1009 | |
| 1010 | return None |
| 1011 | |
| 1012 | |
| 1013 | async def update_customermeta_table(customer_code: str, session: AsyncSession, tenant_id: str) -> None: |
| 1014 | """ |
| 1015 | Updates the `customer_meta` table to set the `office365_tenant_id` column to the given tenant_id. |
| 1016 | |
| 1017 | Args: |
| 1018 | customer_code (str): The customer code. |
| 1019 | session (AsyncSession): The async session object for making HTTP requests. |
| 1020 | """ |
| 1021 | await session.execute( |
| 1022 | update(CustomersMeta).where(CustomersMeta.customer_code == customer_code).values(customer_meta_office365_organization_id=tenant_id), |
| 1023 | ) |
| 1024 | await session.commit() |
| 1025 | |
| 1026 | return None |
| 1027 | |
| 1028 | |
| 1029 | async def create_integration_meta_entry( |
| 1030 | customer_integration_meta: CustomerIntegrationsMetaSchema, |
| 1031 | session: AsyncSession, |
| 1032 | ) -> None: |
| 1033 | """ |
| 1034 | Creates an entry for the customer integration meta in the database. |
| 1035 | |
| 1036 | Args: |
| 1037 | customer_integration_meta (CustomerIntegrationsMetaSchema): The customer integration meta object. |
| 1038 | session (AsyncSession): The async session object for database operations. |
| 1039 | """ |
| 1040 | await create_integration_meta(customer_integration_meta, session) |
| 1041 | logger.info( |
| 1042 | f"Integration meta entry created for customer {customer_integration_meta.customer_code}.", |
| 1043 | ) |