[KMS] Implement PMK-aware validation for CMK

src/aks-preview/azext_aks_preview/_validators.py

@@ -31,7 +31,7 @@
                                        RequiredArgumentMissingError)
 from azure.cli.core.commands.validators import validate_tag
 from azure.cli.core.util import CLIError
-from azure.mgmt.core.tools import is_valid_resource_id
+from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
 from knack.log import get_logger
 
 logger = get_logger(__name__)
@@ -701,26 +701,94 @@ def validate_crg_id(namespace):
 def validate_azure_keyvault_kms_key_id(namespace):
     key_id = namespace.azure_keyvault_kms_key_id
     if key_id:
-        err_msg = (
-            "--azure-keyvault-kms-key-id is not a valid Key Vault key ID. "
-            "See https://docs.microsoft.com/en-us/azure/key-vault/general/about-keys-secrets-certificates#vault-name-and-object-name"  # pylint: disable=line-too-long
+        # Check if PMK (Platform-Managed Keys) is enabled
+        is_pmk_enabled = (
+            hasattr(namespace, 'kms_infrastructure_encryption') and
+            namespace.kms_infrastructure_encryption == "Enabled"
         )
 
         https_prefix = "https://"
         if not key_id.startswith(https_prefix):
+            err_msg = (
+                "--azure-keyvault-kms-key-id is not a valid Key Vault key ID. "
+                "See https://docs.microsoft.com/en-us/azure/key-vault/general/about-keys-secrets-certificates#vault-name-and-object-name"  # pylint: disable=line-too-long
+            )
             raise InvalidArgumentValueError(err_msg)
 
         segments = key_id[len(https_prefix):].split("/")
-        if len(segments) != 4 or segments[1] != "keys":
-            raise InvalidArgumentValueError(err_msg)
+
+        if is_pmk_enabled:
+            # PMK enabled (K2P): Only accept versionless key ID (3 segments: vault.net/keys/key-name)
+            if len(segments) != 3 or segments[1] != "keys":
+                err_msg = (
+                    "--azure-keyvault-kms-key-id is not a valid versionless Key Vault key ID for PMK. "
+                    "Valid format is https://{key-vault-url}/keys/{key-name}. "
+                    "See https://docs.microsoft.com/en-us/azure/key-vault/general/about-keys-secrets-certificates#vault-name-and-object-name"  # pylint: disable=line-too-long
+                )
+                raise InvalidArgumentValueError(err_msg)
+        else:
+            # PMK disabled (KMS v2): Accept versioned key ID (4 segments)
+            if len(segments) != 4 or segments[1] != "keys":
+                err_msg = (
+                    "--azure-keyvault-kms-key-id is not a valid Key Vault key ID. "
+                    "See https://docs.microsoft.com/en-us/azure/key-vault/general/about-keys-secrets-certificates#vault-name-and-object-name"  # pylint: disable=line-too-long
+                )
+                raise InvalidArgumentValueError(err_msg)
 
 
 def validate_azure_keyvault_kms_key_vault_resource_id(namespace):
     key_vault_resource_id = namespace.azure_keyvault_kms_key_vault_resource_id
+
+    # Check if PMK (Platform-Managed Keys) is enabled
+    is_pmk_enabled = (
+        hasattr(namespace, 'kms_infrastructure_encryption') and
+        namespace.kms_infrastructure_encryption == "Enabled"
+    )
+
     if key_vault_resource_id is None or key_vault_resource_id == '':
-        return
+        if is_pmk_enabled:
+            raise RequiredArgumentMissingError(
+                "--azure-keyvault-kms-key-vault-resource-id is required when "
+                "--kms-infrastructure-encryption is set to Enabled (PMK)."
+            )
+        else:
+            return
+
+    is_cmk_enabled = (
+        hasattr(namespace, 'enable_azure_keyvault_kms') and
+        namespace.enable_azure_keyvault_kms
+    )
+    if not is_cmk_enabled:
+        raise RequiredArgumentMissingError(
+            '"--azure-keyvault-kms-key-vault-resource-id" requires "--enable-azure-keyvault-kms".'
+        )
+
     if not is_valid_resource_id(key_vault_resource_id):
-        raise InvalidArgumentValueError("--azure-keyvault-kms-key-vault-resource-id is not a valid Azure resource ID.")
+        raise InvalidArgumentValueError(
+            "--azure-keyvault-kms-key-vault-resource-id is not a valid Azure resource ID."
+        )
+
+    try:
+        parsed = parse_resource_id(key_vault_resource_id)
+        provider = parsed.get('namespace', '').lower()
+        if provider != 'microsoft.keyvault':
+            raise InvalidArgumentValueError(
+                "--azure-keyvault-kms-key-vault-resource-id must reference a "
+                "Microsoft.KeyVault resource."
+            )
+        resource_type = parsed.get('type', '').lower()
+        if resource_type not in ['vaults', 'managedhsms']:
+            raise InvalidArgumentValueError(
+                "--azure-keyvault-kms-key-vault-resource-id must reference a Key Vault "
+                "(vaults) or Managed HSM (managedHSMs)."
+            )
+    except InvalidArgumentValueError:
+        # Re-raise our validation errors
+        raise
+    except Exception as ex:
+        raise InvalidArgumentValueError(
+            f"--azure-keyvault-kms-key-vault-resource-id parsing failed: {str(ex)}"
+        )
 
 
 def validate_bootstrap_container_registry_resource_id(namespace):

src/aks-preview/azext_aks_preview/managed_cluster_decorator.py

@@ -1288,6 +1288,29 @@ def get_kms_infrastructure_encryption(self) -> str:
 
         return kms_infrastructure_encryption
 
+    def get_azure_keyvault_kms_key_vault_resource_id(self) -> Union[str, None]:
+        """Obtain the value of azure_keyvault_kms_key_vault_resource_id.
+
+        :return: bool
+        """
+        # read the original value passed by the command
+        azure_keyvault_kms_key_vault_resource_id = self.raw_param.get(
+            "azure_keyvault_kms_key_vault_resource_id"
+        )
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if (
+                self.mc and
+                hasattr(self.mc, "security_profile") and  # backward compatibility
+                self.mc.security_profile and
+                self.mc.security_profile.azure_key_vault_kms and
+                self.mc.security_profile.azure_key_vault_kms.key_vault_resource_id is not None
+            ):
+                azure_keyvault_kms_key_vault_resource_id = (
+                    self.mc.security_profile.azure_key_vault_kms.key_vault_resource_id
+                )
+
+        return azure_keyvault_kms_key_vault_resource_id
+
     def get_cluster_snapshot_id(self) -> Union[str, None]:
         """Obtain the values of cluster_snapshot_id.
 
@@ -3313,8 +3336,8 @@ def set_up_image_integrity(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
-    def set_up_kms_infrastructure_encryption(self, mc: ManagedCluster) -> ManagedCluster:
-        """Set up security profile KubernetesResourceObjectEncryptionProfile for the ManagedCluster object.
+    def set_up_kms_pmk_and_cmk(self, mc: ManagedCluster) -> ManagedCluster:
+        """Set up security profile KubernetesResourceObjectEncryptionProfile and AzureKeyVaultKms for the ManagedCluster object.
 
         :return: the ManagedCluster object
         """
@@ -3335,6 +3358,17 @@ def set_up_kms_infrastructure_encryption(self, mc: ManagedCluster) -> ManagedClu
             # pylint: disable=line-too-long
             mc.security_profile.kubernetes_resource_object_encryption_profile.infrastructure_encryption = kms_infrastructure_encryption
 
+        if self.context.get_enable_azure_keyvault_kms():
+            key_id = self.context.get_azure_keyvault_kms_key_id()
+            if key_id:
+                if mc.security_profile is None:
+                    mc.security_profile = self.models.ManagedClusterSecurityProfile()
+                mc.security_profile.azure_key_vault_kms = self.models.AzureKeyVaultKms(
+                    enabled=True,
+                    key_id=key_id,
+                    key_vault_resource_id=self.context.get_azure_keyvault_kms_key_vault_resource_id(),
+                )
+
         return mc
 
     def set_up_creationdata_of_cluster_snapshot(self, mc: ManagedCluster) -> ManagedCluster:
@@ -3906,7 +3940,7 @@ def construct_mc_profile_preview(self, bypass_restore_defaults: bool = False) ->
         # set up image integrity
         mc = self.set_up_image_integrity(mc)
         # set up KMS infrastructure encryption
-        mc = self.set_up_kms_infrastructure_encryption(mc)
+        mc = self.set_up_kms_pmk_and_cmk(mc)
         # set up cluster snapshot
         mc = self.set_up_creationdata_of_cluster_snapshot(mc)
         # set up app routing profile