diff --git a/poetry.lock b/poetry.lock index 9eea3c056f..58abe3af1c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.2 and should not be changed by hand. [[package]] name = "about-time" @@ -574,6 +574,23 @@ azure-mgmt-core = ">=1.3.2" isodate = ">=0.6.1" typing-extensions = ">=4.6.0" +[[package]] +name = "azure-mgmt-databricks" +version = "2.0.0" +description = "Microsoft Azure Data Bricks Management Client Library for Python" +optional = false +python-versions = ">=3.7" +groups = ["main"] +files = [ + {file = "azure-mgmt-databricks-2.0.0.zip", hash = "sha256:70d11362dc2d17f5fb1db0cfe65c1af55b8f136f1a0db9a5b51e7acf760cf5b9"}, + {file = "azure_mgmt_databricks-2.0.0-py3-none-any.whl", hash = "sha256:0c29434a7339e74231bd171a6c08dcdf8153abaebd332658d7f66b8ea143fa17"}, +] + +[package.dependencies] +azure-common = ">=1.1,<2.0" +azure-mgmt-core = ">=1.3.2,<2.0.0" +isodate = ">=0.6.1,<1.0.0" + [[package]] name = "azure-mgmt-keyvault" version = "10.3.1" @@ -6586,4 +6603,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">3.9.1,<3.13" -content-hash = "e88f0c925f435849bf13d3a38978f71213033c2407708dcf3007fa0f79bc29a0" +content-hash = "d72c55b52949ba94f0c68004d5b778edb69514a05bbb7aba8d641b5058a99fd5" diff --git a/prowler/CHANGELOG.md b/prowler/CHANGELOG.md index 393a1d50d2..7db3191276 100644 --- a/prowler/CHANGELOG.md +++ b/prowler/CHANGELOG.md @@ -34,6 +34,7 @@ All notable changes to the **Prowler SDK** are documented in this file. - ISO 27001 2022 for M365 provider. [(#7985)](https://github.com/prowler-cloud/prowler/pull/7985) - `codebuild_project_uses_allowed_github_organizations` check for AWS provider [(#7595)](https://github.com/prowler-cloud/prowler/pull/7595) - IaC provider [(#7852)](https://github.com/prowler-cloud/prowler/pull/7852) +- Azure Databricks service integration for Azure provider, including the `databricks_workspace_vnet_injection_enabled` check [(#8008)](https://github.com/prowler-cloud/prowler/pull/8008) --- diff --git a/prowler/providers/azure/services/databricks/__init__.py b/prowler/providers/azure/services/databricks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/prowler/providers/azure/services/databricks/databricks_client.py b/prowler/providers/azure/services/databricks/databricks_client.py new file mode 100644 index 0000000000..178db476fa --- /dev/null +++ b/prowler/providers/azure/services/databricks/databricks_client.py @@ -0,0 +1,4 @@ +from prowler.providers.azure.services.databricks.databricks_service import Databricks +from prowler.providers.common.provider import Provider + +databricks_client = Databricks(Provider.get_global_provider()) diff --git a/prowler/providers/azure/services/databricks/databricks_service.py b/prowler/providers/azure/services/databricks/databricks_service.py new file mode 100644 index 0000000000..1f35eb74ff --- /dev/null +++ b/prowler/providers/azure/services/databricks/databricks_service.py @@ -0,0 +1,80 @@ +from typing import Optional + +from azure.mgmt.databricks import AzureDatabricksManagementClient +from pydantic import BaseModel + +from prowler.lib.logger import logger +from prowler.providers.azure.azure_provider import AzureProvider +from prowler.providers.azure.lib.service.service import AzureService + + +class Databricks(AzureService): + """ + Service class for interacting with Azure Databricks workspaces. + + This class initializes the Azure Databricks Management Client for each subscription + and retrieves all Databricks workspaces within those subscriptions. + """ + + def __init__(self, provider: AzureProvider): + """ + Initialize the Databricks service with the given Azure provider. + + Args: + provider: The Azure provider instance containing credentials and configuration. + """ + super().__init__(AzureDatabricksManagementClient, provider) + self.workspaces = self._get_workspaces() + + def _get_workspaces(self) -> dict: + """ + Retrieve all Databricks workspaces for each subscription. + + Returns: + A dictionary mapping subscription IDs to their Databricks workspaces. + """ + logger.info("Databricks - Getting workspaces...") + workspaces = {} + for subscription, client in self.clients.items(): + try: + workspaces[subscription] = {} + + for workspace in client.workspaces.list_by_subscription(): + workspace_parameters = getattr(workspace, "parameters", None) + + workspaces[subscription][workspace.id] = DatabricksWorkspace( + id=workspace.id, + name=workspace.name, + location=workspace.location, + custom_managed_vnet_id=( + getattr( + workspace_parameters, "custom_virtual_network_id", None + ).value + if getattr( + workspace_parameters, "custom_virtual_network_id", None + ) + else None + ), + ) + except Exception as error: + logger.error( + f"Subscription: {subscription} -- {error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}" + ) + return workspaces + + +class DatabricksWorkspace(BaseModel): + """ + Pydantic model representing an Azure Databricks workspace. + + Attributes: + id: The unique identifier of the workspace. + name: The name of the workspace. + location: The Azure region where the workspace is deployed. + custom_managed_vnet_id: The ID of the custom managed virtual network, if configured. + """ + + id: str + name: str + location: str + custom_managed_vnet_id: Optional[str] diff --git a/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/__init__.py b/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.metadata.json b/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.metadata.json new file mode 100644 index 0000000000..d66e789e7d --- /dev/null +++ b/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.metadata.json @@ -0,0 +1,30 @@ +{ + "Provider": "azure", + "CheckID": "databricks_workspace_vnet_injection_enabled", + "CheckTitle": "Ensure Azure Databricks workspaces are deployed in a customer-managed VNet (VNet Injection)", + "CheckType": [], + "ServiceName": "databricks", + "SubServiceName": "", + "ResourceIdTemplate": "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Databricks/workspaces/{workspaceName}", + "Severity": "medium", + "ResourceType": "AzureDatabricksWorkspace", + "Description": "Checks whether Azure Databricks workspaces are deployed in a customer-managed Virtual Network (VNet Injection) instead of a Databricks-managed VNet.", + "Risk": "Using a Databricks-managed VNet limits control over network security policies, firewall configurations, and routing, increasing the risk of unauthorized access or data exfiltration.", + "RelatedUrl": "https://learn.microsoft.com/en-us/azure/databricks/administration-guide/cloud-configurations/azure/vnet-inject", + "Remediation": { + "Code": { + "CLI": "az databricks workspace create --name --resource-group --location --managed-resource-group --enable-no-public-ip true --network-security-group-rule \"NoAzureServices\" --public-network-access Disabled --custom-virtual-network-id /subscriptions//resourceGroups//providers/Microsoft.Network/virtualNetworks/", + "NativeIaC": "", + "Other": "", + "Terraform": "" + }, + "Recommendation": { + "Text": "Deploy Databricks workspaces into a customer-managed VNet to ensure better control over network security and compliance.", + "Url": "https://www.trendmicro.com/cloudoneconformity/knowledge-base/azure/Databricks/check-for-vnet-injection.html" + } + }, + "Categories": [], + "DependsOn": [], + "RelatedTo": [], + "Notes": "" +} diff --git a/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.py b/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.py new file mode 100644 index 0000000000..f667342dab --- /dev/null +++ b/prowler/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled.py @@ -0,0 +1,32 @@ +from prowler.lib.check.models import Check, Check_Report_Azure +from prowler.providers.azure.services.databricks.databricks_client import ( + databricks_client, +) + + +class databricks_workspace_vnet_injection_enabled(Check): + """ + Ensure Azure Databricks workspaces are deployed in a customer-managed VNet (VNet Injection). + + This check evaluates whether each Azure Databricks workspace in the subscription is configured to use VNet Injection, meaning it is deployed in a customer-managed virtual network (VNet). + + - PASS: The workspace is deployed in a customer-managed VNet (custom_managed_vnet_id is set). + - FAIL: The workspace is not deployed in a customer-managed VNet (VNet Injection is not enabled). + """ + + def execute(self): + findings = [] + for subscription, workspaces in databricks_client.workspaces.items(): + for workspace in workspaces.values(): + report = Check_Report_Azure( + metadata=self.metadata(), resource=workspace + ) + report.subscription = subscription + if workspace.custom_managed_vnet_id: + report.status = "PASS" + report.status_extended = f"Databricks workspace {workspace.name} in subscription {subscription} is deployed in a customer-managed VNet ({workspace.custom_managed_vnet_id})." + else: + report.status = "FAIL" + report.status_extended = f"Databricks workspace {workspace.name} in subscription {subscription} is not deployed in a customer-managed VNet (VNet Injection is not enabled)." + findings.append(report) + return findings diff --git a/pyproject.toml b/pyproject.toml index f78af76f10..401e88d341 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "azure-mgmt-containerregistry==12.0.0", "azure-mgmt-containerservice==34.1.0", "azure-mgmt-cosmosdb==9.7.0", + "azure-mgmt-databricks==2.0.0", "azure-mgmt-keyvault==10.3.1", "azure-mgmt-monitor==6.0.2", "azure-mgmt-network==28.1.0", diff --git a/tests/providers/azure/services/databricks/databricks_service_test.py b/tests/providers/azure/services/databricks/databricks_service_test.py new file mode 100644 index 0000000000..df4cf920f0 --- /dev/null +++ b/tests/providers/azure/services/databricks/databricks_service_test.py @@ -0,0 +1,50 @@ +from unittest.mock import patch + +from prowler.providers.azure.services.databricks.databricks_service import ( + Databricks, + DatabricksWorkspace, +) +from tests.providers.azure.azure_fixtures import ( + AZURE_SUBSCRIPTION_ID, + set_mocked_azure_provider, +) + + +def mock_databricks_get_workspaces(_): + return { + AZURE_SUBSCRIPTION_ID: { + "test-workspace-id": DatabricksWorkspace( + id="test-workspace-id", + name="test-workspace", + location="eastus", + custom_managed_vnet_id="test-vnet-id", + ) + } + } + + +@patch( + "prowler.providers.azure.services.databricks.databricks_service.Databricks._get_workspaces", + new=mock_databricks_get_workspaces, +) +class Test_Databricks_Service: + def test_get_client(self): + databricks = Databricks(set_mocked_azure_provider()) + assert ( + databricks.clients[AZURE_SUBSCRIPTION_ID].__class__.__name__ + == "AzureDatabricksManagementClient" + ) + + def test_get_workspaces(self): + databricks = Databricks(set_mocked_azure_provider()) + assert ( + databricks.workspaces[AZURE_SUBSCRIPTION_ID][ + "test-workspace-id" + ].__class__.__name__ + == "DatabricksWorkspace" + ) + workspace = databricks.workspaces[AZURE_SUBSCRIPTION_ID]["test-workspace-id"] + assert workspace.id == "test-workspace-id" + assert workspace.name == "test-workspace" + assert workspace.location == "eastus" + assert workspace.custom_managed_vnet_id == "test-vnet-id" diff --git a/tests/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled_test.py b/tests/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled_test.py new file mode 100644 index 0000000000..912ee363ac --- /dev/null +++ b/tests/providers/azure/services/databricks/databricks_workspace_vnet_injection_enabled/databricks_workspace_vnet_injection_enabled_test.py @@ -0,0 +1,119 @@ +from unittest import mock +from uuid import uuid4 + +from prowler.providers.azure.services.databricks.databricks_service import ( + DatabricksWorkspace, +) +from tests.providers.azure.azure_fixtures import ( + AZURE_SUBSCRIPTION_ID, + set_mocked_azure_provider, +) + + +class Test_databricks_workspace_vnet_injection_enabled: + def test_databricks_no_workspaces(self): + databricks_client = mock.MagicMock + databricks_client.workspaces = {} + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=set_mocked_azure_provider(), + ), + mock.patch( + "prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled.databricks_client", + new=databricks_client, + ), + ): + from prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled import ( + databricks_workspace_vnet_injection_enabled, + ) + + check = databricks_workspace_vnet_injection_enabled() + result = check.execute() + assert len(result) == 0 + + def test_databricks_workspace_vnet_injection_disabled(self): + workspace_id = str(uuid4()) + workspace_name = "test-workspace" + databricks_client = mock.MagicMock + databricks_client.workspaces = { + AZURE_SUBSCRIPTION_ID: { + workspace_id: DatabricksWorkspace( + id=workspace_id, + name=workspace_name, + location="eastus", + custom_managed_vnet_id=None, + ) + } + } + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=set_mocked_azure_provider(), + ), + mock.patch( + "prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled.databricks_client", + new=databricks_client, + ), + ): + from prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled import ( + databricks_workspace_vnet_injection_enabled, + ) + + check = databricks_workspace_vnet_injection_enabled() + result = check.execute() + assert len(result) == 1 + assert result[0].status == "FAIL" + assert ( + result[0].status_extended + == f"Databricks workspace {workspace_name} in subscription {AZURE_SUBSCRIPTION_ID} is not deployed in a customer-managed VNet (VNet Injection is not enabled)." + ) + assert result[0].subscription == AZURE_SUBSCRIPTION_ID + assert result[0].resource_name == workspace_name + assert result[0].resource_id == workspace_id + assert result[0].location == "eastus" + + def test_databricks_workspace_vnet_injection_enabled(self): + workspace_id = str(uuid4()) + workspace_name = "test-workspace" + vnet_id = "test-vnet-id" + databricks_client = mock.MagicMock + databricks_client.workspaces = { + AZURE_SUBSCRIPTION_ID: { + workspace_id: DatabricksWorkspace( + id=workspace_id, + name=workspace_name, + location="eastus", + custom_managed_vnet_id=vnet_id, + ) + } + } + + with ( + mock.patch( + "prowler.providers.common.provider.Provider.get_global_provider", + return_value=set_mocked_azure_provider(), + ), + mock.patch( + "prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled.databricks_client", + new=databricks_client, + ), + ): + from prowler.providers.azure.services.databricks.databricks_workspace_vnet_injection_enabled.databricks_workspace_vnet_injection_enabled import ( + databricks_workspace_vnet_injection_enabled, + ) + + check = databricks_workspace_vnet_injection_enabled() + result = check.execute() + assert len(result) == 1 + assert result[0].status == "PASS" + assert ( + result[0].status_extended + == f"Databricks workspace {workspace_name} in subscription {AZURE_SUBSCRIPTION_ID} is deployed in a customer-managed VNet ({vnet_id})." + ) + assert result[0].subscription == AZURE_SUBSCRIPTION_ID + assert result[0].resource_name == workspace_name + assert result[0].resource_id == workspace_id + assert result[0].location == "eastus"