{Compute} az vm identity: Migrate commands to aaz-based implementation

src/azure-cli/azure/cli/command_modules/sqlvm/tests/latest/test_sqlvm_commands.py

@@ -48,7 +48,7 @@ def __init__(self, name_prefix=sqlvm_name_prefix, location='westus',
     def create_resource(self, name, **kwargs):
         group = self._get_resource_group(**kwargs)
         template = ('az vm create -l {} -g {} -n {} --admin-username {} --admin-password {} --image {}'
-                    ' --size Standard_DS2_v2 --nsg-rule NONE')
+                    ' --size Standard_B2ms --nsg-rule NONE')
         execute(DummyCli(), template.format(self.location, group, name, self.vm_user, self.vm_password, self.image))
         return {self.parameter_name: name}
 

src/azure-cli/azure/cli/command_modules/vm/_vm_utils.py

@@ -7,10 +7,13 @@
 import os
 import re
 import importlib
+from enum import Enum
 
 from urllib.parse import urlparse
 
 from azure.cli.core.commands.arm import ArmTemplateBuilder
+from azure.cli.core.commands.client_factory import get_mgmt_service_client
+from azure.cli.core.profiles import ResourceType, get_sdk
 
 from knack.log import get_logger
 from knack.util import CLIError
@@ -32,7 +35,7 @@ def get_target_network_api(cli_ctx):
     if cli_ctx.cloud.profile == 'latest':
         version = '2022-01-01'
     else:
-        from azure.cli.core.profiles import get_api_version, ResourceType
+        from azure.cli.core.profiles import get_api_version
         version = get_api_version(cli_ctx, ResourceType.MGMT_NETWORK)
     return version
 
@@ -46,8 +49,6 @@ def read_content_if_is_file(string_or_file):
 
 
 def _resolve_api_version(cli_ctx, provider_namespace, resource_type, parent_path):
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from azure.cli.core.profiles import ResourceType
     client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES)
     provider = client.providers.get(provider_namespace)
 
@@ -75,10 +76,8 @@ def log_pprint_template(template):
 def check_existence(cli_ctx, value, resource_group, provider_namespace, resource_type,
                     parent_name=None, parent_type=None, static_version=None):
     # check for name or ID and set the type flags
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     from azure.core.exceptions import HttpResponseError
     from azure.mgmt.core.tools import parse_resource_id
-    from azure.cli.core.profiles import ResourceType
     id_parts = parse_resource_id(value)
     resource_client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES,
                                               subscription_id=id_parts.get('subscription', None)).resources
@@ -414,8 +413,6 @@ def _update(model, lun, value):
 
 
 def get_storage_blob_uri(cli_ctx, storage):
-    from azure.cli.core.profiles._shared import ResourceType
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
     if urlparse(storage).scheme:
         storage_uri = storage
     else:
@@ -757,3 +754,83 @@ def _open(filename, mode):
         f.write(public_bytes)
 
     return public_bytes.decode()
+
+
+def _gen_guid():
+    import uuid
+    return uuid.uuid4()
+
+
+def assign_identity(cli_ctx, getter, setter, identity_role=None, identity_scope=None):
+    import time
+    from azure.core.exceptions import HttpResponseError
+
+    # get
+    resource = getter()
+    resource = setter(resource)
+
+    # create role assignment:
+    if identity_scope:
+        principal_id = resource.get('identity', {}).get('principalId') or resource.get('identity', {}).get('principal_id')
+
+        identity_role_id = resolve_role_id(cli_ctx, identity_role, identity_scope)
+        assignments_client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_assignments
+        RoleAssignmentCreateParameters = get_sdk(cli_ctx, ResourceType.MGMT_AUTHORIZATION,
+                                                 'RoleAssignmentCreateParameters', mod='models',
+                                                 operation_group='role_assignments')
+        parameters = RoleAssignmentCreateParameters(role_definition_id=identity_role_id, principal_id=principal_id,
+                                                    principal_type=None)
+
+        logger.info("Creating an assignment with a role '%s' on the scope of '%s'", identity_role_id, identity_scope)
+        retry_times = 36
+        assignment_name = _gen_guid()
+        for retry_time in range(0, retry_times):
+            try:
+                assignments_client.create(scope=identity_scope, role_assignment_name=assignment_name,
+                                          parameters=parameters)
+                break
+            except HttpResponseError as ex:
+                if ex.error.code == 'RoleAssignmentExists':
+                    logger.info('Role assignment already exists')
+                    break
+                if retry_time < retry_times and ' does not exist in the directory ' in ex.message:
+                    time.sleep(5)
+                    logger.warning('Retrying role assignment creation: %s/%s', retry_time + 1,
+                                   retry_times)
+                    continue
+                raise
+    return resource
+
+
+def resolve_role_id(cli_ctx, role, scope):
+    import uuid
+    client = get_mgmt_service_client(cli_ctx, ResourceType.MGMT_AUTHORIZATION).role_definitions
+
+    role_id = None
+    if re.match(r'/subscriptions/[^/]+/providers/Microsoft.Authorization/roleDefinitions/',
+                role, re.I):
+        role_id = role
+    else:
+        try:
+            uuid.UUID(role)
+            role_id = '/subscriptions/{}/providers/Microsoft.Authorization/roleDefinitions/{}'.format(
+                client.config.subscription_id, role)
+        except ValueError:
+            pass
+        if not role_id:  # retrieve role id
+            role_defs = list(client.list(scope, "roleName eq '{}'".format(role)))
+            if not role_defs:
+                raise CLIError("Role '{}' doesn't exist.".format(role))
+            if len(role_defs) > 1:
+                ids = [r.id for r in role_defs]
+                err = "More than one role matches the given name '{}'. Please pick an id from '{}'"
+                raise CLIError(err.format(role, ids))
+            role_id = role_defs[0].id
+    return role_id
+
+
+class IdentityType(Enum):
+    SYSTEM_ASSIGNED = 'SystemAssigned'
+    USER_ASSIGNED = 'UserAssigned'
+    SYSTEM_ASSIGNED_USER_ASSIGNED = 'SystemAssigned, UserAssigned'
+    NONE = 'None'

src/azure-cli/azure/cli/command_modules/vm/commands.py

@@ -288,11 +288,12 @@ def load_command_table(self, _):
         from .operations.snapshot import SnapshotUpdate
         self.command_table['snapshot update'] = SnapshotUpdate(loader=self)
 
-    with self.command_group('vm', compute_vm_sdk) as g:
-        g.custom_command('identity assign', 'assign_vm_identity', validator=process_assign_identity_namespace)
-        g.custom_command('identity remove', 'remove_vm_identity', validator=process_remove_identity_namespace, min_api='2017-12-01')
-        g.custom_show_command('identity show', 'show_vm_identity')
+    with self.command_group('vm identity') as g:
+        g.custom_command('assign', 'assign_vm_identity', validator=process_assign_identity_namespace)
+        g.custom_command('remove', 'remove_vm_identity', validator=process_remove_identity_namespace, min_api='2017-12-01')
+        g.custom_show_command('show', 'show_vm_identity')
 
+    with self.command_group('vm', compute_vm_sdk) as g:
         g.custom_command('application set', 'set_vm_applications', validator=process_set_applications_namespace, min_api='2021-07-01')
         g.custom_command('application list', 'list_vm_applications', min_api='2021-07-01')