[AKS] Support use custom kubelet identity (#18615)

* Add custom kubelet identity implementation

* Add recording file

* Fix lint

* Fix lint 2

* Apply code review

* Fix errors

* Apply code review

* Fix lint

* Fix lint 2

linter_exclusions.yml

@@ -260,6 +260,9 @@ aks create:
       rule_exclusions:
       - option_length_too_long
     enable_encryption_at_host:
+      rule_exclusions:
+      - option_length_too_long  
+    assign_kubelet_identity:
       rule_exclusions:
       - option_length_too_long
 aks enable-addons:

src/azure-cli/azure/cli/command_modules/acs/_consts.py

@@ -49,3 +49,6 @@
 }
 
 CONST_CANIPULL_IMAGE = "mcr.microsoft.com/aks/canipull:0.0.2-alpha"
+
+CONST_MANAGED_IDENTITY_OPERATOR_ROLE = 'Managed Identity Operator'
+CONST_MANAGED_IDENTITY_OPERATOR_ROLE_ID = 'f1a07417-d97a-45cb-824c-7a7467783830'

src/azure-cli/azure/cli/command_modules/acs/_help.py

@@ -413,6 +413,9 @@
   - name: --assign-identity
     type: string
     short-summary: Specify an existing user assigned identity for control plane's usage in order to manage cluster resource group.
+  - name: --assign-kubelet-identity
+    type: string
+    short-summary: Specify an existing user assigned identity for kubelet's usage, which is typically used to pull image from ACR.
   - name: --node-osdisk-diskencryptionset-id -d
     type: string
     short-summary: ResourceId of the disk encryption set to use for enabling encryption at rest on agent node os disk.
@@ -489,6 +492,8 @@
     text: az aks create -g MyResourceGroup -n MyManagedCluster --enable-ultra-ssd
   - name: Create a kubernetes cluster with Azure RBAC enabled.
     text: az aks create -g MyResourceGroup -n MyManagedCluster --enable-aad --enable-azure-rbac
+  - name: Create a kubernetes cluster with custom control plane identity and kubelet identity.
+    text: az aks create -g MyResourceGroup -n MyManagedCluster --assign-identity <control-plane-identity-resource-id> --assign-kubelet-identity <kubelet-identity-resource-id>
 """
 
 helps['aks update'] = """

src/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -24,7 +24,7 @@
     validate_priority, validate_eviction_policy, validate_spot_max_price,
     validate_load_balancer_outbound_ip_prefixes, validate_taints, validate_ip_ranges, validate_acr, validate_nodepool_tags,
     validate_load_balancer_outbound_ports, validate_load_balancer_idle_timeout, validate_vnet_subnet_id, validate_nodepool_labels,
-    validate_ppg, validate_assign_identity, validate_max_surge)
+    validate_ppg, validate_assign_identity, validate_max_surge, validate_assign_kubelet_identity)
 from ._consts import CONST_OUTBOUND_TYPE_LOAD_BALANCER, CONST_OUTBOUND_TYPE_USER_DEFINED_ROUTING, \
     CONST_SCALE_SET_PRIORITY_REGULAR, CONST_SCALE_SET_PRIORITY_SPOT, \
     CONST_SPOT_EVICTION_POLICY_DELETE, CONST_SPOT_EVICTION_POLICY_DEALLOCATE, \
@@ -264,6 +264,7 @@ def load_arguments(self, _):
                    '--appgw-subnet-id'], arg_group='Application Gateway')
         c.argument('appgw_watch_namespace', options_list=[
                    '--appgw-watch-namespace'], arg_group='Application Gateway')
+        c.argument('assign_kubelet_identity', validator=validate_assign_kubelet_identity)
         c.argument('yes', options_list=[
                    '--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')
         c.argument('enable_sgxquotehelper', action='store_true')

src/azure-cli/azure/cli/command_modules/acs/_validators.py

@@ -421,3 +421,12 @@ def validate_assign_identity(namespace):
         from msrestazure.tools import is_valid_resource_id
         if not is_valid_resource_id(namespace.assign_identity):
             raise InvalidArgumentValueError("--assign-identity is not a valid Azure resource ID.")
+
+
+def validate_assign_kubelet_identity(namespace):
+    if namespace.assign_kubelet_identity is not None:
+        if namespace.assign_kubelet_identity == '':
+            return
+        from msrestazure.tools import is_valid_resource_id
+        if not is_valid_resource_id(namespace.assign_kubelet_identity):
+            raise InvalidArgumentValueError("--assign-kubelet-identity is not a valid Azure resource ID.")

src/azure-cli/azure/cli/command_modules/acs/custom.py

@@ -47,11 +47,12 @@
 from azure.cli.command_modules.acs._params import regions_in_preview, regions_in_prod
 from azure.cli.core.api import get_config_dir
 from azure.cli.core.azclierror import (ResourceNotFoundError,
-                                       ArgumentUsageError,
                                        ClientRequestError,
+                                       ArgumentUsageError,
                                        InvalidArgumentValueError,
                                        MutuallyExclusiveArgumentError,
-                                       ValidationError)
+                                       ValidationError,
+                                       UnauthorizedError)
 from azure.cli.core._profile import Profile
 from azure.cli.core.profiles import ResourceType
 from azure.cli.core.commands.client_factory import get_mgmt_service_client, get_subscription_id
@@ -100,6 +101,7 @@
 from ._consts import ADDONS
 from ._consts import CONST_CANIPULL_IMAGE
 from ._consts import CONST_PRIVATE_DNS_ZONE_SYSTEM
+from ._consts import CONST_MANAGED_IDENTITY_OPERATOR_ROLE, CONST_MANAGED_IDENTITY_OPERATOR_ROLE_ID
 
 logger = get_logger(__name__)
 
@@ -784,10 +786,16 @@ def _generate_properties(api_version, orchestrator_type, orchestrator_version, m
     return properties
 
 
-def _get_user_assigned_identity_client_id(cli_ctx, resource_id):
-    pattern = '/subscriptions/(.*?)/resourcegroups/(.*?)/providers/microsoft.managedidentity/userassignedidentities/(.*)'  # pylint: disable=line-too-long
+def _get_user_assigned_identity_resource_id_regular_expression():
+    return re.compile(
+        r'/subscriptions/(.*?)/resourcegroups/(.*?)/providers/microsoft.managedidentity/userassignedidentities/(.*)',
+        flags=re.IGNORECASE)
+
+
+def _get_user_assigned_identity(cli_ctx, resource_id):
     resource_id = resource_id.lower()
-    match = re.search(pattern, resource_id)
+    _re_user_assigned_identity_resource_id = _get_user_assigned_identity_resource_id_regular_expression()
+    match = _re_user_assigned_identity_resource_id.search(resource_id)
     if match:
         subscription_id = match.group(1)
         resource_group_name = match.group(2)
@@ -798,14 +806,21 @@ def _get_user_assigned_identity_client_id(cli_ctx, resource_id):
                                                                resource_name=identity_name)
         except CloudError as ex:
             if 'was not found' in ex.message:
-                raise ResourceNotFoundError(
-                    "Identity {} not found.".format(resource_id))
+                raise ResourceNotFoundError("Identity {} not found.".format(resource_id))
             raise ClientRequestError(ex.message)
-        return identity.client_id
+        return identity
     raise InvalidArgumentValueError(
         "Cannot parse identity name from provided resource id {}.".format(resource_id))
 
 
+def _get_user_assigned_identity_client_id(cli_ctx, resource_id):
+    return _get_user_assigned_identity(cli_ctx, resource_id).client_id
+
+
+def _get_user_assigned_identity_object_id(cli_ctx, resource_id):
+    return _get_user_assigned_identity(cli_ctx, resource_id).principal_id
+
+
 # pylint: disable=too-many-locals
 def acs_create(cmd, client, resource_group_name, deployment_name, name, ssh_key_value, dns_name_prefix=None,
                location=None, admin_username="azureuser", api_version=None, master_profile=None,
@@ -1958,6 +1973,7 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                appgw_watch_namespace=None,
                enable_sgxquotehelper=False,
                enable_encryption_at_host=False,
+               assign_kubelet_identity=None,
                enable_ultra_ssd=False,
                no_wait=False,
                yes=False,
@@ -1992,6 +2008,9 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
     ManagedClusterIdentity = cmd.get_models('ManagedClusterIdentity',
                                             resource_type=ResourceType.MGMT_CONTAINERSERVICE,
                                             operation_group='managed_clusters')
+    ManagedClusterPropertiesIdentityProfileValue = cmd.get_models('ManagedClusterPropertiesIdentityProfileValue',
+                                                                  resource_type=ResourceType.MGMT_CONTAINERSERVICE,
+                                                                  operation_group='managed_clusters')
     ManagedCluster = cmd.get_models('ManagedCluster',
                                     resource_type=ResourceType.MGMT_CONTAINERSERVICE,
                                     operation_group='managed_clusters')
@@ -2291,6 +2310,25 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
             user_assigned_identities=user_assigned_identity
         )
 
+    identity_profile = None
+    if assign_kubelet_identity:
+        if not assign_identity:
+            # pylint: disable=line-too-long
+            raise ArgumentUsageError('--assign-kubelet-identity can only be specified when --assign-identity is specified')
+        kubelet_identity = _get_user_assigned_identity(cmd.cli_ctx, assign_kubelet_identity)
+        identity_profile = {
+            'kubeletidentity': ManagedClusterPropertiesIdentityProfileValue(
+                resource_id=assign_kubelet_identity,
+                client_id=kubelet_identity.client_id,
+                object_id=kubelet_identity.principal_id
+            )
+        }
+        cluster_identity_object_id = _get_user_assigned_identity_object_id(cmd.cli_ctx, assign_identity)
+        # ensure the cluster identity has "Managed Identity Operator" role at the scope of kubelet identity
+        _ensure_cluster_identity_permission_on_kubelet_identity(
+            cmd.cli_ctx,
+            cluster_identity_object_id)
+
     mc = ManagedCluster(
         location=location,
         tags=tags,
@@ -2307,7 +2345,8 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
         auto_scaler_profile=cluster_autoscaler_profile,
         api_server_access_profile=api_server_access_profile,
         identity=identity,
-        disk_encryption_set_id=node_osdisk_diskencryptionset_id
+        disk_encryption_set_id=node_osdisk_diskencryptionset_id,
+        identity_profile=identity_profile
     )
 
     use_custom_private_dns_zone = False
@@ -4694,3 +4733,23 @@ def _put_managed_cluster_ensuring_permission(
                               headers=headers)
 
     return cluster
+
+
+def _ensure_cluster_identity_permission_on_kubelet_identity(cli_ctx, cluster_identity_object_id, scope):
+    factory = get_auth_management_client(cli_ctx, scope)
+    assignments_client = factory.role_assignments
+
+    for i in assignments_client.list_for_scope(scope=scope, filter='atScope()'):
+        if i.scope.lower() != scope.lower():
+            continue
+        if not i.role_definition_id.lower().endswith(CONST_MANAGED_IDENTITY_OPERATOR_ROLE_ID):
+            continue
+        if i.principal_id.lower() != cluster_identity_object_id.lower():
+            continue
+        # already assigned
+        return
+
+    if not _add_role_assignment(cli_ctx, CONST_MANAGED_IDENTITY_OPERATOR_ROLE, cluster_identity_object_id,
+                                is_service_principal=False, scope=scope):
+        raise UnauthorizedError('Could not grant Managed Identity Operator '
+                                'permission to cluster identity at scope {}'.format(scope))