feat(azure): add Databricks service and check for workspace VNet injection (#8008)

Co-authored-by: MrCloudSec <hello@mistercloudsec.com>
2026-07-24 13:01:56 +00:00 · 2025-06-18 11:38:09 +02:00
parent 7de7122c3b
commit 76f0d890e9
11 changed files with 336 additions and 2 deletions
@@ -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"
@@ -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)

 ---

@@ -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())
@@ -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]
@@ -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 <databricks-workspace-name> --resource-group <resource-group-name> --location <region> --managed-resource-group <managed-rg-name> --enable-no-public-ip true --network-security-group-rule \"NoAzureServices\" --public-network-access Disabled --custom-virtual-network-id /subscriptions/<subscription-id>/resourceGroups/<resource-group-name>/providers/Microsoft.Network/virtualNetworks/<vnet-name>",
+      "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": ""
+}
@@ -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
@@ -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",
@@ -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"
@@ -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"