[Profile] az login: Add --client-id, --object-id and --resource-id for authenticating with user-assigned managed identity (#30525)

Author: jiasli

src/azure-cli-core/azure/cli/core/_profile.py

@@ -60,6 +60,11 @@
 
 _AZ_LOGIN_MESSAGE = "Please run 'az login' to setup account."
 
+MANAGED_IDENTITY_ID_WARNING = (
+    "Passing the managed identity ID with --username is deprecated and will be removed in a future release. "
+    "Please use --client-id, --object-id or --resource-id instead."
+)
+
 
 def load_subscriptions(cli_ctx, all_clouds=False, refresh=False):
     profile = Profile(cli_ctx=cli_ctx)
@@ -219,7 +224,8 @@ def login(self,
         self._set_subscriptions(consolidated)
         return deepcopy(consolidated)
 
-    def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=None):
+    def login_with_managed_identity(self, identity_id=None, client_id=None, object_id=None, resource_id=None,
+                                    allow_no_subscriptions=None):
         if _on_azure_arc():
             return self.login_with_managed_identity_azure_arc(
                 identity_id=identity_id, allow_no_subscriptions=allow_no_subscriptions)
@@ -229,7 +235,28 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N
         from azure.cli.core.auth.adal_authentication import MSIAuthenticationWrapper
         resource = self.cli_ctx.cloud.endpoints.active_directory_resource_id
 
-        if identity_id:
+        id_arg_count = len([arg for arg in (client_id, object_id, resource_id, identity_id) if arg])
+        if id_arg_count > 1:
+            raise CLIError('Usage error: Provide only one of --client-id, --object-id, --resource-id, or --username.')
+
+        if id_arg_count == 0:
+            identity_type = MsiAccountTypes.system_assigned
+            msi_creds = MSIAuthenticationWrapper(resource=resource)
+        elif client_id:
+            identity_type = MsiAccountTypes.user_assigned_client_id
+            identity_id = client_id
+            msi_creds = MSIAuthenticationWrapper(resource=resource, client_id=client_id)
+        elif object_id:
+            identity_type = MsiAccountTypes.user_assigned_object_id
+            identity_id = object_id
+            msi_creds = MSIAuthenticationWrapper(resource=resource, object_id=object_id)
+        elif resource_id:
+            identity_type = MsiAccountTypes.user_assigned_resource_id
+            identity_id = resource_id
+            msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=resource_id)
+        # The old way of re-using the same --username for 3 types of ID
+        elif identity_id:
+            logger.warning(MANAGED_IDENTITY_ID_WARNING)
             if is_valid_resource_id(identity_id):
                 msi_creds = MSIAuthenticationWrapper(resource=resource, msi_res_id=identity_id)
                 identity_type = MsiAccountTypes.user_assigned_resource_id
@@ -260,10 +287,6 @@ def login_with_managed_identity(self, identity_id=None, allow_no_subscriptions=N
                 if not authenticated:
                     raise CLIError('Failed to connect to MSI, check your managed service identity id.')
 
-        else:
-            identity_type = MsiAccountTypes.system_assigned
-            msi_creds = MSIAuthenticationWrapper(resource=resource)
-
         token_entry = msi_creds.token
         token = token_entry['access_token']
         logger.info('MSI: token was retrieved. Now trying to initialize local accounts...')

src/azure-cli-core/azure/cli/core/tests/test_profile.py

@@ -498,7 +498,7 @@ def test_login_in_cloud_shell(self, cloud_shell_credential_mock, create_subscrip
 
     @mock.patch('requests.get', autospec=True)
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
-    def test_find_subscriptions_in_vm_with_mi_system_assigned(self, create_subscription_client_mock, mock_get):
+    def test_login_with_mi_system_assigned(self, create_subscription_client_mock, mock_get):
         mock_subscription_client = mock.MagicMock()
         mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)]
         create_subscription_client_mock.return_value = mock_subscription_client
@@ -531,7 +531,7 @@ def test_find_subscriptions_in_vm_with_mi_system_assigned(self, create_subscript
 
     @mock.patch('requests.get', autospec=True)
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
-    def test_find_subscriptions_in_vm_with_mi_no_subscriptions(self, create_subscription_client_mock, mock_get):
+    def test_login_with_mi_no_subscriptions(self, create_subscription_client_mock, mock_get):
         mock_subscription_client = mock.MagicMock()
         mock_subscription_client.subscriptions.list.return_value = []
         create_subscription_client_mock.return_value = mock_subscription_client
@@ -566,8 +566,7 @@ def test_find_subscriptions_in_vm_with_mi_no_subscriptions(self, create_subscrip
 
     @mock.patch('requests.get', autospec=True)
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
-    def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, create_subscription_client_mock,
-                                                                           mock_get):
+    def test_login_with_mi_user_assigned_client_id(self, create_subscription_client_mock, mock_get):
         mock_subscription_client = mock.MagicMock()
         mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)]
         create_subscription_client_mock.return_value = mock_subscription_client
@@ -587,6 +586,19 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, cre
         good_response.content = encoded_test_token
         mock_get.return_value = good_response
 
+        subscriptions = profile.login_with_managed_identity(client_id=test_client_id)
+
+        self.assertEqual(len(subscriptions), 1)
+        s = subscriptions[0]
+        self.assertEqual(s['name'], self.display_name1)
+        self.assertEqual(s['id'], self.id1.split('/')[-1])
+        self.assertEqual(s['tenantId'], self.test_mi_tenant)
+
+        self.assertEqual(s['user']['name'], 'userAssignedIdentity')
+        self.assertEqual(s['user']['type'], 'servicePrincipal')
+        self.assertEqual(s['user']['assignedIdentityInfo'], 'MSIClient-{}'.format(test_client_id))
+
+        # Old way of using identity_id
         subscriptions = profile.login_with_managed_identity(identity_id=test_client_id)
 
         self.assertEqual(len(subscriptions), 1)
@@ -601,7 +613,7 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_client_id(self, cre
 
     @mock.patch('azure.cli.core.auth.adal_authentication.MSIAuthenticationWrapper', autospec=True)
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
-    def test_find_subscriptions_in_vm_with_mi_user_assigned_with_object_id(self, create_subscription_client_mock,
+    def test_login_with_mi_user_assigned_object_id(self, create_subscription_client_mock,
                                                                            mock_msi_auth):
         mock_subscription_client = mock.MagicMock()
         mock_subscription_client.subscriptions.list.return_value = [deepcopy(self.subscription1_raw)]
@@ -632,6 +644,14 @@ def set_token(self):
         mock_msi_auth.side_effect = AuthStub
         test_object_id = '54826b22-38d6-4fb2-bad9-b7b93a3e9999'
 
+        subscriptions = profile.login_with_managed_identity(object_id=test_object_id)
+
+        s = subscriptions[0]
+        self.assertEqual(s['user']['name'], 'userAssignedIdentity')
+        self.assertEqual(s['user']['type'], 'servicePrincipal')
+        self.assertEqual(s['user']['assignedIdentityInfo'], 'MSIObject-{}'.format(test_object_id))
+
+        # Old way of using identity_id
         subscriptions = profile.login_with_managed_identity(identity_id=test_object_id)
 
         s = subscriptions[0]
@@ -641,7 +661,7 @@ def set_token(self):
 
     @mock.patch('requests.get', autospec=True)
     @mock.patch('azure.cli.core._profile.SubscriptionFinder._create_subscription_client', autospec=True)
-    def test_find_subscriptions_in_vm_with_mi_user_assigned_with_res_id(self, create_subscription_client_mock,
+    def test_login_with_mi_user_assigned_resource_id(self, create_subscription_client_mock,
                                                                         mock_get):
 
         mock_subscription_client = mock.MagicMock()
@@ -665,6 +685,14 @@ def test_find_subscriptions_in_vm_with_mi_user_assigned_with_res_id(self, create
         good_response.content = encoded_test_token
         mock_get.return_value = good_response
 
+        subscriptions = profile.login_with_managed_identity(resource_id=test_res_id)
+
+        s = subscriptions[0]
+        self.assertEqual(s['user']['name'], 'userAssignedIdentity')
+        self.assertEqual(s['user']['type'], 'servicePrincipal')
+        self.assertEqual(subscriptions[0]['user']['assignedIdentityInfo'], 'MSIResource-{}'.format(test_res_id))
+
+        # Old way of using identity_id
         subscriptions = profile.login_with_managed_identity(identity_id=test_res_id)
 
         s = subscriptions[0]

src/azure-cli/azure/cli/command_modules/profile/__init__.py

@@ -75,6 +75,12 @@ def load_arguments(self, command):
             # Managed identity
             c.argument('identity', options_list=('-i', '--identity'), action='store_true',
                        help="Log in using managed identity", arg_group='Managed Identity')
+            c.argument('client_id',
+                       help="Client ID of the user-assigned managed identity", arg_group='Managed Identity')
+            c.argument('object_id',
+                       help="Object ID of the user-assigned managed identity", arg_group='Managed Identity')
+            c.argument('resource_id',
+                       help="Resource ID of the user-assigned managed identity", arg_group='Managed Identity')
 
         with self.argument_context('logout') as c:
             c.argument('username', help='account user, if missing, logout the current active account')

src/azure-cli/azure/cli/command_modules/profile/_help.py

@@ -43,8 +43,10 @@
       text: az login --service-principal --username APP_ID --certificate /path/to/cert.pem --tenant TENANT_ID
     - name: Log in with a system-assigned managed identity.
       text: az login --identity
-    - name: Log in with a user-assigned managed identity. You must specify the client ID, object ID or resource ID of the user-assigned managed identity with --username.
-      text: az login --identity --username 00000000-0000-0000-0000-000000000000
+    - name: Log in with a user-assigned managed identity's client ID.
+      text: az login --identity --client-id 00000000-0000-0000-0000-000000000000
+    - name: Log in with a user-assigned managed identity's resource ID.
+      text: az login --identity --resource-id /subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/MyResourceGroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/MyIdentity
 """
 
 helps['account'] = """

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

@@ -114,14 +114,14 @@ def account_clear(cmd):
     profile.logout_all()
 
 
-# pylint: disable=inconsistent-return-statements, too-many-branches
+# pylint: disable=too-many-branches, too-many-locals
 def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_subscriptions=False,
           # Device code flow
           use_device_code=False,
           # Service principal
           service_principal=None, certificate=None, use_cert_sn_issuer=None, client_assertion=None,
           # Managed identity
-          identity=False):
+          identity=False, client_id=None, object_id=None, resource_id=None):
     """Log in to access Azure subscriptions"""
 
     # quick argument usage check
@@ -143,7 +143,9 @@ def login(cmd, username=None, password=None, tenant=None, scopes=None, allow_no_
     if identity:
         if in_cloud_console():
             return profile.login_in_cloud_shell()
-        return profile.login_with_managed_identity(username, allow_no_subscriptions)
+        return profile.login_with_managed_identity(
+            identity_id=username, client_id=client_id, object_id=object_id, resource_id=resource_id,
+            allow_no_subscriptions=allow_no_subscriptions)
     if in_cloud_console():  # tell users they might not need login
         logger.warning(_CLOUD_CONSOLE_LOGIN_WARNING)
 

src/azure-cli/azure/cli/command_modules/profile/tests/latest/test_profile_custom.py

@@ -88,15 +88,15 @@ def test_get_raw_token(self, get_raw_token_mock):
         get_raw_token_mock.assert_called_with(mock.ANY, None, None, None, tenant_id)
 
     @mock.patch('azure.cli.command_modules.profile.custom.Profile', autospec=True)
-    def test_get_login(self, profile_mock):
+    def test_login_with_mi(self, profile_mock):
         invoked = []
 
-        def test_login(msi_port, identity_id=None):
+        def login_with_managed_identity_mock(*args, **kwargs):
             invoked.append(True)
 
         # mock the instance
         profile_instance = mock.MagicMock()
-        profile_instance.login_with_managed_identity = test_login
+        profile_instance.login_with_managed_identity = login_with_managed_identity_mock
         # mock the constructor
         profile_mock.return_value = profile_instance