[Role] az role assignment create/update: auto complete `assignee_pr…

…incipal_type` (Azure#17669)

* auto complete for assignee_principle_type

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

@@ -154,10 +154,12 @@ def create_role_assignment(cmd, role, assignee=None, assignee_object_id=None, re
     if condition_version and not condition:
         raise CLIError('usage error: When --condition-version is set, --condition must be set as well.')
 
+    object_id, principal_type = _resolve_assignee_object(cmd.cli_ctx, assignee, assignee_object_id,
+                                                         assignee_principal_type)
+
     try:
-        return _create_role_assignment(cmd.cli_ctx, role, assignee or assignee_object_id, resource_group_name, scope,
-                                       resolve_assignee=(not assignee_object_id),
-                                       assignee_principal_type=assignee_principal_type, description=description,
+        return _create_role_assignment(cmd.cli_ctx, role, object_id, resource_group_name, scope, resolve_assignee=False,
+                                       assignee_principal_type=principal_type, description=description,
                                        condition=condition, condition_version=condition_version)
     except Exception as ex:  # pylint: disable=broad-except
         if _error_caused_by_role_assignment_exists(ex):  # for idempotent
@@ -281,6 +283,9 @@ def update_role_assignment(cmd, role_assignment):
             original_assignment.condition_version.startswith('2.') and assignment.condition_version.startswith('1.')):
         raise CLIError("Condition version cannot be downgraded to '1.X'.")
 
+    if not assignment.principal_type:
+        assignment.principal_type = original_assignment.principal_type
+
     return assignments_client.create(scope, name, parameters=assignment)
 
 
@@ -1479,7 +1484,8 @@ def create_service_principal_for_rbac(
             logger.warning("Creating '%s' role assignment under scope '%s'", role, scope)
             for retry_time in range(0, _RETRY_TIMES):
                 try:
-                    _create_role_assignment(cmd.cli_ctx, role, sp_oid, None, scope, resolve_assignee=False)
+                    _create_role_assignment(cmd.cli_ctx, role, sp_oid, None, scope, resolve_assignee=False,
+                                            assignee_principal_type='ServicePrincipal')
                     break
                 except Exception as ex:
                     if retry_time < _RETRY_TIMES and ' does not exist in the directory ' in str(ex):
@@ -1780,6 +1786,45 @@ def _encode_custom_key_description(key_description):
     return key_description.encode('utf-16')
 
 
+def _resolve_assignee_object(cli_ctx, assignee, assignee_object_id, assignee_principal_type):
+    client = _graph_client_factory(cli_ctx)
+    result = None
+
+    # resolve assignee (same as _resolve_object_id)
+    if assignee:
+        if assignee.find('@') >= 0:  # looks like a user principal name
+            result = list(client.users.list(filter="userPrincipalName eq '{}'".format(assignee)))
+        if not result:
+            result = list(client.service_principals.list(
+                filter="servicePrincipalNames/any(c:c eq '{}')".format(assignee)))
+        if not result and is_guid(assignee):  # assume an object id, let us verify it
+            result = _get_object_stubs(client, [assignee])
+
+        # 2+ matches should never happen, so we only check 'no match' here
+        if not result:
+            raise CLIError("Cannot find user or service principal in graph database for '{assignee}'. "
+                           "If the assignee is an appId, make sure the corresponding service principal is created "
+                           "with 'az ad sp create --id {assignee}'.".format(assignee=assignee))
+
+        return result[0].object_id, result[0].object_type
+
+    # try to resolve assignee object id
+    try:
+        result = _get_object_stubs(client, [assignee_object_id])
+        if result:
+            return result[0].object_id, result[0].object_type
+    except CloudError:
+        pass
+
+    # If failed to verify assignee object id, DO NOT raise exception
+    # since --assignee-object-id is exposed to bypass Graph API
+    if not assignee_principal_type:
+        logger.warning('Failed to query --assignee-principal-type for %s by invoking Graph API.\n'
+                       'RBAC server might reject creating role assignment without --assignee-principal-type '
+                       'in the future. Better to specify --assignee-principal-type manually.', assignee_object_id)
+    return assignee_object_id, assignee_principal_type
+
+
 def _resolve_object_id(cli_ctx, assignee, fallback_to_object_id=False):
     client = _graph_client_factory(cli_ctx)
     result = None