| 1 | import os |
| 2 | |
| 3 | from dotenv import load_dotenv |
| 4 | from loguru import logger |
| 5 | from sqlalchemy import and_ |
| 6 | from sqlalchemy.ext.asyncio import AsyncSession |
| 7 | from sqlalchemy.future import select |
| 8 | |
| 9 | from app.auth.models.users import Role |
| 10 | from app.connectors.models import Connectors |
| 11 | from app.integrations.models.customer_integration_settings import AvailableIntegrations |
| 12 | from app.integrations.models.customer_integration_settings import ( |
| 13 | AvailableIntegrationsAuthKeys, |
| 14 | ) |
| 15 | from app.network_connectors.models.network_connectors import AvailableNetworkConnectors |
| 16 | from app.network_connectors.models.network_connectors import ( |
| 17 | AvailableNetworkConnectorsKeys, |
| 18 | ) |
| 19 | |
| 20 | load_dotenv() |
| 21 | |
| 22 | |
| 23 | def load_connector_data( |
| 24 | connector_name, |
| 25 | connector_type, |
| 26 | accepts_key, |
| 27 | description, |
| 28 | extra_data_key=None, |
| 29 | ): |
| 30 | """ |
| 31 | Load connector data from environment variables. |
| 32 | |
| 33 | Args: |
| 34 | connector_name (str): The name of the connector. |
| 35 | connector_type (str): The type of the connector. |
| 36 | accepts_key (str): The type of key the connector accepts. |
| 37 | description (str): The description of the connector. |
| 38 | extra_data_key (str, optional): The key for extra data. Defaults to None. |
| 39 | |
| 40 | Returns: |
| 41 | dict: A dictionary containing the connector data. |
| 42 | """ |
| 43 | env_prefix = connector_name.upper().replace("-", "_").replace(" ", "_") |
| 44 | url = os.getenv(f"{env_prefix}_URL") |
| 45 | logger.info( |
| 46 | f"Loading connector data for {connector_name} from environment variables with URL: {url}", |
| 47 | ) |
| 48 | return { |
| 49 | "connector_name": connector_name, |
| 50 | "connector_type": connector_type, |
| 51 | "connector_url": os.getenv(f"{env_prefix}_URL"), |
| 52 | "connector_username": os.getenv(f"{env_prefix}_USERNAME"), |
| 53 | "connector_password": os.getenv(f"{env_prefix}_PASSWORD"), |
| 54 | "connector_api_key": os.getenv(f"{env_prefix}_API_KEY"), |
| 55 | "connector_description": description, |
| 56 | "connector_supports": os.getenv(f"{env_prefix}_SUPPORTS", "Not specified."), |
| 57 | "connector_configured": True, |
| 58 | "connector_verified": bool(os.getenv(f"{env_prefix}_VERIFIED", False)), |
| 59 | "connector_accepts_host_only": accepts_key == "host_only", |
| 60 | "connector_accepts_api_key": accepts_key == "api_key", |
| 61 | "connector_accepts_username_password": accepts_key == "username_password", |
| 62 | "connector_accepts_file": accepts_key == "file", |
| 63 | "connector_accepts_extra_data": True if extra_data_key else False, |
| 64 | "connector_extra_data": os.getenv(extra_data_key) if extra_data_key else None, |
| 65 | } |
| 66 | |
| 67 | |
| 68 | def get_connectors_list(): |
| 69 | """ |
| 70 | Get a list of connectors with their respective versions and authentication methods. |
| 71 | |
| 72 | Returns: |
| 73 | list: A list of connector data, where each item contains the connector name, version, and authentication method. |
| 74 | """ |
| 75 | connectors = [ |
| 76 | ("Wazuh-Indexer", "4.4.1", "username_password", "Connection to Wazuh-Indexer."), |
| 77 | ( |
| 78 | "Wazuh-Manager", |
| 79 | "4.4.1", |
| 80 | "username_password", |
| 81 | "Connection to Wazuh-Manager. Default is wazuh-wui:wazuh-wui", |
| 82 | ), |
| 83 | ("Graylog", "5.0.7", "username_password", "Connection to Graylog."), |
| 84 | ("Shuffle", "1.1.0", "api_key", "Connection to Shuffle."), |
| 85 | # ("DFIR-IRIS", "2.0", "api_key", "Connection to DFIR-IRIS."), |
| 86 | ( |
| 87 | "Velociraptor", |
| 88 | "0.6.8", |
| 89 | "file", |
| 90 | "Connection to Velociraptor. Make sure you have generated the api file first.", |
| 91 | ), |
| 92 | ("Sublime", "3", "api_key", "Connection to Sublime."), |
| 93 | ( |
| 94 | "InfluxDB", |
| 95 | "3", |
| 96 | "api_key", |
| 97 | "Connection to InfluxDB.", |
| 98 | "INFLUXDB_ORG_AND_BUCKET", |
| 99 | ), |
| 100 | # ( |
| 101 | # "AskSocfortress", |
| 102 | # "3", |
| 103 | # "api_key", |
| 104 | # "Connection to AskSocfortress. Make sure you have requested an API key.", |
| 105 | # ), |
| 106 | # ( |
| 107 | # "SocfortressThreatIntel", |
| 108 | # "3", |
| 109 | # "api_key", |
| 110 | # "Connection to Socfortress Threat Intel. Make sure you have requested an API key.", |
| 111 | # ), |
| 112 | # ( |
| 113 | # "Cortex", |
| 114 | # "3", |
| 115 | # "api_key", |
| 116 | # "Connection to Cortex. Make sure you have created an API key.", |
| 117 | # ), |
| 118 | ("Grafana", "3", "username_password", "Connection to Grafana."), |
| 119 | # ! TODO - LOOK TO REMOVE WAZUH WORKER PROVISIONING FROM CONNECTORS LIST ! # |
| 120 | ( |
| 121 | "Wazuh Worker Provisioning", |
| 122 | "3", |
| 123 | "host_only", |
| 124 | ( |
| 125 | "Connection to Wazuh Worker Provisioning. Make sure you have " |
| 126 | "deployed the Wazuh Worker Provisioning Application provided by " |
| 127 | "SOCFortress: https://github.com/socfortress/Customer-Provisioning-Worker" |
| 128 | ), |
| 129 | ), |
| 130 | ( |
| 131 | "Event Shipper", |
| 132 | "3", |
| 133 | "host_only", |
| 134 | "Connection to Graylog GELF Input to receive events from integrations. Make sure you have created a GELF Input in Graylog.", |
| 135 | "GELF_INPUT_PORT", |
| 136 | ), |
| 137 | ( |
| 138 | "HAProxy Provisioning", |
| 139 | "3", |
| 140 | "host_only", |
| 141 | ( |
| 142 | "Connection to HAProxy Provisioning. Make sure you have deployed the " |
| 143 | "HAProxy Provisioning Application provided by " |
| 144 | "SOCFortress: https://github.com/socfortress/Customer-Provisioning-Worker" |
| 145 | ), |
| 146 | ), |
| 147 | ( |
| 148 | "VirusTotal", |
| 149 | "3", |
| 150 | "api_key", |
| 151 | "Connection to VirusTotal. Make sure you have created an API key.", |
| 152 | ), |
| 153 | ("Portainer", "3", "username_password", "Connection to Portainer.", "PORTAINER_ENDPOINT_ID"), |
| 154 | ( |
| 155 | "Graylog-Network", |
| 156 | "5.0.7", |
| 157 | "username_password", |
| 158 | "Connection to Graylog. If you only have one Graylog instance, set this to the same as Graylog.", |
| 159 | ), |
| 160 | ( |
| 161 | "Talon", |
| 162 | "3", |
| 163 | "api_key", |
| 164 | "Talon is an automated AI SOC analyst built by SOCfortress for the CoPilot stack. It runs as a background service alongside CoPilot — pulling raw events from your Wazuh/OpenSearch SIEM, enriching them with threat intelligence, correlating across your environment, and writing structured investigation reports with severity assessments and recommended actions directly back into CoPilot.", |
| 165 | ), |
| 166 | # ... Add more connectors as needed ... |
| 167 | ] |
| 168 | |
| 169 | return [load_connector_data(*connector) for connector in connectors] |
| 170 | |
| 171 | |
| 172 | def delete_connectors_list(): |
| 173 | """ |
| 174 | Get a list of connectors with their respective versions and authentication methods. |
| 175 | |
| 176 | Returns: |
| 177 | list: A list of connector data, where each item contains the connector name, version, and authentication method. |
| 178 | """ |
| 179 | connectors = [ |
| 180 | "DFIR-IRIS", |
| 181 | "AskSocfortress", |
| 182 | "SocfortressThreatIntel", |
| 183 | "Cortex", |
| 184 | ] |
| 185 | |
| 186 | return connectors |
| 187 | |
| 188 | |
| 189 | async def add_connectors_if_not_exist(session: AsyncSession): |
| 190 | """ |
| 191 | Adds connectors to the database if they do not already exist. |
| 192 | |
| 193 | Args: |
| 194 | session (AsyncSession): The database session. |
| 195 | |
| 196 | Returns: |
| 197 | None |
| 198 | """ |
| 199 | connector_list = get_connectors_list() |
| 200 | logger.info("Checking for existence of connectors.") |
| 201 | |
| 202 | for connector_data in connector_list: |
| 203 | logger.info(f"Checking for existence of connector {connector_data['connector_name']}") |
| 204 | query = select(Connectors).where( |
| 205 | Connectors.connector_name == connector_data["connector_name"], |
| 206 | ) |
| 207 | result = await session.execute(query) |
| 208 | existing_connector = result.scalars().first() |
| 209 | |
| 210 | if existing_connector is None: |
| 211 | new_connector = Connectors(**connector_data) |
| 212 | session.add(new_connector) |
| 213 | logger.info(f"Added new connector: {connector_data['connector_name']}") |
| 214 | |
| 215 | await session.commit() |
| 216 | |
| 217 | |
| 218 | async def delete_connectors_if_exist(session: AsyncSession): |
| 219 | """ |
| 220 | Deletes connectors from the database if they already exist. |
| 221 | |
| 222 | Args: |
| 223 | session (AsyncSession): The database session. |
| 224 | |
| 225 | Returns: |
| 226 | None |
| 227 | """ |
| 228 | connector_list = delete_connectors_list() |
| 229 | logger.info("Checking for existence of connectors. This connector will be deleted.") |
| 230 | |
| 231 | for connector_data in connector_list: |
| 232 | logger.info(f"Checking for existence of connector {connector_data}") |
| 233 | query = select(Connectors).where( |
| 234 | Connectors.connector_name == connector_data, |
| 235 | ) |
| 236 | result = await session.execute(query) |
| 237 | existing_connector = result.scalars().first() |
| 238 | |
| 239 | if existing_connector is not None: |
| 240 | await session.delete(existing_connector) |
| 241 | logger.info(f"Deleted connector: {connector_data}") |
| 242 | |
| 243 | await session.commit() |
| 244 | |
| 245 | |
| 246 | async def add_roles_if_not_exist(session: AsyncSession) -> None: |
| 247 | """ |
| 248 | Adds roles to the database if they do not already exist. |
| 249 | |
| 250 | Args: |
| 251 | session (AsyncSession): The database session. |
| 252 | |
| 253 | Returns: |
| 254 | None |
| 255 | """ |
| 256 | # List of roles to add |
| 257 | role_list = [ |
| 258 | {"name": "admin", "description": "Administrator"}, |
| 259 | {"name": "analyst", "description": "SOC Analyst"}, |
| 260 | {"name": "scheduler", "description": "Scheduler for automated tasks"}, |
| 261 | {"name": "customer_user", "description": "Customer user with limited access to their own data"}, |
| 262 | ] |
| 263 | |
| 264 | for role_data in role_list: |
| 265 | logger.info(f"Checking for existence of role {role_data['name']}") |
| 266 | query = select(Role).where(Role.name == role_data["name"]) |
| 267 | result = await session.execute(query) |
| 268 | existing_role = result.scalars().first() |
| 269 | |
| 270 | if existing_role is None: |
| 271 | new_role = Role(**role_data) |
| 272 | session.add(new_role) # Use session.add() to add new objects |
| 273 | logger.info(f"Added new role: {role_data['name']}") |
| 274 | |
| 275 | await session.commit() # Commit the transaction |
| 276 | logger.info("Role check and addition completed.") |
| 277 | |
| 278 | |
| 279 | # ! AVAILABLE THIRD PARTY INTEGRATIONS ! # |
| 280 | def load_available_integrations_data( |
| 281 | integration_name: str, |
| 282 | description: str, |
| 283 | integration_details: str, |
| 284 | ): |
| 285 | """ |
| 286 | Load available integrations data from environment variables. |
| 287 | |
| 288 | Args: |
| 289 | integration_name (str): The name of the integration. |
| 290 | description (str): The description of the integration. |
| 291 | |
| 292 | Returns: |
| 293 | dict: A dictionary containing the integration data. |
| 294 | """ |
| 295 | logger.info(f"Loading available integrations data for {integration_name}.") |
| 296 | return { |
| 297 | "integration_name": integration_name, |
| 298 | "description": description, |
| 299 | "integration_details": integration_details, |
| 300 | } |
| 301 | |
| 302 | |
| 303 | def load_markdown_for_integration(integration_name: str) -> str: |
| 304 | """ |
| 305 | Load markdown content for a given integration from a file. |
| 306 | |
| 307 | Args: |
| 308 | integration_name (str): The name of the integration. |
| 309 | |
| 310 | Returns: |
| 311 | str: The content of the markdown file. |
| 312 | """ |
| 313 | # file_path = os.path.join("integrations_markdown", f"{integration_name.lower()}.md") |
| 314 | # if space in the integration name, replace it with underscore |
| 315 | if " " in integration_name: |
| 316 | integration_name = integration_name.replace(" ", "_") |
| 317 | file_path = os.path.join( |
| 318 | "app", |
| 319 | "integrations", |
| 320 | "markdown", |
| 321 | f"{integration_name.lower()}.md", |
| 322 | ) |
| 323 | try: |
| 324 | with open(file_path, "r") as file: |
| 325 | return file.read() |
| 326 | except FileNotFoundError: |
| 327 | return "No deployment instructions available." |
| 328 | |
| 329 | |
| 330 | def get_available_integrations_list(): |
| 331 | """ |
| 332 | Get a list of available integrations. |
| 333 | |
| 334 | Returns: |
| 335 | list: A list of available integrations data, where each item contains the integration name, description, and markdown details. |
| 336 | """ |
| 337 | available_integrations = [ |
| 338 | ("Office365", "Integrate Office365 with SOCFortress."), |
| 339 | ("Mimecast", "Integrate Mimecast with SOCFortress."), |
| 340 | ("SAP SIEM", "Integrate SAP SIEM with SOCFortress."), |
| 341 | ("Huntress", "Integrate Huntress with SOCFortress."), |
| 342 | ("CarbonBlack", "Integrate CarbonBlack with SOCFortress."), |
| 343 | ("Crowdstrike", "Integrate Crowdstrike with SOCFortress."), |
| 344 | ("DUO", "Integrate DUO with SOCFortress."), |
| 345 | ("Darktrace", "Integrate Darktrace with SOCFortress."), |
| 346 | ("BitDefender", "Integrate BitDefender with SOCFortress."), |
| 347 | ("CATO", "Integrate CATO NETWORKS with SOCFortress."), |
| 348 | ("DefenderForEndpoint", "Integrate DefenderForEndpoint with SOCFortress."), |
| 349 | ("SOCFortress MDR", "Forward alerts to the SOCFortress MDR server for this customer."), |
| 350 | # ... Add more available integrations as needed ... |
| 351 | ] |
| 352 | |
| 353 | return [ |
| 354 | load_available_integrations_data( |
| 355 | integration_name, |
| 356 | description, |
| 357 | load_markdown_for_integration(integration_name), |
| 358 | ) |
| 359 | for integration_name, description in available_integrations |
| 360 | ] |
| 361 | |
| 362 | |
| 363 | async def add_available_integrations_if_not_exist(session: AsyncSession): |
| 364 | """ |
| 365 | Adds available integrations to the database if they do not already exist. |
| 366 | |
| 367 | Args: |
| 368 | session (AsyncSession): The database session. |
| 369 | |
| 370 | Returns: |
| 371 | None |
| 372 | """ |
| 373 | available_integrations_list = get_available_integrations_list() |
| 374 | |
| 375 | for available_integration_data in available_integrations_list: |
| 376 | try: |
| 377 | query = select(AvailableIntegrations).where( |
| 378 | AvailableIntegrations.integration_name == available_integration_data["integration_name"], |
| 379 | ) |
| 380 | result = await session.execute(query) |
| 381 | existing_available_integration = result.scalars().first() |
| 382 | |
| 383 | if existing_available_integration is None: |
| 384 | new_available_integration = AvailableIntegrations( |
| 385 | **available_integration_data, |
| 386 | ) |
| 387 | logger.info(f"New available integration: {available_integration_data}") |
| 388 | session.add(new_available_integration) |
| 389 | logger.info( |
| 390 | f"Added new available integration: {available_integration_data['integration_name']}", |
| 391 | ) |
| 392 | else: |
| 393 | # Check if the integration details need to be updated |
| 394 | if existing_available_integration.integration_details == "No deployment instructions available.": |
| 395 | new_integration_details = load_markdown_for_integration(available_integration_data["integration_name"]) |
| 396 | if new_integration_details != "No deployment instructions available.": |
| 397 | existing_available_integration.integration_details = new_integration_details |
| 398 | logger.info( |
| 399 | f"Updated integration details for {available_integration_data['integration_name']}", |
| 400 | ) |
| 401 | except Exception as e: |
| 402 | logger.error(f"Error adding available integration: {e}") |
| 403 | await session.rollback() |
| 404 | raise e |
| 405 | await session.commit() |
| 406 | # Close the session |
| 407 | await session.close() |
| 408 | |
| 409 | |
| 410 | def load_available_integrations_auth_keys( |
| 411 | integration_id: int, |
| 412 | integration_name: str, |
| 413 | auth_key_name: str, |
| 414 | ): |
| 415 | """ |
| 416 | Load available integrations auth keys from environment variables. |
| 417 | |
| 418 | Args: |
| 419 | integration_id (int): The ID of the integration. |
| 420 | integration_name (str): The name of the integration. |
| 421 | auth_key_name (str): The name of the auth key. |
| 422 | |
| 423 | Returns: |
| 424 | dict: A dictionary containing the auth key data. |
| 425 | """ |
| 426 | logger.info( |
| 427 | f"Loading available integrations auth keys data for {integration_name}.", |
| 428 | ) |
| 429 | return { |
| 430 | "integration_id": integration_id, |
| 431 | "integration_name": integration_name, |
| 432 | "auth_key_name": auth_key_name, |
| 433 | } |
| 434 | |
| 435 | |
| 436 | async def get_available_integrations_auth_keys_list(session: AsyncSession): |
| 437 | """ |
| 438 | Get a list of available integrations auth keys with their corresponding integration IDs. |
| 439 | |
| 440 | Args: |
| 441 | session (AsyncSession): The database session. |
| 442 | |
| 443 | Returns: |
| 444 | list: A list of available integrations auth keys data, where each item contains the integration ID, integration name, and auth key name. |
| 445 | """ |
| 446 | available_integrations_auth_keys = [] |
| 447 | available_integrations = [ |
| 448 | ("Office365", "TENANT_ID"), |
| 449 | ("Office365", "CLIENT_ID"), |
| 450 | ("Office365", "CLIENT_SECRET"), |
| 451 | ("Office365", "API_TYPE"), |
| 452 | ("Mimecast", "APP_ID"), |
| 453 | ("Mimecast", "APP_KEY"), |
| 454 | ("Mimecast", "EMAIL_ADDRESS"), |
| 455 | ("Mimecast", "ACCESS_KEY"), |
| 456 | ("Mimecast", "SECRET_KEY"), |
| 457 | ("SAP SIEM", "API_KEY"), |
| 458 | ("SAP SIEM", "SECRET_KEY"), |
| 459 | ("SAP SIEM", "USER_KEY"), |
| 460 | ("SAP SIEM", "API_DOMAIN"), |
| 461 | ("Huntress", "API_KEY"), |
| 462 | ("Huntress", "API_SECRET"), |
| 463 | ("CarbonBlack", "API_KEY"), |
| 464 | ("CarbonBlack", "API_URL"), |
| 465 | ("CarbonBlack", "API_ID"), |
| 466 | ("CarbonBlack", "ORGANIZATION_KEY"), |
| 467 | ("Crowdstrike", "CLIENT_ID"), |
| 468 | ("Crowdstrike", "CLIENT_SECRET"), |
| 469 | ("Crowdstrike", "BASE_URL"), |
| 470 | ("Crowdstrike", "SYSLOG_PORT"), |
| 471 | ("DUO", "API_HOSTNAME"), |
| 472 | ("DUO", "INTEGRATION_KEY"), |
| 473 | ("DUO", "SECRET_KEY"), |
| 474 | ("Darktrace", "PUBLIC_TOKEN"), |
| 475 | ("Darktrace", "PRIVATE_TOKEN"), |
| 476 | ("Darktrace", "HOST"), |
| 477 | ("Darktrace", "PORT"), |
| 478 | ("BitDefender", "BASIC_AUTH_USERNAME"), |
| 479 | ("BitDefender", "BASIC_AUTH_PASSWORD"), |
| 480 | ("BitDefender", "WEBSERVER_HOSTNAME"), |
| 481 | ("BitDefender", "WEBSERVER_PORT"), |
| 482 | ("BitDefender", "GRAYLOG_PORT"), |
| 483 | ("BitDefender", "API_KEY"), |
| 484 | ("CATO", "API_KEY"), |
| 485 | ("CATO", "ACCOUNT_ID"), |
| 486 | ("CATO", "EVENT_TYPES"), |
| 487 | ("CATO", "EVENT_SUB_TYPES"), |
| 488 | ("DefenderForEndpoint", "TENANT_ID"), |
| 489 | ("DefenderForEndpoint", "CLIENT_ID"), |
| 490 | ("DefenderForEndpoint", "CLIENT_SECRET"), |
| 491 | ("DefenderForEndpoint", "SYSLOG_PORT"), |
| 492 | ("SOCFortress MDR", "COLLECTOR_UUID"), |
| 493 | # ... Add more available integrations auth keys as needed ... |
| 494 | ] |
| 495 | logger.info("Getting available integrations auth keys.") |
| 496 | try: |
| 497 | for integration_name, auth_key_name in available_integrations: |
| 498 | query = select(AvailableIntegrations.id).where( |
| 499 | AvailableIntegrations.integration_name == integration_name, |
| 500 | ) |
| 501 | result = await session.execute(query) |
| 502 | integration_id = result.scalars().first() |
| 503 | logger.info(f"Integration ID for {integration_name}: {integration_id}") |
| 504 | if integration_id: |
| 505 | logger.info(f"Found integration ID for {integration_name}: {integration_id}") |
| 506 | available_integrations_auth_keys.append( |
| 507 | load_available_integrations_auth_keys( |
| 508 | integration_id, |
| 509 | integration_name, |
| 510 | auth_key_name, |
| 511 | ), |
| 512 | ) |
| 513 | |
| 514 | return available_integrations_auth_keys |
| 515 | except Exception as e: |
| 516 | logger.error(f"Error getting available integrations auth keys: {e}") |
| 517 | await session.rollback() |
| 518 | raise e |
| 519 | |
| 520 | |
| 521 | async def add_available_integrations_auth_keys_if_not_exist(session: AsyncSession): |
| 522 | """ |
| 523 | Adds available integrations auth keys to the database if they do not already exist. |
| 524 | |
| 525 | Args: |
| 526 | session (AsyncSession): The database session. |
| 527 | |
| 528 | Returns: |
| 529 | None |
| 530 | """ |
| 531 | logger.info("Checking for existence of available integrations auth keys.") |
| 532 | available_integrations_auth_keys_list = await get_available_integrations_auth_keys_list(session=session) |
| 533 | logger.info("Adding available integrations auth keys to the database.") |
| 534 | for available_integration_auth_keys_data in available_integrations_auth_keys_list: |
| 535 | try: |
| 536 | query = select(AvailableIntegrations).where( |
| 537 | AvailableIntegrations.integration_name == available_integration_auth_keys_data["integration_name"], |
| 538 | ) |
| 539 | result = await session.execute(query) |
| 540 | existing_integration = result.scalars().first() |
| 541 | |
| 542 | if existing_integration: |
| 543 | available_integration_auth_keys_data["integration_id"] = existing_integration.id |
| 544 | auth_key_query = select(AvailableIntegrationsAuthKeys).where( |
| 545 | and_( |
| 546 | AvailableIntegrationsAuthKeys.integration_id == existing_integration.id, |
| 547 | AvailableIntegrationsAuthKeys.auth_key_name == available_integration_auth_keys_data["auth_key_name"], |
| 548 | ), |
| 549 | ) |
| 550 | auth_key_result = await session.execute(auth_key_query) |
| 551 | existing_auth_key = auth_key_result.scalars().first() |
| 552 | |
| 553 | if existing_auth_key is None: |
| 554 | new_auth_key = AvailableIntegrationsAuthKeys( |
| 555 | **available_integration_auth_keys_data, |
| 556 | ) |
| 557 | session.add(new_auth_key) |
| 558 | logger.info( |
| 559 | f"Added new available integration auth keys: " |
| 560 | f"{available_integration_auth_keys_data['auth_key_name']} for " |
| 561 | f"{available_integration_auth_keys_data['integration_name']}", |
| 562 | ) |
| 563 | except Exception as e: |
| 564 | logger.error(f"Error adding available integration auth keys: {e}") |
| 565 | raise e |
| 566 | await session.commit() |
| 567 | |
| 568 | |
| 569 | # ! AVAILABLE NETWORK CONNECTORS ! # |
| 570 | def load_available_network_connectors_data( |
| 571 | network_connector_name: str, |
| 572 | description: str, |
| 573 | network_connector_details: str, |
| 574 | ): |
| 575 | """ |
| 576 | Load available network_connectors data from environment variables. |
| 577 | |
| 578 | Args: |
| 579 | network_connector_name (str): The name of the network_connector. |
| 580 | description (str): The description of the network_connector. |
| 581 | |
| 582 | Returns: |
| 583 | dict: A dictionary containing the network_connector data. |
| 584 | """ |
| 585 | logger.info(f"Loading available network_connectors data for {network_connector_name}.") |
| 586 | return { |
| 587 | "network_connector_name": network_connector_name, |
| 588 | "description": description, |
| 589 | "network_connector_details": network_connector_details, |
| 590 | } |
| 591 | |
| 592 | |
| 593 | def load_markdown_for_network_connector(network_connector_name: str) -> str: |
| 594 | """ |
| 595 | Load markdown content for a given network_connector from a file. |
| 596 | |
| 597 | Args: |
| 598 | network_connector_name (str): The name of the network_connector. |
| 599 | |
| 600 | Returns: |
| 601 | str: The content of the markdown file. |
| 602 | """ |
| 603 | # file_path = os.path.join("network_connectors_markdown", f"{network_connector_name.lower()}.md") |
| 604 | # if space in the network_connector name, replace it with underscore |
| 605 | if " " in network_connector_name: |
| 606 | network_connector_name = network_connector_name.replace(" ", "_") |
| 607 | file_path = os.path.join( |
| 608 | "app", |
| 609 | "network_connectors", |
| 610 | "markdown", |
| 611 | f"{network_connector_name.lower()}.md", |
| 612 | ) |
| 613 | try: |
| 614 | with open(file_path, "r") as file: |
| 615 | return file.read() |
| 616 | except FileNotFoundError: |
| 617 | return "No deployment instructions available." |
| 618 | |
| 619 | |
| 620 | def get_available_network_connectors_list(): |
| 621 | """ |
| 622 | Get a list of available network_connectors. |
| 623 | |
| 624 | Returns: |
| 625 | list: A list of available network_connectors data, where each item contains the network_connector name, description, and markdown details. |
| 626 | """ |
| 627 | available_network_connectors = [ |
| 628 | ("Fortinet", "Integrate Fortinet with SOCFortress."), |
| 629 | ("Sonicwall", "Integrate Sonicwall with SOCFortress."), |
| 630 | ("Sentinelone", "Integrate Sentinelone with SOCFortress."), |
| 631 | # ... Add more available network_connectors as needed ... |
| 632 | ] |
| 633 | |
| 634 | return [ |
| 635 | load_available_network_connectors_data( |
| 636 | network_connector_name, |
| 637 | description, |
| 638 | load_markdown_for_network_connector(network_connector_name), |
| 639 | ) |
| 640 | for network_connector_name, description in available_network_connectors |
| 641 | ] |
| 642 | |
| 643 | |
| 644 | async def add_available_network_connectors_if_not_exist(session: AsyncSession): |
| 645 | """ |
| 646 | Adds available network_connectors to the database if they do not already exist. |
| 647 | |
| 648 | Args: |
| 649 | session (AsyncSession): The database session. |
| 650 | |
| 651 | Returns: |
| 652 | None |
| 653 | """ |
| 654 | available_network_connectors_list = get_available_network_connectors_list() |
| 655 | |
| 656 | for available_network_connector_data in available_network_connectors_list: |
| 657 | try: |
| 658 | query = select(AvailableNetworkConnectors).where( |
| 659 | AvailableNetworkConnectors.network_connector_name == available_network_connector_data["network_connector_name"], |
| 660 | ) |
| 661 | result = await session.execute(query) |
| 662 | existing_available_network_connector = result.scalars().first() |
| 663 | |
| 664 | if existing_available_network_connector is None: |
| 665 | new_available_network_connector = AvailableNetworkConnectors( |
| 666 | **available_network_connector_data, |
| 667 | ) |
| 668 | logger.info(f"New available network_connector: {available_network_connector_data}") |
| 669 | session.add(new_available_network_connector) |
| 670 | logger.info( |
| 671 | f"Added new available network_connector: {available_network_connector_data['network_connector_name']}", |
| 672 | ) |
| 673 | except Exception as e: |
| 674 | logger.error(f"Error adding available network_connector: {e}") |
| 675 | await session.rollback() |
| 676 | raise e |
| 677 | await session.commit() |
| 678 | # Close the session |
| 679 | await session.close() |
| 680 | |
| 681 | |
| 682 | def load_available_network_connectors_auth_keys( |
| 683 | network_connector_id: int, |
| 684 | network_connector_name: str, |
| 685 | auth_key_name: str, |
| 686 | ): |
| 687 | """ |
| 688 | Load available network_connectors auth keys from environment variables. |
| 689 | |
| 690 | Args: |
| 691 | network_connector_id (int): The ID of the network_connector. |
| 692 | network_connector_name (str): The name of the network_connector. |
| 693 | auth_key_name (str): The name of the auth key. |
| 694 | |
| 695 | Returns: |
| 696 | dict: A dictionary containing the auth key data. |
| 697 | """ |
| 698 | logger.info( |
| 699 | f"Loading available network_connectors auth keys data for {network_connector_name}.", |
| 700 | ) |
| 701 | return { |
| 702 | "network_connector_id": network_connector_id, |
| 703 | "network_connector_name": network_connector_name, |
| 704 | "auth_key_name": auth_key_name, |
| 705 | } |
| 706 | |
| 707 | |
| 708 | async def get_available_network_connectors_auth_keys_list(session: AsyncSession): |
| 709 | """ |
| 710 | Get a list of available network_connectors auth keys with their corresponding network_connector IDs. |
| 711 | |
| 712 | Args: |
| 713 | session (AsyncSession): The database session. |
| 714 | |
| 715 | Returns: |
| 716 | list: A list of available network_connectors auth keys data, where each item contains the network_connector ID, network_connector name, and auth key name. |
| 717 | """ |
| 718 | available_network_connectors_auth_keys = [] |
| 719 | available_network_connectors = [ |
| 720 | ("Fortinet", "SYSLOG_PORT"), |
| 721 | ("Sonicwall", "SYSLOG_PORT"), |
| 722 | ("Sonicwall", "TLS_CERT_FILE"), |
| 723 | ("Sonicwall", "TLS_KEY_FILE"), |
| 724 | ("Sentinelone", "SYSLOG_PORT"), |
| 725 | ("Sentinelone", "TLS_CERT_FILE"), |
| 726 | ("Sentinelone", "TLS_KEY_FILE"), |
| 727 | # ... Add more available network_connectors auth keys as needed ... |
| 728 | ] |
| 729 | logger.info("Getting available network_connectors auth keys.") |
| 730 | try: |
| 731 | for network_connector_name, auth_key_name in available_network_connectors: |
| 732 | query = select(AvailableNetworkConnectors.id).where( |
| 733 | AvailableNetworkConnectors.network_connector_name == network_connector_name, |
| 734 | ) |
| 735 | result = await session.execute(query) |
| 736 | network_connector_id = result.scalars().first() |
| 737 | logger.info(f"Network Connector ID for {network_connector_name}: {network_connector_id}") |
| 738 | if network_connector_id: |
| 739 | logger.info(f"Found network_connector ID for {network_connector_name}: {network_connector_id}") |
| 740 | available_network_connectors_auth_keys.append( |
| 741 | load_available_network_connectors_auth_keys( |
| 742 | network_connector_id, |
| 743 | network_connector_name, |
| 744 | auth_key_name, |
| 745 | ), |
| 746 | ) |
| 747 | |
| 748 | return available_network_connectors_auth_keys |
| 749 | except Exception as e: |
| 750 | logger.error(f"Error getting available network_connectors auth keys: {e}") |
| 751 | await session.rollback() |
| 752 | raise e |
| 753 | |
| 754 | |
| 755 | async def add_available_network_connectors_auth_keys_if_not_exist(session: AsyncSession): |
| 756 | """ |
| 757 | Adds available network_connectors auth keys to the database if they do not already exist. |
| 758 | |
| 759 | Args: |
| 760 | session (AsyncSession): The database session. |
| 761 | |
| 762 | Returns: |
| 763 | None |
| 764 | """ |
| 765 | logger.info("Checking for existence of available network_connectors auth keys.") |
| 766 | available_network_connectors_auth_keys_list = await get_available_network_connectors_auth_keys_list(session=session) |
| 767 | logger.info("Adding available network_connectors auth keys to the database.") |
| 768 | for available_network_connector_auth_keys_data in available_network_connectors_auth_keys_list: |
| 769 | try: |
| 770 | query = select(AvailableNetworkConnectors).where( |
| 771 | AvailableNetworkConnectors.network_connector_name == available_network_connector_auth_keys_data["network_connector_name"], |
| 772 | ) |
| 773 | result = await session.execute(query) |
| 774 | existing_network_connector = result.scalars().first() |
| 775 | |
| 776 | if existing_network_connector: |
| 777 | available_network_connector_auth_keys_data["network_connector_id"] = existing_network_connector.id |
| 778 | auth_key_query = select(AvailableNetworkConnectorsKeys).where( |
| 779 | and_( |
| 780 | AvailableNetworkConnectorsKeys.network_connector_id == existing_network_connector.id, |
| 781 | AvailableNetworkConnectorsKeys.auth_key_name == available_network_connector_auth_keys_data["auth_key_name"], |
| 782 | ), |
| 783 | ) |
| 784 | auth_key_result = await session.execute(auth_key_query) |
| 785 | existing_auth_key = auth_key_result.scalars().first() |
| 786 | |
| 787 | if existing_auth_key is None: |
| 788 | new_auth_key = AvailableNetworkConnectorsKeys( |
| 789 | **available_network_connector_auth_keys_data, |
| 790 | ) |
| 791 | session.add(new_auth_key) |
| 792 | logger.info( |
| 793 | f"Added new available network_connector auth keys: " |
| 794 | f"{available_network_connector_auth_keys_data['auth_key_name']} for " |
| 795 | f"{available_network_connector_auth_keys_data['network_connector_name']}", |
| 796 | ) |
| 797 | except Exception as e: |
| 798 | logger.error(f"Error adding available network_connector auth keys: {e}") |
| 799 | raise e |
| 800 | await session.commit() |