From bf29a5716eb483df12db4e490633301de7f350e1 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Fri, 7 Feb 2020 14:21:08 -0800
Subject: [PATCH 1/9] msi changes for GF and BF for omsagent (#1)

msi role assignment changes for greenfield and brownfield onboarding for omsagent
---
 .../azure/cli/command_modules/acs/custom.py   | 82 +++++++++++++++----
 1 file changed, 66 insertions(+), 16 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index 79b39263d45..b17fa798bb4 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -551,7 +551,7 @@ def _build_service_principal(rbac_client, cli_ctx, name, url, client_secret):
     return service_principal
 
 
-def _add_role_assignment(cli_ctx, role, service_principal, delay=2, scope=None):
+def _add_role_assignment(cli_ctx, role, service_principal_msi_id, isServicePrincipal=True, delay=2, scope=None):
     # AAD can have delays in propagating data, so sleep and retry
     hook = cli_ctx.get_progress_controller(True)
     hook.add(message='Waiting for AAD role to propagate', value=0, total_val=1.0)
@@ -560,7 +560,7 @@ def _add_role_assignment(cli_ctx, role, service_principal, delay=2, scope=None):
         hook.add(message='Waiting for AAD role to propagate', value=0.1 * x, total_val=1.0)
         try:
             # TODO: break this out into a shared utility library
-            create_role_assignment(cli_ctx, role, service_principal, scope=scope)
+            create_role_assignment(cli_ctx, role, service_principal_msi_id, isServicePrincipal, scope=scope)
             break
         except CloudError as ex:
             if ex.message == 'The role assignment already exists.':
@@ -1395,11 +1395,11 @@ def create_service_principal(cli_ctx, identifier, resolve_app=True, rbac_client=
     return rbac_client.service_principals.create(ServicePrincipalCreateParameters(app_id=app_id, account_enabled=True))
 
 
-def create_role_assignment(cli_ctx, role, assignee, resource_group_name=None, scope=None):
-    return _create_role_assignment(cli_ctx, role, assignee, resource_group_name, scope)
+def create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name=None, scope=None):
+    return _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name, scope)
 
 
-def _create_role_assignment(cli_ctx, role, assignee, resource_group_name=None, scope=None, resolve_assignee=True):
+def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name=None, scope=None, resolve_assignee=True):
     from azure.cli.core.profiles import ResourceType, get_sdk
     factory = get_auth_management_client(cli_ctx, scope)
     assignments_client = factory.role_assignments
@@ -1408,7 +1408,13 @@ def _create_role_assignment(cli_ctx, role, assignee, resource_group_name=None, s
     scope = _build_role_scope(resource_group_name, scope, assignments_client.config.subscription_id)
 
     role_id = _resolve_role_id(role, scope, definitions_client)
-    object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
+
+    # If the cluster has service principal resolve the service principal client id to get the object id, if not use MSI object id.
+    if isServicePrincipal:
+        object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
+    else:
+        object_id = assignee
+
     RoleAssignmentCreateParameters = get_sdk(cli_ctx, ResourceType.MGMT_AUTHORIZATION,
                                              'RoleAssignmentCreateParameters', mod='models',
                                              operation_group='role_assignments')
@@ -1640,6 +1646,7 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                api_server_authorized_ip_ranges=None,
                attach_acr=None,
                no_wait=False):
+
     _validate_ssh_key(no_ssh_key, ssh_key_value)
 
     subscription_id = get_subscription_id(cmd.cli_ctx)
@@ -1791,6 +1798,10 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                                  client.create_or_update,
                                  resource_group_name=resource_group_name,
                                  resource_name=name, parameters=mc)
+
+            # adding a wait here since we rely on the result for role assignment
+            result = LongRunningOperation(cmd.cli_ctx)(result)
+
             # add cluster spn with Monitoring Metrics Publisher role assignment to the cluster resource
             # mdm metrics supported only in azure public cloud so add the  role assignment only in this cloud
             cloud_name = cmd.cli_ctx.cloud.name
@@ -1802,10 +1813,28 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                     namespace='Microsoft.ContainerService', type='managedClusters',
                     name=name
                 )
-                if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
-                                            service_principal_profile.client_id, scope=cluster_resource_id):
-                    logger.warning('Could not create a role assignment for monitoring addon. '
-                                   'Are you an Owner on this subscription?')
+
+                service_principal_msi_id = None
+                # Check if service principal exists, if it does, assign permissions to service principal
+                # Else, provide permissions to MSI
+                if hasattr(result, 'service_principal_profile') and hasattr(result.service_principal_profile, 'client_id'):
+                    logger.warning('service principal exists, using it')
+                    service_principal_msi_id = result.service_principal_profile.client_id
+                    isServicePrincipal = True
+                elif ((hasattr(result, 'addon_profiles')) and
+                     ('omsagent' in result.addon_profiles) and
+                     (hasattr(result.addon_profiles['omsagent'], 'identity')) and
+                     (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))):
+                    logger.warning('omsagent MSI exists, using it')
+                    service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
+                    isServicePrincipal = False
+
+                if service_principal_msi_id is not None:
+                    if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
+                                                service_principal_msi_id, isServicePrincipal,
+                                                scope=cluster_resource_id):
+                        logger.warning('Could not create a role assignment for monitoring addon. '
+                                       'Are you an Owner on this subscription?')
             return result
         except CloudError as ex:
             retry_exception = ex
@@ -1833,17 +1862,23 @@ def aks_disable_addons(cmd, client, resource_group_name, name, addons, no_wait=F
     # send the managed cluster representation to update the addon profiles
     return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, name, instance)
 
-
 def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_resource_id=None,
                       subnet_name=None, no_wait=False):
     instance = client.get(resource_group_name, name)
     subscription_id = get_subscription_id(cmd.cli_ctx)
-    service_principal_client_id = instance.service_principal_profile.client_id
+
     instance = _update_addons(cmd, instance, subscription_id, resource_group_name, addons, enable=True,
                               workspace_resource_id=workspace_resource_id, subnet_name=subnet_name, no_wait=no_wait)
 
     if 'omsagent' in instance.addon_profiles:
         _ensure_container_insights_for_monitoring(cmd, instance.addon_profiles['omsagent'])
+        
+    # send the managed cluster representation to update the addon profiles
+    result = sdk_no_wait(no_wait, client.create_or_update,
+                          resource_group_name, name, instance)
+    result = LongRunningOperation(cmd.cli_ctx)(result)
+
+    if 'omsagent' in instance.addon_profiles:
         cloud_name = cmd.cli_ctx.cloud.name
         # mdm metrics supported only in Azure Public cloud so add the role assignment only in this cloud
         if cloud_name.lower() == 'azurecloud':
@@ -1854,13 +1889,28 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
                 namespace='Microsoft.ContainerService', type='managedClusters',
                 name=name
             )
+        service_principal_msi_id = None
+        # Check if service principal exists, if it does, assign permissions to service principal
+        # Else, provide permissions to MSI
+        if hasattr(result, 'service_principal_profile') and hasattr(result.service_principal_profile, 'client_id'):
+            logger.warning('service principal exists, using it')
+            service_principal_msi_id = result.service_principal_profile.client_id
+            isServicePrincipal = True
+        elif ((hasattr(result, 'addon_profiles')) and
+             ('omsagent' in result.addon_profiles) and
+             (hasattr(result.addon_profiles['omsagent'], 'identity')) and
+             (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))):
+            logger.warning('omsagent MSI exists, using it')
+            service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
+            isServicePrincipal = False
+
+        if service_principal_msi_id is not None:
             if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
-                                        service_principal_client_id, scope=cluster_resource_id):
+                                        service_principal_msi_id, isServicePrincipal, scope=cluster_resource_id):
                 logger.warning('Could not create a role assignment for Monitoring addon. '
-                               'Are you an Owner on this subscription?')
+                    'Are you an Owner on this subscription?')
 
-    # send the managed cluster representation to update the addon profiles
-    return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, name, instance)
+    return result
 
 
 def aks_get_versions(cmd, client, location):

From 57d48afd3e65678d28809c375d16e84a1814b4f3 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Fri, 7 Feb 2020 18:03:30 -0800
Subject: [PATCH 2/9] add role assignment to user assigned msi for omsagent
 (#2)

{omsagent}: fixing lint errors
---
 .../azure/cli/command_modules/acs/custom.py   | 45 ++++++++++++-------
 1 file changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index b17fa798bb4..b367618f64e 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1399,7 +1399,8 @@ def create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource
     return _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name, scope)
 
 
-def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name=None, scope=None, resolve_assignee=True):
+def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal,
+                            resource_group_name=None, scope=None, resolve_assignee=True):
     from azure.cli.core.profiles import ResourceType, get_sdk
     factory = get_auth_management_client(cli_ctx, scope)
     assignments_client = factory.role_assignments
@@ -1409,7 +1410,8 @@ def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resourc
 
     role_id = _resolve_role_id(role, scope, definitions_client)
 
-    # If the cluster has service principal resolve the service principal client id to get the object id, if not use MSI object id.
+    # If the cluster has service principal resolve the service principal client id to get the object id,
+    # if not use MSI object id.
     if isServicePrincipal:
         object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
     else:
@@ -1817,14 +1819,19 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                 service_principal_msi_id = None
                 # Check if service principal exists, if it does, assign permissions to service principal
                 # Else, provide permissions to MSI
-                if hasattr(result, 'service_principal_profile') and hasattr(result.service_principal_profile, 'client_id'):
+                if (
+                        hasattr(result, 'service_principal_profile') and
+                        hasattr(result.service_principal_profile, 'client_id')
+                ):
                     logger.warning('service principal exists, using it')
                     service_principal_msi_id = result.service_principal_profile.client_id
                     isServicePrincipal = True
-                elif ((hasattr(result, 'addon_profiles')) and
-                     ('omsagent' in result.addon_profiles) and
-                     (hasattr(result.addon_profiles['omsagent'], 'identity')) and
-                     (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))):
+                elif (
+                        (hasattr(result, 'addon_profiles')) and
+                        ('omsagent' in result.addon_profiles) and
+                        (hasattr(result.addon_profiles['omsagent'], 'identity')) and
+                        (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))
+                ):
                     logger.warning('omsagent MSI exists, using it')
                     service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
                     isServicePrincipal = False
@@ -1872,10 +1879,12 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
 
     if 'omsagent' in instance.addon_profiles:
         _ensure_container_insights_for_monitoring(cmd, instance.addon_profiles['omsagent'])
-        
     # send the managed cluster representation to update the addon profiles
-    result = sdk_no_wait(no_wait, client.create_or_update,
-                          resource_group_name, name, instance)
+    result = sdk_no_wait(
+        no_wait, client.create_or_update,
+        resource_group_name, name, instance
+    )
+
     result = LongRunningOperation(cmd.cli_ctx)(result)
 
     if 'omsagent' in instance.addon_profiles:
@@ -1896,10 +1905,12 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
             logger.warning('service principal exists, using it')
             service_principal_msi_id = result.service_principal_profile.client_id
             isServicePrincipal = True
-        elif ((hasattr(result, 'addon_profiles')) and
-             ('omsagent' in result.addon_profiles) and
-             (hasattr(result.addon_profiles['omsagent'], 'identity')) and
-             (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))):
+        elif (
+                (hasattr(result, 'addon_profiles')) and
+                ('omsagent' in result.addon_profiles) and
+                (hasattr(result.addon_profiles['omsagent'], 'identity')) and
+                (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))
+        ):
             logger.warning('omsagent MSI exists, using it')
             service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
             isServicePrincipal = False
@@ -1907,8 +1918,10 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
         if service_principal_msi_id is not None:
             if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
                                         service_principal_msi_id, isServicePrincipal, scope=cluster_resource_id):
-                logger.warning('Could not create a role assignment for Monitoring addon. '
-                    'Are you an Owner on this subscription?')
+                logger.warning(
+                    'Could not create a role assignment for Monitoring addon. '
+                    'Are you an Owner on this subscription?'
+                )
 
     return result
 

From d9715821ceeb07dca26de7fc2edd59128a399412 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Tue, 11 Feb 2020 14:41:59 -0800
Subject: [PATCH 3/9] fixing tests to add the new parameter

---
 .../cli/command_modules/acs/tests/latest/test_custom.py     | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
index 8df3e825578..c9652d4cba9 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
@@ -56,7 +56,7 @@ def test_add_role_assignment_basic(self):
         with mock.patch(
                 'azure.cli.command_modules.acs.custom.create_role_assignment') as create_role_assignment:
             ok = _add_role_assignment(cli_ctx, role, sp, delay=0)
-            create_role_assignment.assert_called_with(cli_ctx, role, sp, scope=None)
+            create_role_assignment.assert_called_with(cli_ctx, role, sp, True, scope=None)
             self.assertTrue(ok, 'Expected _add_role_assignment to succeed')
 
     def test_add_role_assignment_exists(self):
@@ -74,7 +74,7 @@ def test_add_role_assignment_exists(self):
             create_role_assignment.side_effect = err
             ok = _add_role_assignment(cli_ctx, role, sp, delay=0)
 
-            create_role_assignment.assert_called_with(cli_ctx, role, sp, scope=None)
+            create_role_assignment.assert_called_with(cli_ctx, role, sp, True, scope=None)
             self.assertTrue(ok, 'Expected _add_role_assignment to succeed')
 
     def test_add_role_assignment_fails(self):
@@ -92,7 +92,7 @@ def test_add_role_assignment_fails(self):
             create_role_assignment.side_effect = err
             ok = _add_role_assignment(cli_ctx, role, sp, delay=0)
 
-            create_role_assignment.assert_called_with(cli_ctx, role, sp, scope=None)
+            create_role_assignment.assert_called_with(cli_ctx, role, sp, True, scope=None)
             self.assertFalse(ok, 'Expected _add_role_assignment to fail')
 
     @mock.patch('azure.cli.core.commands.client_factory.get_subscription_id')

From c99fb9f2201e92ac4a59d645643ebe7bf9c742c8 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Tue, 11 Feb 2020 15:38:50 -0800
Subject: [PATCH 4/9] adding a new line after function to fix styling error

---
 src/azure-cli/azure/cli/command_modules/acs/custom.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index b367618f64e..e5299d8c5fa 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1869,6 +1869,7 @@ def aks_disable_addons(cmd, client, resource_group_name, name, addons, no_wait=F
     # send the managed cluster representation to update the addon profiles
     return sdk_no_wait(no_wait, client.create_or_update, resource_group_name, name, instance)
 
+
 def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_resource_id=None,
                       subnet_name=None, no_wait=False):
     instance = client.get(resource_group_name, name)

From b257a3669f05ec441cf16d0893566ae58f6c2bd4 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Wed, 12 Feb 2020 18:41:55 -0800
Subject: [PATCH 5/9] addressing PR comments

---
 .../azure/cli/command_modules/acs/custom.py   | 103 ++++++++----------
 .../acs/tests/latest/test_custom.py           |  11 ++
 2 files changed, 56 insertions(+), 58 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index e5299d8c5fa..acd1e29233f 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -551,7 +551,7 @@ def _build_service_principal(rbac_client, cli_ctx, name, url, client_secret):
     return service_principal
 
 
-def _add_role_assignment(cli_ctx, role, service_principal_msi_id, isServicePrincipal=True, delay=2, scope=None):
+def _add_role_assignment(cli_ctx, role, service_principal_msi_id, is_service_principal=True, delay=2, scope=None):
     # AAD can have delays in propagating data, so sleep and retry
     hook = cli_ctx.get_progress_controller(True)
     hook.add(message='Waiting for AAD role to propagate', value=0, total_val=1.0)
@@ -560,7 +560,7 @@ def _add_role_assignment(cli_ctx, role, service_principal_msi_id, isServicePrinc
         hook.add(message='Waiting for AAD role to propagate', value=0.1 * x, total_val=1.0)
         try:
             # TODO: break this out into a shared utility library
-            create_role_assignment(cli_ctx, role, service_principal_msi_id, isServicePrincipal, scope=scope)
+            create_role_assignment(cli_ctx, role, service_principal_msi_id, is_service_principal, scope=scope)
             break
         except CloudError as ex:
             if ex.message == 'The role assignment already exists.':
@@ -1395,11 +1395,11 @@ def create_service_principal(cli_ctx, identifier, resolve_app=True, rbac_client=
     return rbac_client.service_principals.create(ServicePrincipalCreateParameters(app_id=app_id, account_enabled=True))
 
 
-def create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name=None, scope=None):
-    return _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal, resource_group_name, scope)
+def create_role_assignment(cli_ctx, role, assignee, is_service_principal, resource_group_name=None, scope=None):
+    return _create_role_assignment(cli_ctx, role, assignee, is_service_principal, resource_group_name, scope)
 
 
-def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal,
+def _create_role_assignment(cli_ctx, role, assignee, is_service_principal,
                             resource_group_name=None, scope=None, resolve_assignee=True):
     from azure.cli.core.profiles import ResourceType, get_sdk
     factory = get_auth_management_client(cli_ctx, scope)
@@ -1412,7 +1412,7 @@ def _create_role_assignment(cli_ctx, role, assignee, isServicePrincipal,
 
     # If the cluster has service principal resolve the service principal client id to get the object id,
     # if not use MSI object id.
-    if isServicePrincipal:
+    if is_service_principal:
         object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
     else:
         object_id = assignee
@@ -1603,6 +1603,37 @@ def _validate_ssh_key(no_ssh_key, ssh_key_value):
             raise CLIError('Provided ssh key ({}) is invalid or non-existent'.format(shortened_key))
 
 
+def _add_monitoring_role_assignment(result, cluster_resource_id, cmd):
+    service_principal_msi_id = None
+    # Check if service principal exists, if it does, assign permissions to service principal
+    # Else, provide permissions to MSI
+    if (
+            hasattr(result, 'service_principal_profile') and
+            hasattr(result.service_principal_profile, 'client_id')
+    ):
+        logger.info('service principal exists, using it')
+        service_principal_msi_id = result.service_principal_profile.client_id
+        is_service_principal = True
+    elif (
+            (hasattr(result, 'addon_profiles')) and
+            ('omsagent' in result.addon_profiles) and
+            (hasattr(result.addon_profiles['omsagent'], 'identity')) and
+            (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))
+    ):
+        logger.info('omsagent MSI exists, using it')
+        service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
+        is_service_principal = False
+
+    if service_principal_msi_id is not None:
+        if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
+                                    service_principal_msi_id, is_service_principal, scope=cluster_resource_id):
+            logger.warning('Could not create a role assignment for Monitoring addon. '
+                           'Are you an Owner on this subscription?')
+    else:
+        logger.warning('Could not find service principal or user assigned MSI for role'
+                       'assignment')
+
+
 # pylint: disable=too-many-statements,too-many-branches
 def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint: disable=too-many-locals
                dns_name_prefix=None,
@@ -1802,12 +1833,13 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                                  resource_name=name, parameters=mc)
 
             # adding a wait here since we rely on the result for role assignment
-            result = LongRunningOperation(cmd.cli_ctx)(result)
 
             # add cluster spn with Monitoring Metrics Publisher role assignment to the cluster resource
             # mdm metrics supported only in azure public cloud so add the  role assignment only in this cloud
             cloud_name = cmd.cli_ctx.cloud.name
             if cloud_name.lower() == 'azurecloud' and monitoring:
+                # adding a wait here since we rely on the result for role assignment
+                result = LongRunningOperation(cmd.cli_ctx)(result)
                 from msrestazure.tools import resource_id
                 cluster_resource_id = resource_id(
                     subscription=subscription_id,
@@ -1816,32 +1848,8 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                     name=name
                 )
 
-                service_principal_msi_id = None
-                # Check if service principal exists, if it does, assign permissions to service principal
-                # Else, provide permissions to MSI
-                if (
-                        hasattr(result, 'service_principal_profile') and
-                        hasattr(result.service_principal_profile, 'client_id')
-                ):
-                    logger.warning('service principal exists, using it')
-                    service_principal_msi_id = result.service_principal_profile.client_id
-                    isServicePrincipal = True
-                elif (
-                        (hasattr(result, 'addon_profiles')) and
-                        ('omsagent' in result.addon_profiles) and
-                        (hasattr(result.addon_profiles['omsagent'], 'identity')) and
-                        (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))
-                ):
-                    logger.warning('omsagent MSI exists, using it')
-                    service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
-                    isServicePrincipal = False
-
-                if service_principal_msi_id is not None:
-                    if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
-                                                service_principal_msi_id, isServicePrincipal,
-                                                scope=cluster_resource_id):
-                        logger.warning('Could not create a role assignment for monitoring addon. '
-                                       'Are you an Owner on this subscription?')
+                _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
+
             return result
         except CloudError as ex:
             retry_exception = ex
@@ -1886,9 +1894,10 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
         resource_group_name, name, instance
     )
 
-    result = LongRunningOperation(cmd.cli_ctx)(result)
+    # result = LongRunningOperation(cmd.cli_ctx)(result)
 
     if 'omsagent' in instance.addon_profiles:
+        result = LongRunningOperation(cmd.cli_ctx)(result)
         cloud_name = cmd.cli_ctx.cloud.name
         # mdm metrics supported only in Azure Public cloud so add the role assignment only in this cloud
         if cloud_name.lower() == 'azurecloud':
@@ -1899,30 +1908,8 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
                 namespace='Microsoft.ContainerService', type='managedClusters',
                 name=name
             )
-        service_principal_msi_id = None
-        # Check if service principal exists, if it does, assign permissions to service principal
-        # Else, provide permissions to MSI
-        if hasattr(result, 'service_principal_profile') and hasattr(result.service_principal_profile, 'client_id'):
-            logger.warning('service principal exists, using it')
-            service_principal_msi_id = result.service_principal_profile.client_id
-            isServicePrincipal = True
-        elif (
-                (hasattr(result, 'addon_profiles')) and
-                ('omsagent' in result.addon_profiles) and
-                (hasattr(result.addon_profiles['omsagent'], 'identity')) and
-                (hasattr(result.addon_profiles['omsagent'].identity, 'object_id'))
-        ):
-            logger.warning('omsagent MSI exists, using it')
-            service_principal_msi_id = result.addon_profiles['omsagent'].identity.object_id
-            isServicePrincipal = False
-
-        if service_principal_msi_id is not None:
-            if not _add_role_assignment(cmd.cli_ctx, 'Monitoring Metrics Publisher',
-                                        service_principal_msi_id, isServicePrincipal, scope=cluster_resource_id):
-                logger.warning(
-                    'Could not create a role assignment for Monitoring addon. '
-                    'Are you an Owner on this subscription?'
-                )
+
+        _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
 
     return result
 
diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
index c9652d4cba9..56b1c424dff 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_custom.py
@@ -59,6 +59,17 @@ def test_add_role_assignment_basic(self):
             create_role_assignment.assert_called_with(cli_ctx, role, sp, True, scope=None)
             self.assertTrue(ok, 'Expected _add_role_assignment to succeed')
 
+    def test_add_role_assignment_msi_basic(self):
+        role = 'Owner'
+        sp = '1234567'
+        cli_ctx = mock.MagicMock()
+
+        with mock.patch(
+                'azure.cli.command_modules.acs.custom.create_role_assignment') as create_role_assignment:
+            ok = _add_role_assignment(cli_ctx, role, sp, False, delay=0)
+            create_role_assignment.assert_called_with(cli_ctx, role, sp, False, scope=None)
+            self.assertTrue(ok, 'Expected _add_role_assignment with msi to succeed')
+
     def test_add_role_assignment_exists(self):
         role = 'Owner'
         sp = '1234567'

From 45446bd5c89e4f3609f9cae65cdd3c5253ba5b09 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Wed, 19 Feb 2020 18:23:42 -0800
Subject: [PATCH 6/9] addressing PR comments

---
 src/azure-cli/azure/cli/command_modules/acs/_help.py |  5 ++++-
 .../azure/cli/command_modules/acs/custom.py          | 12 +++++-------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/_help.py b/src/azure-cli/azure/cli/command_modules/acs/_help.py
index fd630e62126..29571e3b82c 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/_help.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/_help.py
@@ -279,7 +279,9 @@
     long-summary: |-
         These addons are available:
             http_application_routing - configure ingress with automatic public DNS name creation.
-            monitoring - turn on Log Analytics monitoring. Uses the Log Analytics Default Workspace if it exists, else creates one. Specify "--workspace-resource-id" to use an existing workspace.
+            monitoring - turn on Log Analytics monitoring. Uses the Log Analytics Default Workspace if it exists, else creates one.
+                         Specify "--workspace-resource-id" to use an existing workspace.
+                         If monitoring addon is enabled --no-wait argument will have no effect
             virtual-node - enable AKS Virtual Node (PREVIEW). Requires --subnet-name to provide the name of an existing subnet for the Virtual Node to use.
   - name: --disable-rbac
     type: bool
@@ -447,6 +449,7 @@
     These addons are available:
         http_application_routing - configure ingress with automatic public DNS name creation.
         monitoring - turn on Log Analytics monitoring. Requires "--workspace-resource-id".
+                     If monitoring addon is enabled --no-wait argument will have no effect
         virtual-node - enable AKS Virtual Node (PREVIEW). Requires --subnet-name to provide the name of an existing subnet for the Virtual Node to use.
 parameters:
   - name: --addons -a
diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index acd1e29233f..4d91b36bbdd 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1609,9 +1609,10 @@ def _add_monitoring_role_assignment(result, cluster_resource_id, cmd):
     # Else, provide permissions to MSI
     if (
             hasattr(result, 'service_principal_profile') and
-            hasattr(result.service_principal_profile, 'client_id')
+            hasattr(result.service_principal_profile, 'client_id') and
+            result.service_principal_profile.client_id != 'msi'
     ):
-        logger.info('service principal exists, using it')
+        logger.info('valid service principal exists, using it')
         service_principal_msi_id = result.service_principal_profile.client_id
         is_service_principal = True
     elif (
@@ -1679,7 +1680,6 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                api_server_authorized_ip_ranges=None,
                attach_acr=None,
                no_wait=False):
-
     _validate_ssh_key(no_ssh_key, ssh_key_value)
 
     subscription_id = get_subscription_id(cmd.cli_ctx)
@@ -1832,8 +1832,6 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                                  resource_group_name=resource_group_name,
                                  resource_name=name, parameters=mc)
 
-            # adding a wait here since we rely on the result for role assignment
-
             # add cluster spn with Monitoring Metrics Publisher role assignment to the cluster resource
             # mdm metrics supported only in azure public cloud so add the  role assignment only in this cloud
             cloud_name = cmd.cli_ctx.cloud.name
@@ -1888,15 +1886,15 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
 
     if 'omsagent' in instance.addon_profiles:
         _ensure_container_insights_for_monitoring(cmd, instance.addon_profiles['omsagent'])
+
     # send the managed cluster representation to update the addon profiles
     result = sdk_no_wait(
         no_wait, client.create_or_update,
         resource_group_name, name, instance
     )
 
-    # result = LongRunningOperation(cmd.cli_ctx)(result)
-
     if 'omsagent' in instance.addon_profiles:
+        # adding a wait here since we rely on the result for role assignment
         result = LongRunningOperation(cmd.cli_ctx)(result)
         cloud_name = cmd.cli_ctx.cloud.name
         # mdm metrics supported only in Azure Public cloud so add the role assignment only in this cloud

From e071b60e363829ef64950a375530133e102aa5c2 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Sat, 22 Feb 2020 23:08:54 -0800
Subject: [PATCH 7/9] PR comments

---
 .../azure/cli/command_modules/acs/custom.py   | 73 +++++++++----------
 1 file changed, 34 insertions(+), 39 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index 40e99315a1d..16145b290ac 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1396,10 +1396,12 @@ def create_service_principal(cli_ctx, identifier, resolve_app=True, rbac_client=
 
 
 def create_role_assignment(cli_ctx, role, assignee, is_service_principal, resource_group_name=None, scope=None):
-    return _create_role_assignment(cli_ctx, role, assignee, is_service_principal, resource_group_name, scope)
+    return _create_role_assignment(cli_ctx,
+                                   role, assignee, resource_group_name,
+                                   scope, resolve_assignee=is_service_principal)
 
 
-def _create_role_assignment(cli_ctx, role, assignee, is_service_principal,
+def _create_role_assignment(cli_ctx, role, assignee,
                             resource_group_name=None, scope=None, resolve_assignee=True):
     from azure.cli.core.profiles import ResourceType, get_sdk
     factory = get_auth_management_client(cli_ctx, scope)
@@ -1412,10 +1414,7 @@ def _create_role_assignment(cli_ctx, role, assignee, is_service_principal,
 
     # If the cluster has service principal resolve the service principal client id to get the object id,
     # if not use MSI object id.
-    if is_service_principal:
-        object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
-    else:
-        object_id = assignee
+    object_id = _resolve_object_id(cli_ctx, assignee) if resolve_assignee else assignee
 
     RoleAssignmentCreateParameters = get_sdk(cli_ctx, ResourceType.MGMT_AUTHORIZATION,
                                              'RoleAssignmentCreateParameters', mod='models',
@@ -1829,28 +1828,30 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
     retry_exception = Exception(None)
     for _ in range(0, max_retry):
         try:
-            result = sdk_no_wait(no_wait,
-                                 client.create_or_update,
-                                 resource_group_name=resource_group_name,
-                                 resource_name=name, parameters=mc)
-
-            # add cluster spn with Monitoring Metrics Publisher role assignment to the cluster resource
-            # mdm metrics supported only in azure public cloud so add the  role assignment only in this cloud
-            cloud_name = cmd.cli_ctx.cloud.name
-            if cloud_name.lower() == 'azurecloud' and monitoring:
+            if monitoring:
                 # adding a wait here since we rely on the result for role assignment
-                result = LongRunningOperation(cmd.cli_ctx)(result)
-                from msrestazure.tools import resource_id
-                cluster_resource_id = resource_id(
-                    subscription=subscription_id,
-                    resource_group=resource_group_name,
-                    namespace='Microsoft.ContainerService', type='managedClusters',
-                    name=name
-                )
-
-                _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
-
-            return result
+                result = LongRunningOperation(cmd.cli_ctx)(client.create_or_update(
+                    resource_group_name=resource_group_name,
+                    resource_name=name,
+                    parameters=mc))
+                cloud_name = cmd.cli_ctx.cloud.name
+                # add cluster spn/msi Monitoring Metrics Publisher role assignment to publish metrics to MDM
+                # mdm metrics is supported only in azure public cloud, so add the role assignment only in this cloud
+                if cloud_name.lower() == 'azurecloud':
+                    from msrestazure.tools import resource_id
+                    cluster_resource_id = resource_id(
+                        subscription=subscription_id,
+                        resource_group=resource_group_name,
+                        namespace='Microsoft.ContainerService', type='managedClusters',
+                        name=name
+                    )
+                    _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
+                return result
+            else:
+                return sdk_no_wait(no_wait,
+                                   client.create_or_update,
+                                   resource_group_name=resource_group_name,
+                                   resource_name=name, parameters=mc)
         except CloudError as ex:
             retry_exception = ex
             if 'not found in Active Directory tenant' in ex.message:
@@ -1888,16 +1889,8 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
 
     if 'omsagent' in instance.addon_profiles and instance.addon_profiles['omsagent'].enabled:
         _ensure_container_insights_for_monitoring(cmd, instance.addon_profiles['omsagent'])
-
-    # send the managed cluster representation to update the addon profiles
-    result = sdk_no_wait(
-        no_wait, client.create_or_update,
-        resource_group_name, name, instance
-    )
-
-    if 'omsagent' in instance.addon_profiles:
         # adding a wait here since we rely on the result for role assignment
-        result = LongRunningOperation(cmd.cli_ctx)(result)
+        result = LongRunningOperation(cmd.cli_ctx)(client.create_or_update(resource_group_name, name, instance))
         cloud_name = cmd.cli_ctx.cloud.name
         # mdm metrics supported only in Azure Public cloud so add the role assignment only in this cloud
         if cloud_name.lower() == 'azurecloud':
@@ -1908,10 +1901,12 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
                 namespace='Microsoft.ContainerService', type='managedClusters',
                 name=name
             )
+            _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
 
-        _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
-
-    return result
+        return result
+    else:
+        return sdk_no_wait(no_wait, client.create_or_update,
+                           resource_group_name, name, instance)
 
 
 def aks_get_versions(cmd, client, location):

From e52a2360d07d6d89a19636b389d6bd967dc2d8c3 Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Sun, 23 Feb 2020 18:23:59 -0800
Subject: [PATCH 8/9] fixing style errors

---
 .../azure/cli/command_modules/acs/custom.py    | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index 16145b290ac..464c404d3dc 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1846,12 +1846,12 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                         name=name
                     )
                     _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
-                return result
             else:
-                return sdk_no_wait(no_wait,
-                                   client.create_or_update,
-                                   resource_group_name=resource_group_name,
-                                   resource_name=name, parameters=mc)
+                result = sdk_no_wait(no_wait,
+                                     client.create_or_update,
+                                     resource_group_name=resource_group_name,
+                                     resource_name=name, parameters=mc)
+            return result
         except CloudError as ex:
             retry_exception = ex
             if 'not found in Active Directory tenant' in ex.message:
@@ -1902,11 +1902,11 @@ def aks_enable_addons(cmd, client, resource_group_name, name, addons, workspace_
                 name=name
             )
             _add_monitoring_role_assignment(result, cluster_resource_id, cmd)
-
-        return result
     else:
-        return sdk_no_wait(no_wait, client.create_or_update,
-                           resource_group_name, name, instance)
+        result = sdk_no_wait(no_wait, client.create_or_update,
+                             resource_group_name, name, instance)
+
+    return result
 
 
 def aks_get_versions(cmd, client, location):

From d283c04e588a043a2811d60a2a2b69d1e67f627a Mon Sep 17 00:00:00 2001
From: rashmichandrashekar <rashmy@microsoft.com>
Date: Mon, 24 Feb 2020 15:37:36 -0800
Subject: [PATCH 9/9] adding tolower check

---
 src/azure-cli/azure/cli/command_modules/acs/custom.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py
index 464c404d3dc..1efe39010a5 100644
--- a/src/azure-cli/azure/cli/command_modules/acs/custom.py
+++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py
@@ -1611,7 +1611,7 @@ def _add_monitoring_role_assignment(result, cluster_resource_id, cmd):
     if (
             hasattr(result, 'service_principal_profile') and
             hasattr(result.service_principal_profile, 'client_id') and
-            result.service_principal_profile.client_id != 'msi'
+            result.service_principal_profile.client_id.lower() != 'msi'
     ):
         logger.info('valid service principal exists, using it')
         service_principal_msi_id = result.service_principal_profile.client_id
@@ -1682,7 +1682,6 @@ def aks_create(cmd, client, resource_group_name, name, ssh_key_value,  # pylint:
                attach_acr=None,
                no_wait=False):
     _validate_ssh_key(no_ssh_key, ssh_key_value)
-
     subscription_id = get_subscription_id(cmd.cli_ctx)
     if not dns_name_prefix:
         dns_name_prefix = _get_default_dns_prefix(name, resource_group_name, subscription_id)