[Refractor] - Migrated remove_vm_identity function

Author: william051200

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

@@ -7,6 +7,7 @@
 import os
 import re
 import importlib
+from enum import Enum
 
 from urllib.parse import urlparse
 
@@ -831,4 +832,11 @@ def resolve_role_id(cli_ctx, role, scope):
                 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
+    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/custom.py

@@ -2545,18 +2545,81 @@ def _remove_identities(cmd, resource_group_name, name, identities, getter, sette
     return result.identity
 
 
+def _remove_identities_by_aaz(cmd, resource_group_name, name, identities, getter, setter):
+    from ._vm_utils import MSI_LOCAL_ID, IdentityType
+
+    remove_system_assigned_identity = False
+
+    if MSI_LOCAL_ID in identities:
+        remove_system_assigned_identity = True
+        identities.remove(MSI_LOCAL_ID)
+
+    resource = getter(cmd, resource_group_name, name)
+    existing_identity = resource.get('identity')
+
+    if existing_identity is None:
+        return None
+
+    existing_emsis = [x.lower() for x in list((existing_identity.get('userAssignedIdentities') or {}).keys())]
+
+    if identities:
+        emsis_to_remove = [x.lower() for x in identities]
+
+        non_existing = [emsis for emsis in emsis_to_remove if not emsis in existing_emsis]
+        if non_existing:
+            raise CLIError("'{}' are not associated with '{}'".format(','.join(non_existing), name))
+
+        emsis_to_be_retain = [emsis for emsis in existing_emsis if not emsis in emsis_to_remove]
+
+        if not emsis_to_be_retain:  # if all emsis are gone, we need to update the type
+            if existing_identity['type'] == IdentityType.USER_ASSIGNED.value:
+                existing_identity['type'] = IdentityType.NONE.value
+            elif existing_identity['type'] == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value:
+                existing_identity['type'] = IdentityType.SYSTEM_ASSIGNED.value
+
+        existing_identity['userAssignedIdentities'] = {}
+        for emsis in identities:
+            existing_identity['userAssignedIdentities'][emsis] = {}
+    else:
+        existing_identity['userAssignedIdentities'] = None
+
+    if remove_system_assigned_identity:
+        if existing_identity['type'] == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value or existing_identity['type'] == IdentityType.USER_ASSIGNED.value:
+            existing_identity['type'] = IdentityType.USER_ASSIGNED.value
+        else:
+            existing_identity['type'] = IdentityType.NONE.value
+
+    result = LongRunningOperation(cmd.cli_ctx)(setter(resource_group_name, name, resource))
+    return result.get('identity') or None
+
+
 def remove_vm_identity(cmd, resource_group_name, vm_name, identities=None):
     def setter(resource_group_name, vm_name, vm):
-        client = _compute_client_factory(cmd.cli_ctx)
-        VirtualMachineUpdate = cmd.get_models('VirtualMachineUpdate', operation_group='virtual_machines')
-        vm_update = VirtualMachineUpdate(identity=vm.identity)
-        return client.virtual_machines.begin_update(resource_group_name, vm_name, vm_update)
+        command_args = {
+            'resource_group': resource_group_name,
+            'vm_name': vm_name
+        }
+
+        from ._vm_utils import IdentityType
+        if vm.get('identity') and vm.get('identity').get('type') == IdentityType.USER_ASSIGNED.value:
+            command_args['mi_user_assigned'] = [key for key in list((vm.get('identity').get('userAssignedIdentities') or {}).keys())] + ['UserAssigned']
+        elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED.value:
+            command_args['mi_user_assigned'] = []
+            command_args['mi_system_assigned'] = 'True'
+        elif vm.get('identity') and vm.get('identity').get('type') == IdentityType.SYSTEM_ASSIGNED_USER_ASSIGNED.value:
+            command_args['mi_user_assigned'] = [key for key in list((vm.get('identity').get('userAssignedIdentities') or {}).keys())]
+            command_args['mi_system_assigned'] = 'True'
+        else:
+            command_args['mi_user_assigned'] = []
+
+        from .operations.vm import VMIdentityRemove
+        return VMIdentityRemove(cli_ctx=cmd.cli_ctx)(command_args=command_args)
 
     if identities is None:
         from ._vm_utils import MSI_LOCAL_ID
         identities = [MSI_LOCAL_ID]
 
-    return _remove_identities(cmd, resource_group_name, vm_name, identities, get_vm, setter)
+    return _remove_identities_by_aaz(cmd, resource_group_name, vm_name, identities, get_vm_migrated, setter)
 
 
 # region VirtualMachines Images

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

@@ -3,11 +3,15 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 # pylint: disable=no-self-use, line-too-long, protected-access, too-few-public-methods, unused-argument, too-many-statements, too-many-branches, too-many-locals
+from typing import override
+import json
+
 from knack.log import get_logger
 
 from azure.cli.core.aaz import AAZStrType
-from ..aaz.latest.vm import (Show as _VMShow, ListSizes as _VMListSizes,
+from ..aaz.latest.vm import (Show as _VMShow, ListSizes as _VMListSizes, Patch as _VMPatch,
                              Update as _VMUpdate, Capture as _VMCapture, Create as _VMCreate)
+from .._vm_utils import IdentityType
 
 logger = get_logger(__name__)
 
@@ -155,6 +159,76 @@ def _output(self, *args, **kwargs):
         return result
 
 
+class VMIdentityRemove(_VMPatch):
+    def _output(self, *args, **kwargs):
+        result = self.deserialize_output(self.ctx.vars.instance, client_flatten=True)
+
+        identity = result.get('identity')
+        if not identity:
+            return result
+
+        if not identity.get('principalId'):
+            identity['principalId'] = None
+
+        if not identity.get('tenantId'):
+            identity['tenantId'] = None
+
+        if not identity.get('userAssignedIdentities'):
+            identity['userAssignedIdentities'] = None
+
+        return result
+
+    class VirtualMachinesUpdate(_VMPatch.VirtualMachinesUpdate):
+        def _format_content(self, content):
+            if type(content) == str:
+                content = json.loads(content)
+
+            if not content.get('identity'):
+                content['identity'] = {
+                    'userAssignedIdentities': None,
+                    'type': IdentityType.NONE.value
+                }
+                return json.dumps(content)
+
+            identities = content.get('identity', {}).get('userAssignedIdentities')
+            if identities:
+                if 'UserAssigned' in identities.keys():
+                    identities.pop('UserAssigned')
+
+                for key in identities.keys():
+                    identities[key] = None
+
+            if len(content.get('identity', {}).get('userAssignedIdentities', {}).keys()) < 1:
+                content['identity']['userAssignedIdentities'] = None
+
+            return json.dumps(content)
+
+        def __call__(self, *args, **kwargs):
+            request = self.make_request()
+            request.data = self._format_content(request.data)
+            session = self.client.send_request(request=request, stream=False, **kwargs)
+            if session.http_response.status_code in [202]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "azure-async-operation"},
+                    path_format_arguments=self.url_parameters,
+                )
+            if session.http_response.status_code in [200]:
+                return self.client.build_lro_polling(
+                    self.ctx.args.no_wait,
+                    session,
+                    self.on_200,
+                    self.on_error,
+                    lro_options={"final-state-via": "azure-async-operation"},
+                    path_format_arguments=self.url_parameters,
+                )
+
+            return self.on_error(session.http_response)
+
+
 def convert_show_result_to_snake_case(result):
     new_result = {}
     if "extendedLocation" in result: