[AKS] az aks update: Add option --assignee-principal-type to specify the principal type when using --attach-acr (#31464)

Author: jovieir

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

@@ -580,8 +580,10 @@ def load_arguments(self, _):
         c.argument('gmsa_dns_server')
         c.argument('gmsa_root_domain_name')
         c.argument('disable_windows_gmsa', action='store_true')
-        c.argument('attach_acr', acr_arg_type, validator=validate_acr)
-        c.argument('detach_acr', acr_arg_type, validator=validate_acr)
+        c.argument('attach_acr', acr_arg_type, validator=validate_acr, help='Attach an ACR to the AKS cluster.')
+        c.argument('detach_acr', acr_arg_type, validator=validate_acr, help='Detach an ACR from the AKS cluster.')
+        c.argument('assignee_principal_type', validator=validate_acr, options_list=['--assignee-principal-type', '-t'],
+                   help='Use this parameter in conjunction with --attach-acr to explicitly set the principal type in the ACR role assignment. This helps avoid RBAC-related errors by ensuring the correct principal type is applied. Valid values are \'User\', \'Group\' or \'ServicePrincipal\'')
         c.argument('disable_defender', action='store_true', validator=validate_defender_disable_and_enable_parameters)
         c.argument('enable_defender', action='store_true')
         c.argument('defender_config', validator=validate_defender_config_parameter)
@@ -693,7 +695,6 @@ def load_arguments(self, _):
         c.argument('disable_cost_analysis', action='store_true')
         c.argument("if_match")
         c.argument("if_none_match")
-
     with self.argument_context('aks disable-addons', resource_type=ResourceType.MGMT_CONTAINERSERVICE, operation_group='managed_clusters') as c:
         c.argument('addons', options_list=['--addons', '-a'])
 

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

@@ -63,7 +63,8 @@ def build_role_scope(resource_group_name: str, scope: str, subscription_id: str)
     return scope
 
 
-def add_role_assignment_executor(cmd, role, assignee, resource_group_name=None, scope=None, resolve_assignee=True):
+def add_role_assignment_executor(cmd, role, assignee, resource_group_name=None,
+                                 scope=None, resolve_assignee=True, assignee_principal_type=None):
     factory = get_auth_management_client(cmd.cli_ctx, scope)
     assignments_client = factory.role_assignments
     definitions_client = factory.role_definitions
@@ -94,7 +95,7 @@ def add_role_assignment_executor(cmd, role, assignee, resource_group_name=None,
     )
     if cmd.supported_api_version(min_api="2018-01-01-preview", resource_type=ResourceType.MGMT_AUTHORIZATION):
         parameters = RoleAssignmentCreateParameters(role_definition_id=role_id, principal_id=object_id,
-                                                    principal_type=None)
+                                                    principal_type=assignee_principal_type)
         return assignments_client.create(scope, assignment_name, parameters, headers=custom_headers)
 
     # for backward compatibility
@@ -109,7 +110,8 @@ def add_role_assignment_executor(cmd, role, assignee, resource_group_name=None,
     return assignments_client.create(scope, assignment_name, properties, headers=custom_headers)
 
 
-def add_role_assignment(cmd, role, service_principal_msi_id, is_service_principal=True, delay=2, scope=None):
+def add_role_assignment(cmd, role, service_principal_msi_id, is_service_principal=True,
+                        delay=2, scope=None, assignee_principal_type=None):
     # AAD can have delays in propagating data, so sleep and retry
     hook = cmd.cli_ctx.get_progress_controller(True)
     hook.add(message="Waiting for AAD role to propagate", value=0, total_val=1.0)
@@ -124,6 +126,7 @@ def add_role_assignment(cmd, role, service_principal_msi_id, is_service_principa
                 service_principal_msi_id,
                 scope=scope,
                 resolve_assignee=is_service_principal,
+                assignee_principal_type=assignee_principal_type,
             )
             break
         except HttpResponseError as ex:
@@ -305,21 +308,26 @@ def ensure_cluster_identity_permission_on_kubelet_identity(cmd, cluster_identity
         )
 
 
-def ensure_aks_acr_role_assignment(cmd, assignee, registry_id, detach=False, is_service_principal=True):
+def ensure_aks_acr_role_assignment(cmd, assignee, registry_id, detach=False, is_service_principal=True,
+                                   assignee_principal_type=None):
+
     if detach:
         if not delete_role_assignments(
             cmd.cli_ctx, "acrpull", assignee, scope=registry_id, is_service_principal=is_service_principal
         ):
             raise AzCLIError("Could not delete role assignments for ACR. Are you an Owner on this subscription?")
         return
 
-    if not add_role_assignment(cmd, "acrpull", assignee, scope=registry_id, is_service_principal=is_service_principal):
+    if not add_role_assignment(cmd, "acrpull", assignee, scope=registry_id, is_service_principal=is_service_principal,
+                               assignee_principal_type=assignee_principal_type):
         raise AzCLIError("Could not create a role assignment for ACR. Are you an Owner on this subscription?")
     return
 
 
 # pylint: disable=unused-argument
-def ensure_aks_acr(cmd, assignee, acr_name_or_id, subscription_id, detach=False, is_service_principal=True):
+def ensure_aks_acr(cmd, assignee, acr_name_or_id, subscription_id, detach=False,
+                   is_service_principal=True, assignee_principal_type=None):
+
     from azure.mgmt.core.tools import is_valid_resource_id, parse_resource_id
 
     # Check if the ACR exists by resource ID.
@@ -330,7 +338,8 @@ def ensure_aks_acr(cmd, assignee, acr_name_or_id, subscription_id, detach=False,
             registry = acr_client.registries.get(parsed_registry["resource_group"], parsed_registry["name"])
         except HttpResponseError as ex:
             raise AzCLIError(ex.message)
-        ensure_aks_acr_role_assignment(cmd, assignee, registry.id, detach, is_service_principal)
+        ensure_aks_acr_role_assignment(cmd, assignee, registry.id, detach,
+                                       is_service_principal, assignee_principal_type)
         return
 
     # Check if the ACR exists by name accross all resource groups.
@@ -342,5 +351,5 @@ def ensure_aks_acr(cmd, assignee, acr_name_or_id, subscription_id, detach=False,
         if "was not found" in ex.message:
             raise AzCLIError("ACR {} not found. Have you provided the right ACR name?".format(registry_name))
         raise AzCLIError(ex.message)
-    ensure_aks_acr_role_assignment(cmd, assignee, registry.id, detach, is_service_principal)
+    ensure_aks_acr_role_assignment(cmd, assignee, registry.id, detach, is_service_principal, assignee_principal_type)
     return

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

@@ -341,6 +341,17 @@ def validate_spot_max_price(namespace):
 def validate_acr(namespace):
     if namespace.attach_acr and namespace.detach_acr:
         raise CLIError('Cannot specify "--attach-acr" and "--detach-acr" at the same time.')
+    if namespace.assignee_principal_type and not namespace.attach_acr:
+        raise CLIError('The "--assignee-principal-type" argument can only be used with "--attach-acr".')
+    if namespace.attach_acr:
+        # Validate assignee_principal_type if specified
+        if namespace.assignee_principal_type:
+            valid_types = ['User', 'ServicePrincipal', 'Group']
+            if namespace.assignee_principal_type not in valid_types:
+                raise CLIError(
+                    f"Invalid value for --assignee_principal_type. "
+                    f"Allowed values are: {', '.join(valid_types)}"
+                )
 
 
 def validate_nodepool_tags(ns):

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

@@ -756,6 +756,7 @@ def aks_update(
     disable_windows_gmsa=False,
     attach_acr=None,
     detach_acr=None,
+    assignee_principal_type=None,
     nrg_lockdown_restriction_level=None,
     enable_defender=False,
     disable_defender=False,

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

@@ -1872,6 +1872,21 @@ def get_detach_acr(self) -> Union[str, None]:
         # this parameter does not need validation
         return detach_acr
 
+    def get_assignee_principal_type(self) -> Union[str, None]:
+        """Obtain the value of assignee_principal_type.
+
+        This function will return the override value of the assignee principal type (e.g., User, ServicePrincipal, Group)
+        to determine the type of identity being assigned for the ACR role. This is used
+        during the attach ACR operation to ensure proper role assignment based on the identity type.
+
+        :return: string or None
+        """
+        # read the original value passed by the command
+        assignee_principal_type = self.raw_param.get("assignee_principal_type")
+        # this parameter does not need dynamic completion
+        # this parameter does not need validation
+        return assignee_principal_type
+
     def get_http_proxy_config(self) -> Union[Dict, ManagedClusterHTTPProxyConfig, None]:
         """Obtain the value of http_proxy_config.
 
@@ -7006,6 +7021,7 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None:
                     acr_name_or_id=attach_acr,
                     subscription_id=self.context.get_subscription_id(),
                     is_service_principal=False,
+                    assignee_principal_type=self.context.get_assignee_principal_type()
                 )
 
         # azure monitor metrics addon (v2)
@@ -7374,14 +7390,15 @@ def process_attach_detach_acr(self, mc: ManagedCluster) -> None:
         assignee, is_service_principal = self.context.get_assignee_from_identity_or_sp_profile()
         attach_acr = self.context.get_attach_acr()
         detach_acr = self.context.get_detach_acr()
-
+        assignee_principal_type = self.context.get_assignee_principal_type()
         if attach_acr:
             self.context.external_functions.ensure_aks_acr(
                 self.cmd,
                 assignee=assignee,
                 acr_name_or_id=attach_acr,
                 subscription_id=subscription_id,
                 is_service_principal=is_service_principal,
+                assignee_principal_type=assignee_principal_type,
             )
 
         if detach_acr:
@@ -7392,6 +7409,7 @@ def process_attach_detach_acr(self, mc: ManagedCluster) -> None:
                 subscription_id=subscription_id,
                 detach=True,
                 is_service_principal=is_service_principal,
+                assignee_principal_type=assignee_principal_type,
             )
 
     def update_azure_service_mesh_profile(self, mc: ManagedCluster) -> ManagedCluster:
@@ -8925,6 +8943,7 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None:
                     acr_name_or_id=attach_acr,
                     subscription_id=self.context.get_subscription_id(),
                     is_service_principal=False,
+                    assignee_principal_type=self.context.get_assignee_principal_type()
                 )
 
         enable_azure_container_storage = self.context.get_intermediate("enable_azure_container_storage")

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_aks_commands.py

@@ -7291,7 +7291,7 @@ def test_aks_update_attatch_acr(self, resource_group, resource_group_location, s
         assignee_object_id = _get_test_sp_object_id(sp_name)
         role_assignment_check_cmd = (
             "role assignment list --scope {acr_scope} --assignee " +
-            assignee_object_id if assignee_object_id else sp_name
+            assignee_object_id if assignee_object_id else "role assignment list --scope {acr_scope} --assignee " + sp_name
         )
 
         # attach acr
@@ -7308,6 +7308,27 @@ def test_aks_update_attatch_acr(self, resource_group, resource_group_location, s
         # check role assignment
         self.cmd(role_assignment_check_cmd, checks=[self.is_empty()])
 
+        # Test with --assignee-principal-type
+        attach_with_principal_cmd = "aks update --resource-group={resource_group} --name={name} --attach-acr={acr_name} --assignee-principal-type=ServicePrincipal"
+        self.cmd(attach_with_principal_cmd)
+
+        # check role assignment with principal type
+        role_assignment_with_principal_check_cmd = (
+        "role assignment list --scope {acr_scope} --assignee " +
+        assignee_object_id if assignee_object_id else "role assignment list --scope {acr_scope} --assignee " + sp_name
+        )
+        self.cmd(role_assignment_with_principal_check_cmd, checks=[
+            self.check('length(@) == `1`', True),
+            self.check('[0].principalType', 'ServicePrincipal')
+        ])
+        
+        # detach acr
+        attach_cmd = 'aks update --resource-group={resource_group} --name={name} --detach-acr={acr_name}'
+        self.cmd(attach_cmd)
+
+        # check role assignment
+        self.cmd(role_assignment_check_cmd, checks=[self.is_empty()])
+
         # delete
         self.cmd(
             'aks delete -g {resource_group} -n {name} --yes --no-wait', checks=[self.is_empty()])

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_managed_cluster_decorator.py

@@ -6578,7 +6578,7 @@ def test_process_attach_acr(self):
             return_value=registry,
         ), patch("azure.cli.command_modules.acs._roleassignments.ensure_aks_acr_role_assignment") as ensure_assignment:
             dec_3.process_attach_acr(mc_3)
-        ensure_assignment.assert_called_once_with(self.cmd, "test_service_principal", "test_registry_id", False, True)
+        ensure_assignment.assert_called_once_with(self.cmd, "test_service_principal", "test_registry_id", False, True, None)
 
     def test_set_up_network_profile(self):
         # default value in `aks_create`
@@ -8222,6 +8222,7 @@ def test_postprocessing_after_mc_created(self):
             acr_name_or_id="test_attach_acr",
             subscription_id="1234-5678-9012",
             is_service_principal=False,
+            assignee_principal_type=None,
         )
 
     def test_put_mc(self):
@@ -8946,6 +8947,7 @@ def test_process_attach_detach_acr(self):
                         acr_name_or_id="test_attach_acr",
                         subscription_id="test_subscription_id",
                         is_service_principal=False,
+                        assignee_principal_type=None,
                     ),
                     call(
                         self.cmd,
@@ -8954,6 +8956,7 @@ def test_process_attach_detach_acr(self):
                         subscription_id="test_subscription_id",
                         detach=True,
                         is_service_principal=False,
+                        assignee_principal_type=None,
                     ),
                 ]
             )
@@ -11929,6 +11932,7 @@ def test_postprocessing_after_mc_created(self):
             acr_name_or_id="test_attach_acr",
             subscription_id="1234-5678-9012",
             is_service_principal=False,
+            assignee_principal_type=None,
         )
 
     def test_put_mc(self):