feature(azure): Add Azure Key Management Service (KMS) integration

  - Add Azure KMS configuration support to provisioner and VM provider
  - Pre-create Azure KMS keys per region and reuse SCT configuration
    option 'enterprise_disable_kms' to enable EaR with managed identity.
  - Add SCT configuration option 'enable_kms_key_rotation' for key rotation.

Signed-off-by: Lakshmipathi.Ganapathi <[email protected]>

Author: Lakshmipathi

docs/configuration_options.md

@@ -2834,13 +2834,22 @@ options will be used for enable encryption at-rest for tables
 
 ## **kms_key_rotation_interval** / SCT_KMS_KEY_ROTATION_INTERVAL
 
-The time interval in minutes which gets waited before the KMS key rotation happens. Applied when the AWS KMS service is configured to be used.
+The time interval in minutes which gets waited before the KMS key rotation happens. Applied when AWS KMS or Azure KMS service is configured to be used. NOTE: Be aware that Azure Key rotations cost $1/rotation.
 
 **default:** N/A
 
 **type:** int
 
 
+## **enable_kms_key_rotation** / SCT_ENABLE_KMS_KEY_ROTATION
+
+Allows to disable KMS keys rotation. Applicable only to Azure backend. In case of AWS backend its KMS keys will always be rotated as of now.
+
+**default:** N/A
+
+**type:** boolean
+
+
 ## **enterprise_disable_kms** / SCT_ENTERPRISE_DISABLE_KMS
 
 An escape hatch to disable KMS for enterprise run, when needed, we enable kms by default since if we use scylla 2023.1.3 and up

sdcm/cluster.py

@@ -87,6 +87,9 @@
 from sdcm.utils import properties
 from sdcm.utils.adaptive_timeouts import Operations, adaptive_timeout
 from sdcm.utils.aws_kms import AwsKms
+from sdcm.utils.azure_utils import AzureService
+from sdcm.provision.azure.kms_provider import AzureKmsProvider
+from azure.core.exceptions import ResourceNotFoundError as AzureResourceNotFoundError
 from sdcm.utils.cql_utils import cql_quote_if_needed
 from sdcm.utils.benchmarks import ScyllaClusterBenchmarkManager
 from sdcm.utils.common import (
@@ -4790,6 +4793,31 @@ def _rotate_kms_key(kms_key_alias_name, kms_key_rotation_interval, db_cluster):
         kms_key_rotation_thread.start()
         return None
 
+    def start_azure_kms_key_rotation_thread(self) -> None:
+        if self.params.get("cluster_backend") != 'azure':
+            return None
+        if not self.params.get("enable_kms_key_rotation"):
+            return None
+
+        test_id = str(self.test_config.test_id())
+        region = self.params.get('azure_region_name')[0]
+
+        def _rotate():
+            azure_service = AzureService()
+
+            while True:
+                time.sleep(self.params.get("kms_key_rotation_interval") * 60)
+                try:
+                    key_uri = AzureKmsProvider.get_key_uri_for_test(region, test_id)
+                    rotated_key_id = azure_service.rotate_vault_key(key_uri)
+                    self.log.info(f"Azure KMS key rotated for test {test_id}: {rotated_key_id}")
+                except AzureResourceNotFoundError as e:
+                    self.log.error(f"Azure KMS key not found for rotation: {e}")
+
+        threading.Thread(target=_rotate, daemon=True, name='AzureKmsRotationThread').start()
+        self.log.info("Started Azure KMS rotation thread for test: %s", test_id)
+        return None
+
     def scylla_configure_non_root_installation(self, node, devname):
         node.stop_scylla_server(verify_down=False)
         node.remoter.run(f'{node.offline_install_dir}/sbin/scylla_setup --nic {devname} --no-raid-setup',

sdcm/keystore.py

@@ -114,6 +114,9 @@ def get_docker_hub_credentials(self):
     def get_azure_credentials(self):
         return self.get_json("azure.json")
 
+    def get_azure_kms_config(self):
+        return self.get_json("azure_kms_config.json")
+
     def get_argusdb_credentials(self):
         return self.get_json("argusdb_config_v2.json")
 

sdcm/provision/azure/kms_provider.py

@@ -0,0 +1,124 @@
+
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation; either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+#
+# See LICENSE for more details.
+#
+# Copyright (c) 2025 ScyllaDB
+import logging
+from dataclasses import dataclass
+
+from azure.core.exceptions import AzureError
+from sdcm.utils.azure_utils import AzureService
+from sdcm.keystore import KeyStore
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class AzureKmsProvider:
+    _resource_group_name: str
+    _region: str
+    _az: str
+    _azure_service: AzureService = AzureService()
+
+    def __post_init__(self):
+        self._kms_config = KeyStore().get_azure_kms_config()
+
+    @property
+    def managed_identity_config(self):
+        return {
+            'resource_group': self._kms_config['resource_group'],
+            'identity_name': self._kms_config['identity_name'],
+            'principal_id': self._kms_config['managed_identity_principal_id']
+        }
+
+    @property
+    def sct_service_principal_id(self):
+        return self._kms_config['sct_service_principal_id']
+
+    @classmethod
+    def _get_vault_name(cls, region: str) -> str:
+        """Generate vault name for the given region"""
+        kms_config = KeyStore().get_azure_kms_config()
+        return f"{kms_config['shared_vault_name']}-{region}"
+
+    def _get_managed_identity_id(self) -> str:
+        return (
+            f"/subscriptions/{self._azure_service.subscription_id}"
+            f"/resourcegroups/{self.managed_identity_config['resource_group']}"
+            "/providers/Microsoft.ManagedIdentity"
+            f"/userAssignedIdentities/{self.managed_identity_config['identity_name']}"
+        )
+
+    @classmethod
+    def get_key_uri_for_test(cls, region: str, test_id: str) -> str:
+        vault_name = cls._get_vault_name(region)
+        vault_uri = f"https://{vault_name}.vault.azure.net/"
+        kms_config = KeyStore().get_azure_kms_config()
+        num_of_keys = kms_config['num_of_keys']
+        key_number = (hash(test_id) % num_of_keys) + 1
+        return f"{vault_uri}scylla-key-{key_number}"
+
+    def get_or_create_keyvault_and_identity(self, test_id: str):
+        """Use fixed vault with keys"""
+        vault_name = self._get_vault_name(self._region)
+        try:
+            vault = self._azure_service.keyvault.vaults.begin_create_or_update(
+                resource_group_name=self._kms_config['resource_group'], vault_name=vault_name,
+                parameters={
+                    "location": self._region,
+                    "properties": {
+                        "tenant_id": self._azure_service.azure_credentials["tenant_id"],
+                        "sku": {"name": "standard", "family": "A"},
+                        "enabled_for_disk_encryption": True,
+                        "enable_rbac_authorization": False,
+                        "access_policies": [{
+                            "tenant_id": self._azure_service.azure_credentials["tenant_id"],
+                            "object_id": self.managed_identity_config['principal_id'],
+                            "permissions": {
+                                "keys": ["get", "encrypt", "decrypt", "wrapKey", "unwrapKey"],
+                                "secrets": ["get"],
+                                "certificates": ["get"]
+                            }
+                        }, {
+                            # SCT service principal
+                            "tenant_id": self._azure_service.azure_credentials["tenant_id"],
+                            "object_id": self.sct_service_principal_id,
+                            "permissions": {
+                                "keys": ["create", "get", "list", "update", "import", "delete", "rotate"],
+                                "secrets": ["get"],
+                                "certificates": ["get"]
+                            }
+                        }],
+                    }
+                }
+            ).result()
+
+            vault_uri = vault.properties.vault_uri
+
+            # Pick one key, if required create keys.
+            num_of_keys = self._kms_config['num_of_keys']
+            for i in range(1, num_of_keys + 1):
+                key_name = f"scylla-key-{i}"
+                if not self._azure_service.get_vault_key(vault_uri, key_name):
+                    self._azure_service.create_vault_key(vault_uri, key_name)
+                    LOGGER.info(f"Created key: {key_name}")
+
+            key_number = (hash(test_id) % num_of_keys) + 1
+            key_uri = f"{vault_uri}scylla-key-{key_number}"
+            vault_info = {
+                'identity_id': self._get_managed_identity_id(),
+                'vault_uri': vault_uri,
+                'key_uri': key_uri
+            }
+            return vault_info
+        except AzureError as e:
+            LOGGER.warning(f"Failed to setup Azure KMS: {e}")
+            return None

sdcm/provision/azure/provisioner.py

@@ -38,9 +38,11 @@ class AzureProvisioner(Provisioner):
     """Provides api for VM provisioning in Azure cloud, tuned for Scylla QA. """
 
     def __init__(self, test_id: str, region: str, availability_zone: str,
-                 azure_service: AzureService = AzureService(), **_):
+                 azure_service: AzureService = AzureService(), **config):
         availability_zone = self._convert_az_to_zone(availability_zone)
         super().__init__(test_id, region, availability_zone)
+        # NOTE: disable Azure KMS if not configured explicitly
+        self._enable_azure_kms = not config.get('enterprise_disable_kms') in (True, None)
         self._azure_service: AzureService = azure_service
         self._cache: Dict[str, VmInstance] = {}
         LOGGER.debug("getting resources for %s...", self._resource_group_name)
@@ -54,7 +56,7 @@ def __init__(self, test_id: str, region: str, availability_zone: str,
         self._ip_provider = IpAddressProvider(self._resource_group_name, self._region, self._az, self._azure_service)
         self._nic_provider = NetworkInterfaceProvider(self._resource_group_name, self._region, self._azure_service)
         self._vm_provider = VirtualMachineProvider(
-            self._resource_group_name, self._region, self._az, self._azure_service)
+            self._resource_group_name, self._region, self._az, self._enable_azure_kms, self._azure_service)
         for v_m in self._vm_provider.list():
             try:
                 self._cache[v_m.name] = self._vm_to_instance(v_m)
@@ -145,7 +147,7 @@ def get_or_create_instances(self,
                 self._resource_group_name, self._region, self._az, self._azure_service)
             self._nic_provider = NetworkInterfaceProvider(self._resource_group_name, self._region, self._azure_service)
             self._vm_provider = VirtualMachineProvider(
-                self._resource_group_name, self._region, self._az, self._azure_service)
+                self._resource_group_name, self._region, self._az, self._enable_azure_kms, self._azure_service)
             raise
         for definition, v_m in zip(definitions, v_ms):
             instance = self._vm_to_instance(v_m)

sdcm/provision/azure/virtual_machine_provider.py

@@ -24,18 +24,22 @@
 from azure.mgmt.compute.models import VirtualMachine, RunCommandInput
 from invoke import Result
 
+from sdcm.provision.azure.kms_provider import AzureKmsProvider
 from sdcm.provision.provisioner import InstanceDefinition, PricingModel, ProvisionError, OperationPreemptedError
 from sdcm.provision.user_data import UserDataBuilder
 from sdcm.utils.azure_utils import AzureService
 
 LOGGER = logging.getLogger(__name__)
 
+SCT_RESOURCE_GROUP_PREFIX = "SCT-"
+
 
 @dataclass
 class VirtualMachineProvider:
     _resource_group_name: str
     _region: str
     _az: str
+    _enable_azure_kms: bool = False
     _azure_service: AzureService = AzureService()
     _cache: Dict[str, VirtualMachine] = field(default_factory=dict)
 
@@ -80,6 +84,19 @@ def get_or_create(self, definitions: List[InstanceDefinition], nics_ids: List[st
                     }],
                 },
             }
+
+            if self._enable_azure_kms:
+                self._kms_provider = AzureKmsProvider(
+                    self._resource_group_name, self._region, self._az, self._azure_service)
+                # Extract test_id from resource group name
+                test_id = self._resource_group_name.split(SCT_RESOURCE_GROUP_PREFIX)[-1][:36]
+                vault_info = self._kms_provider.get_or_create_keyvault_and_identity(test_id)
+                params["identity"] = {"type": "UserAssigned",
+                                      "user_assigned_identities": {vault_info['identity_id']: {}}}
+                LOGGER.info(f"Azure Key Vault enabled for {definition.name}")
+            else:
+                LOGGER.info(f"Azure Key Vault disabled for {definition.name}")
+
             if definition.user_data is None:
                 # in case we use specialized image, we don't change things like computer_name, usernames, ssh_keys
                 os_profile = {}

sdcm/sct_config.py

@@ -1440,7 +1440,14 @@ class SCTConfiguration(dict):
 
         dict(name="kms_key_rotation_interval", env="SCT_KMS_KEY_ROTATION_INTERVAL", type=int,
              help="The time interval in minutes which gets waited before the KMS key rotation happens."
-                  " Applied when the AWS KMS service is configured to be used."),
+                  " Applied when AWS KMS or Azure KMS service is configured to be used."
+                  " NOTE: Be aware that Azure Key rotations cost $1/rotation."),
+
+        # TODO: AWS KMS needs to support the enable_kms_key_rotation config option
+
+        dict(name="enable_kms_key_rotation", env="SCT_ENABLE_KMS_KEY_ROTATION", type=boolean,
+             help="Allows to disable KMS keys rotation. Applicable only to Azure backend. "
+                  "In case of AWS backend its KMS keys will always be rotated as of now."),
 
         dict(name="enterprise_disable_kms", env="SCT_ENTERPRISE_DISABLE_KMS", type=boolean,
              help="An escape hatch to disable KMS for enterprise run, when needed, "

sdcm/tester.py

@@ -49,6 +49,7 @@
 from argus.common.enums import TestStatus
 from sdcm import nemesis, cluster_docker, cluster_k8s, cluster_baremetal, db_stats, wait
 from sdcm.cloud_api_client import ScyllaCloudAPIClient
+from sdcm.provision.azure.kms_provider import AzureKmsProvider
 from sdcm.cluster import BaseCluster, NoMonitorSet, SCYLLA_DIR, TestConfig, UserRemoteCredentials, BaseLoaderSet, BaseMonitorSet, \
     BaseScyllaCluster, BaseNode, MINUTE_IN_SEC
 from sdcm.cluster_azure import ScyllaAzureCluster, LoaderSetAzure, MonitorSetAzure
@@ -885,6 +886,56 @@ def prepare_kms_host(self) -> None:
         self.params["append_scylla_yaml"] = append_scylla_yaml
         return None
 
+    def prepare_azure_kms(self) -> None:
+        scylla_version = self.params.scylla_version
+        if not scylla_version:
+            return None
+        version_supports_kms = ComparableScyllaVersion(scylla_version) >= '2025.4.0~dev'
+        backend_support_kms = self.params.get('cluster_backend') in ('azure',)
+        kms_configured_in_sct = self.params.get('scylla_encryption_options')
+        test_uses_oracle = self.params.get("db_type") == "mixed_scylla"
+        should_enable_kms = (version_supports_kms and
+                             backend_support_kms and
+                             not kms_configured_in_sct and
+                             not test_uses_oracle and
+                             not self.params.get('enterprise_disable_kms'))
+
+        if should_enable_kms:
+            self.params['scylla_encryption_options'] = "{ 'cipher_algorithm' : 'AES/ECB/PKCS5Padding', 'secret_key_strength' : 128, 'key_provider': 'AzureKeyProviderFactory', 'azure_host': 'scylla-azure-kms'}"
+        if not (scylla_encryption_options := self.params.get("scylla_encryption_options") or ''):
+            return None
+        azure_host = (yaml.safe_load(scylla_encryption_options) or {}).get("azure_host") or ''
+        if not azure_host:
+            return None
+
+        test_id = str(self.test_config.test_id())
+        append_scylla_yaml = self.params.get("append_scylla_yaml") or {}
+        if "azure_hosts" not in append_scylla_yaml:
+            append_scylla_yaml["azure_hosts"] = {}
+
+        regions = self.params.get('azure_region_name')
+        if len(regions) > 1:
+            raise NotImplementedError("Azure KMS multi-dc support is not yet implemented")
+        key_uri = AzureKmsProvider.get_key_uri_for_test(regions[0], test_id)
+
+        append_scylla_yaml["azure_hosts"][azure_host] = {
+            'master_key': key_uri
+        }
+
+        append_scylla_yaml['user_info_encryption'] = {
+            'enabled': True,
+            'key_provider': 'AzureKeyProviderFactory',
+            'azure_host': azure_host,
+        }
+        append_scylla_yaml['system_info_encryption'] = {
+            'enabled': True,
+            'key_provider': 'AzureKeyProviderFactory',
+            'azure_host': azure_host,
+        }
+
+        self.params["append_scylla_yaml"] = append_scylla_yaml
+        return None
+
     def kafka_configure(self):
         if self.kafka_cluster:
             for connector_config in self.params.get('kafka_connectors'):
@@ -1014,6 +1065,7 @@ def setUp(self):  # noqa: PLR0912, PLR0915
         if self.is_encrypt_keys_needed:
             self.download_encrypt_keys()
         self.prepare_kms_host()
+        self.prepare_azure_kms()
 
         self.nemesis_allocator = NemesisNodeAllocator(self)
 
@@ -1108,6 +1160,7 @@ def _db_post_validation():
             for db_cluster in self.db_clusters_multitenant:
                 if db_cluster:
                     db_cluster.start_kms_key_rotation_thread()
+                    db_cluster.start_azure_kms_key_rotation_thread()
 
             for future in as_completed(futures):
                 future.result()