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]))
 
@@ -401,6 +401,15 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
                 None if tenant else str(account[_SUBSCRIPTION_ID]),
                 str(tenant if tenant else account[_TENANT_ID]))
 
+    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(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
+
     def _normalize_properties(self, user, subscriptions, is_service_principal, cert_sn_issuer_auth=None,
                               assigned_identity_info=None):
         consolidated = []

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/msal_credentials.py

@@ -43,16 +43,23 @@ def __init__(self, client_id, username, **kwargs):
 
         self._account = accounts[0]
 
-    def acquire_token(self, scopes, claims_challenge=None, **kwargs):
+    def acquire_token(self, scopes, claims_challenge=None, data=None, **kwargs):
         # scopes must be a list.
         # For acquiring SSH certificate, scopes is ['https://pas.windows.net/CheckMyAccess/Linux/.default']
+        # data is only used for acquiring VM SSH certificate. DO NOT use it for other purposes.
         # kwargs is already sanitized by CredentialAdaptor, so it can be safely passed to MSAL
         logger.debug("UserCredential.acquire_token: scopes=%r, claims_challenge=%r, kwargs=%r",
-                     scopes, claims_challenge, kwargs)
+                     scopes, claims_challenge)
 
         if claims_challenge:
             logger.warning('Acquiring new access token silently for tenant %s with claims challenge: %s',
                            self._msal_app.authority.tenant, claims_challenge)
+
+        # Only pass data to MSAL if it is set. Passing data=None will cause failure in MSAL:
+        #   AttributeError: 'NoneType' object has no attribute 'get'
+        if data is not None:
+            kwargs['data'] = data
+
         result = self._msal_app.acquire_token_silent_with_error(
             scopes, self._account, claims_challenge=claims_challenge, **kwargs)
 
@@ -105,8 +112,12 @@ def __init__(self, client_id, client_credential, **kwargs):
         """
         self._msal_app = ConfidentialClientApplication(client_id, client_credential=client_credential, **kwargs)
 
-    def acquire_token(self, scopes, **kwargs):
+    def acquire_token(self, scopes, data=None, **kwargs):
         logger.debug("ServicePrincipalCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
+
+        if data is not None:
+            kwargs['data'] = data
+
         result = self._msal_app.acquire_token_for_client(scopes, **kwargs)
         check_result(result)
         return result
@@ -126,8 +137,12 @@ def __init__(self):
             #   token_cache=...
         )
 
-    def acquire_token(self, scopes, **kwargs):
+    def acquire_token(self, scopes, data=None, **kwargs):
         logger.debug("CloudShellCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
+
+        if data is not None:
+            kwargs['data'] = data
+
         result = self._msal_app.acquire_token_interactive(scopes, prompt="none", **kwargs)
         check_result(result, scopes=scopes)
         return result
@@ -147,9 +162,13 @@ def __init__(self, client_id=None, resource_id=None, object_id=None):
             managed_identity = SystemAssignedManagedIdentity()
         self._msal_client = ManagedIdentityClient(managed_identity, http_client=requests.Session())
 
-    def acquire_token(self, scopes, **kwargs):
+    def acquire_token(self, scopes, data=None, **kwargs):
         logger.debug("ManagedIdentityCredential.acquire_token: scopes=%r, kwargs=%r", scopes, kwargs)
 
+        if data is not None:
+            from azure.cli.core.azclierror import AuthenticationError
+            raise AuthenticationError("VM SSH currently doesn't support managed identity.")
+
         from .util import scopes_to_resource
         result = self._msal_client.acquire_token_for_client(resource=scopes_to_resource(scopes))
         check_result(result)

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):