[AKS] Implement workload identity support (#4685)

* feat: implement workload identity related flags

* feat: bump version

* fix: setup workload identity before oidc issuer

* test: define live tests

* fix: construct cmd with proper way

* fix: typo

* fix: typo

* test: add recordings

* fix: lint issues

* test: add test coverage

* chore: update version

* fix: update custom headers

src/aks-preview/HISTORY.md

@@ -2,6 +2,12 @@
 
 Release History
 ===============
+
+0.5.62
+++++++
+
+* Add support for managing workload identity feature.
+
 0.5.61
 ++++++
 * Add support for `--format` parameter in `az aks get-credentials` command.  

src/aks-preview/azext_aks_preview/_help.py

@@ -350,6 +350,9 @@
         - name: --enable-pod-identity-with-kubenet
           type: bool
           short-summary: (PREVIEW) Enable pod identity addon for cluster using Kubnet network plugin.
+        - name: --enable-workload-identity
+          type: bool
+          short-summary: (PREVIEW) Enable workload identity addon.
         - name: --aci-subnet-name
           type: string
           short-summary: The name of a subnet in an existing VNet into which to deploy the virtual nodes.
@@ -618,6 +621,12 @@
         - name: --disable-pod-identity
           type: bool
           short-summary: (PREVIEW) Disable Pod Identity addon for cluster.
+        - name: --enable-workload-identity
+          type: bool
+          short-summary: (PREVIEW) Enable Workload Identity addon for cluster.
+        - name: --disable-workload-identity
+          type: bool
+          short-summary: (PREVIEW) Disable Workload Identity addon for cluster.
         - name: --enable-secret-rotation
           type: bool
           short-summary: Enable secret rotation. Use with azure-keyvault-secrets-provider addon.

src/aks-preview/azext_aks_preview/_params.py

@@ -169,6 +169,7 @@ def load_arguments(self, _):
         c.argument('http_proxy_config', options_list=[
                    '--http-proxy-config'], type=str)
         c.argument('enable_pod_identity', action='store_true')
+        c.argument('enable_workload_identity', arg_type=get_three_state_flag(), is_preview=True)
         c.argument('appgw_name', options_list=[
                    '--appgw-name'], arg_group='Application Gateway')
         c.argument('appgw_subnet_prefix', options_list=[
@@ -255,6 +256,8 @@ def load_arguments(self, _):
                    validator=validate_assign_identity)
         c.argument('enable_pod_identity', action='store_true')
         c.argument('disable_pod_identity', action='store_true')
+        c.argument('enable_workload_identity', arg_type=get_three_state_flag(), is_preview=True)
+        c.argument('disable_workload_identity', arg_type=get_three_state_flag(), is_preview=True)
         c.argument('enable_secret_rotation', action='store_true')
         c.argument('disable_secret_rotation', action='store_true')
         c.argument('rotation_poll_interval', type=str)

src/aks-preview/azext_aks_preview/custom.py

@@ -753,6 +753,8 @@ def aks_create(cmd,
                auto_upgrade_channel=None,
                enable_pod_identity=False,
                enable_pod_identity_with_kubenet=False,
+               # NOTE: for workload identity flags, we need to know if it's set to True/False or not set (None)
+               enable_workload_identity=None,
                enable_encryption_at_host=False,
                enable_ultra_ssd=False,
                edge_zone=None,
@@ -834,6 +836,9 @@ def aks_update(cmd,     # pylint: disable=too-many-statements,too-many-branches,
                enable_pod_identity=False,
                enable_pod_identity_with_kubenet=False,
                disable_pod_identity=False,
+               # NOTE: for workload identity flags, we need to know if it's set to True/False or not set (None)
+               enable_workload_identity=None,
+               disable_workload_identity=None,
                enable_secret_rotation=False,
                disable_secret_rotation=False,
                rotation_poll_interval=None,

src/aks-preview/azext_aks_preview/decorator.py

@@ -7,7 +7,7 @@
 import os
 import time
 from types import SimpleNamespace
-from typing import Dict, List, Tuple, TypeVar, Union
+from typing import Dict, List, Tuple, TypeVar, Union, Optional
 
 from azure.cli.command_modules.acs._consts import (
     DecoratorEarlyExitException,
@@ -81,6 +81,7 @@
 ContainerServiceNetworkProfile = TypeVar("ContainerServiceNetworkProfile")
 ManagedClusterAddonProfile = TypeVar("ManagedClusterAddonProfile")
 ManagedClusterOIDCIssuerProfile = TypeVar('ManagedClusterOIDCIssuerProfile')
+ManagedClusterSecurityProfileWorkloadIdentity = TypeVar('ManagedClusterSecurityProfileWorkloadIdentity')
 Snapshot = TypeVar("Snapshot")
 AzureKeyVaultKms = TypeVar('AzureKeyVaultKms')
 
@@ -120,6 +121,11 @@ def __init__(self, cmd: AzCommandsLoader, resource_type: ResourceType):
             resource_type=self.resource_type,
             operation_group="managed_clusters",
         )
+        self.ManagedClusterSecurityProfileWorkloadIdentity = self.__cmd.get_models(
+            "ManagedClusterSecurityProfileWorkloadIdentity",
+            resource_type=self.resource_type,
+            operation_group="managed_clusters",
+        )
         self.ManagedClusterSecurityProfile = self.__cmd.get_models(
             "ManagedClusterSecurityProfile",
             resource_type=self.resource_type,
@@ -1579,6 +1585,56 @@ def get_oidc_issuer_profile(self) -> ManagedClusterOIDCIssuerProfile:
 
         return profile
 
+    def get_workload_identity_profile(self) -> Optional[ManagedClusterSecurityProfileWorkloadIdentity]:
+        """Obtrain the value of security_profile.workload_identity.
+
+        :return: Optional[ManagedClusterSecurityProfileWorkloadIdentity]
+        """
+        enable_workload_identity = self.raw_param.get("enable_workload_identity")
+        disable_workload_identity = self.raw_param.get("disable_workload_identity")
+        if self.decorator_mode == DecoratorMode.CREATE:
+            # CREATE mode has no --disable-workload-identity flag
+            disable_workload_identity = None
+
+        if enable_workload_identity is None and disable_workload_identity is None:
+            # no flags have been set, return None; server side will backfill the default/existing value
+            return None
+
+        if enable_workload_identity and disable_workload_identity:
+            raise MutuallyExclusiveArgumentError(
+                "Cannot specify --enable-workload-identity and "
+                "--disable-workload-identity at the same time."
+            )
+
+        profile = self.models.ManagedClusterSecurityProfileWorkloadIdentity()
+        if self.decorator_mode == DecoratorMode.CREATE:
+            profile.enabled = bool(enable_workload_identity)
+        elif self.decorator_mode == DecoratorMode.UPDATE:
+            if self.mc.security_profile is not None and self.mc.security_profile.workload_identity is not None:
+                profile = self.mc.security_profile.workload_identity
+            if enable_workload_identity:
+                profile.enabled = True
+            elif disable_workload_identity:
+                profile.enabled = False
+
+        if profile.enabled:
+            # in enable case, we need to check if OIDC issuer has been enabled
+            oidc_issuer_profile = self.get_oidc_issuer_profile()
+            if self.decorator_mode == DecoratorMode.UPDATE and oidc_issuer_profile is None:
+                # if the cluster has enabled OIDC issuer before, in update call:
+                #
+                #    az aks update --enable-workload-identity
+                #
+                # we need to use previous OIDC issuer profile
+                oidc_issuer_profile = self.mc.oidc_issuer_profile
+            oidc_issuer_enabled = oidc_issuer_profile is not None and oidc_issuer_profile.enabled
+            if not oidc_issuer_enabled:
+                raise RequiredArgumentMissingError(
+                    "Enabling workload identity requires enabling OIDC issuer (--enable-oidc-issuer)."
+                )
+
+        return profile
+
     def get_crg_id(self) -> str:
         """Obtain the values of crg_id.
 
@@ -1992,6 +2048,24 @@ def set_up_oidc_issuer_profile(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def set_up_workload_identity_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Set up workload identity for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        profile = self.context.get_workload_identity_profile()
+        if profile is None:
+            if mc.security_profile is not None:
+                # set the value to None to let server side to fill in the default value
+                mc.security_profile.workload_identity = None
+            return mc
+
+        if mc.security_profile is None:
+            mc.security_profile = self.models.ManagedClusterSecurityProfile()
+        mc.security_profile.workload_identity = profile
+
+        return mc
+
     def set_up_azure_keyvault_kms(self, mc: ManagedCluster) -> ManagedCluster:
         """Set up security profile azureKeyVaultKms for the ManagedCluster object.
 
@@ -2027,7 +2101,15 @@ def construct_mc_preview_profile(self) -> ManagedCluster:
         mc = self.set_up_pod_security_policy(mc)
         # set up pod identity profile
         mc = self.set_up_pod_identity_profile(mc)
+
+        # update workload identity & OIDC issuer settings
+        # NOTE: in current implementation, workload identity settings setup requires checking
+        #       previous OIDC issuer profile. However, the OIDC issuer settings setup will
+        #       overrides the previous OIDC issuer profile based on user input. Therefore, we have
+        #       to make sure the workload identity settings setup is done after OIDC issuer settings.
+        mc = self.set_up_workload_identity_profile(mc)
         mc = self.set_up_oidc_issuer_profile(mc)
+
         mc = self.set_up_azure_keyvault_kms(mc)
         return mc
 
@@ -2181,7 +2263,9 @@ def check_raw_parameters(self):
                 '"--nodepool-labels" or '
                 '"--enable-oidc-issuer" or '
                 '"--http-proxy-config" or '
-                '"--enable-azure-keyvault-kms".'
+                '"--enable-azure-keyvault-kms" or '
+                '"--enable-workload-identity" or '
+                '"--disable-workload-identity".'
             )
 
     def update_load_balancer_profile(self, mc: ManagedCluster) -> ManagedCluster:
@@ -2317,6 +2401,26 @@ def update_oidc_issuer_profile(self, mc: ManagedCluster) -> ManagedCluster:
 
         return mc
 
+    def update_workload_identity_profile(self, mc: ManagedCluster) -> ManagedCluster:
+        """Update workload identity profile for the ManagedCluster object.
+
+        :return: the ManagedCluster object
+        """
+        self._ensure_mc(mc)
+
+        profile = self.context.get_workload_identity_profile()
+        if profile is None:
+            if mc.security_profile is not None:
+                # set the value to None to let server side to fill in the default value
+                mc.security_profile.workload_identity = None
+            return mc
+
+        if mc.security_profile is None:
+            mc.security_profile = self.models.ManagedClusterSecurityProfile()
+        mc.security_profile.workload_identity = profile
+
+        return mc
+
     def update_azure_keyvault_kms(self, mc: ManagedCluster) -> ManagedCluster:
         """Update security profile azureKeyvaultKms for the ManagedCluster object.
 
@@ -2367,7 +2471,15 @@ def update_mc_preview_profile(self) -> ManagedCluster:
         mc = self.update_nat_gateway_profile(mc)
         # update pod identity profile
         mc = self.update_pod_identity_profile(mc)
+
+        # update workload identity & OIDC issuer settings
+        # NOTE: in current implementation, workload identity settings setup requires checking
+        #       previous OIDC issuer profile. However, the OIDC issuer settings setup will
+        #       overrides the previous OIDC issuer profile based on user input. Therefore, we have
+        #       to make sure the workload identity settings setup is done after OIDC issuer settings.
+        mc = self.update_workload_identity_profile(mc)
         mc = self.update_oidc_issuer_profile(mc)
+
         mc = self.update_http_proxy_config(mc)
         mc = self.update_azure_keyvault_kms(mc)
         return mc