get_msal_token

Author: jiasli

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

@@ -299,26 +299,24 @@ def logout_all(self):
         identity.logout_all_users()
         identity.logout_all_service_principal()
 
-    def get_login_credentials(self, subscription_id=None, aux_subscriptions=None, aux_tenants=None):
+    def get_login_credentials(self, subscription_id=None, aux_subscriptions=None, aux_tenants=None,
+                              sdk_credential=True):
         """Get a credential compatible with Track 2 SDK."""
         if aux_tenants and aux_subscriptions:
             raise CLIError("Please specify only one of aux_subscriptions and aux_tenants, not both")
 
         account = self.get_subscription(subscription_id)
 
         managed_identity_type, managed_identity_id = Profile._parse_managed_identity_account(account)
-
+        external_credentials = None
         if in_cloud_console() and account[_USER_ENTITY].get(_CLOUD_SHELL_ID):
             # Cloud Shell
             from .auth.msal_credentials import CloudShellCredential
-            # The credential must be wrapped by CredentialAdaptor so that it can work with SDK.
-            sdk_cred = CredentialAdaptor(CloudShellCredential())
+            cred = CloudShellCredential()
 
         elif managed_identity_type:
             # managed identity
-            # The credential must be wrapped by CredentialAdaptor so that it can work with SDK.
             cred = ManagedIdentityAuth.credential_factory(managed_identity_type, managed_identity_id)
-            sdk_cred = CredentialAdaptor(cred)
 
         else:
             # user and service principal
@@ -332,13 +330,15 @@ def get_login_credentials(self, subscription_id=None, aux_subscriptions=None, au
                     if sub[_TENANT_ID] != account[_TENANT_ID]:
                         external_tenants.append(sub[_TENANT_ID])
 
-            credential = self._create_credential(account)
+            cred = self._create_credential(account)
             external_credentials = []
             for external_tenant in external_tenants:
                 external_credentials.append(self._create_credential(account, tenant_id=external_tenant))
-            sdk_cred = CredentialAdaptor(credential, auxiliary_credentials=external_credentials)
 
-        return (sdk_cred,
+        # Wrapping the credential with CredentialAdaptor makes it compatible with SDK.
+        cred_result = CredentialAdaptor(cred, auxiliary_credentials=external_credentials) if sdk_credential else cred
+
+        return (cred_result,
                 str(account[_SUBSCRIPTION_ID]),
                 str(account[_TENANT_ID]))
 
@@ -404,8 +404,9 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
     def get_msal_token(self, scopes, data):
         """Get VM SSH certificate. Do not use it for other purposes. To get an access token, use get_raw_token instead.
         """
-        credential, _, _ = self.get_login_credentials()
-        certificate_string = credential.get_token(*scopes, data=data).token
+        credential, _, _ = self.get_login_credentials(sdk_credential=False)
+        from .auth.constants import ACCESS_TOKEN
+        certificate_string = credential.acquire_token(scopes, data=data)[ACCESS_TOKEN]
         # The first value used to be username, but it is no longer used.
         return None, certificate_string
 

src/azure-cli-core/azure/cli/core/auth/credential_adaptor.py

@@ -57,10 +57,6 @@ def _prepare_msal_kwargs(options=None):
     # Both get_token's kwargs and get_token_info's options are accepted as their schema is the same (at least for now).
     msal_kwargs = {}
     if options:
-        # For VM SSH. 'data' support is a CLI-specific extension.
-        # SDK doesn't support 'data': https://github.com/Azure/azure-sdk-for-python/pull/16397
-        if 'data' in options:
-            msal_kwargs['data'] = options['data']
         # For CAE
         if 'claims' in options:
             msal_kwargs['claims_challenge'] = options['claims']

src/azure-cli-core/azure/cli/core/auth/tests/test_credential_adaptor.py

@@ -44,7 +44,7 @@ def _now_timestamp_mock():
 
 class TestCredentialAdaptor(unittest.TestCase):
 
-    @mock.patch('azure.cli.core.auth.util._now_timestamp', new=_now_timestamp_mock)
+    @mock.patch('azure.cli.core.auth.util.now_timestamp', new=_now_timestamp_mock)
     def test_get_token(self):
         msal_cred = MsalCredentialStub()
         sdk_cred = CredentialAdaptor(msal_cred)
@@ -56,15 +56,11 @@ def test_get_token(self):
         assert access_token.token == MOCK_ACCESS_TOKEN
         assert access_token.expires_on == 1630920323
 
-        # Note that SDK doesn't support 'data'. This is a CLI-specific extension.
-        sdk_cred.get_token('https://management.core.windows.net//.default', data=MOCK_DATA)
-        assert msal_cred.acquire_token_kwargs['data'] == MOCK_DATA
-
         sdk_cred.get_token('https://management.core.windows.net//.default', claims=MOCK_CLAIMS)
         assert msal_cred.acquire_token_claims_challenge == MOCK_CLAIMS
 
 
-    @mock.patch('azure.cli.core.auth.util._now_timestamp', new=_now_timestamp_mock)
+    @mock.patch('azure.cli.core.auth.util.now_timestamp', new=_now_timestamp_mock)
     def test_get_token_info(self):
         msal_cred = MsalCredentialStub()
         sdk_cred = CredentialAdaptor(msal_cred)
@@ -78,10 +74,6 @@ def test_get_token_info(self):
 
         assert msal_cred.acquire_token_scopes == ['https://management.core.windows.net//.default']
 
-        # Note that SDK doesn't support 'data'. If 'data' were supported, it should be tested with:
-        sdk_cred.get_token_info('https://management.core.windows.net//.default', options={'data': MOCK_DATA})
-        assert msal_cred.acquire_token_kwargs['data'] == MOCK_DATA
-
         sdk_cred.get_token_info('https://management.core.windows.net//.default', options={'claims': MOCK_CLAIMS})
         assert msal_cred.acquire_token_claims_challenge == MOCK_CLAIMS
 

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

@@ -50,10 +50,12 @@ def __init__(self, *args, **kwargs):
         # If acquire_token_scopes is checked, make sure to create a new instance of MsalCredentialStub
         # to avoid interference from other tests.
         self.acquire_token_scopes = None
+        self.acquire_token_data=None
         super().__init__()
 
     def acquire_token(self, scopes, **kwargs):
         self.acquire_token_scopes = scopes
+        self.acquire_token_data = kwargs.get('data')
         return {
             'access_token': MOCK_ACCESS_TOKEN,
             'token_type': 'Bearer',
@@ -1287,6 +1289,31 @@ def cloud_shell_credential_factory():
         with self.assertRaisesRegex(CLIError, 'Cloud Shell'):
             profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id)
 
+    @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential')
+    def test_get_msal_token(self, get_user_credential_mock):
+        credential_mock_temp = MsalCredentialStub()
+        get_user_credential_mock.return_value = credential_mock_temp
+        cli = DummyCli()
+
+        storage_mock = {'subscriptions': None}
+        profile = Profile(cli_ctx=cli, storage=storage_mock)
+        consolidated = profile._normalize_properties(self.user1,
+                                                     [self.subscription1],
+                                                     False, None, None)
+        profile._set_subscriptions(consolidated)
+
+        MOCK_DATA = {
+            'key_id': 'test',
+            'req_cnf': 'test',
+            'token_type': 'ssh-cert'
+        }
+        result = profile.get_msal_token(['https://pas.windows.net/CheckMyAccess/Linux/.default'],
+                                        MOCK_DATA)
+
+        assert result == (None, MOCK_ACCESS_TOKEN)
+        assert credential_mock_temp.acquire_token_scopes == ['https://pas.windows.net/CheckMyAccess/Linux/.default']
+        assert credential_mock_temp.acquire_token_data == MOCK_DATA
+
     @mock.patch('azure.cli.core.auth.identity.Identity.logout_service_principal')
     @mock.patch('azure.cli.core.auth.identity.Identity.logout_user')
     def test_logout(self, logout_user_mock, logout_service_principal_mock):