get_token_info

Author: jiasli

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

@@ -24,15 +24,18 @@ def __init__(self, credential, auxiliary_credentials=None):
         self._auxiliary_credentials = auxiliary_credentials
 
     def get_token(self, *scopes, **kwargs):
-        """Get an access token from the main credential."""
+        """Old SDK token protocol. Get an access token from the main credential."""
         logger.debug("CredentialAdaptor.get_token: scopes=%r, kwargs=%r", scopes, kwargs)
 
-        # Discard unsupported kwargs: tenant_id, enable_cae
-        filtered_kwargs = {}
-        if 'data' in kwargs:
-            filtered_kwargs['data'] = kwargs['data']
+        msal_kwargs = _prepare_msal_kwargs(kwargs)
+        return build_sdk_access_token(self._credential.acquire_token(list(scopes), **msal_kwargs))
 
-        return build_sdk_access_token(self._credential.acquire_token(list(scopes), **filtered_kwargs))
+    def get_token_info(self, *scopes, options=None):
+        """New SDK token protocol. Get an access token from the main credential."""
+        logger.debug("CredentialAdaptor.get_token_info: scopes=%r, options=%r", scopes, options)
+
+        msal_kwargs = _prepare_msal_kwargs(options)
+        return _build_sdk_access_token_info(self._credential.acquire_token(list(scopes), **msal_kwargs))
 
     def get_auxiliary_tokens(self, *scopes, **kwargs):
         """Get access tokens from auxiliary credentials."""
@@ -41,3 +44,32 @@ def get_auxiliary_tokens(self, *scopes, **kwargs):
             return [build_sdk_access_token(cred.acquire_token(list(scopes), **kwargs))
                     for cred in self._auxiliary_credentials]
         return None
+
+
+def _prepare_msal_kwargs(options=None):
+    # Preserve supported options and discard unsupported options (tenant_id, enable_cae).
+    # 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. TokenRequestOptions doesn't support 'data'.
+        if 'data' in options:
+            msal_kwargs['data'] = options['data']
+        # For CAE
+        if 'claims' in options:
+            msal_kwargs['claims'] = options['claims']
+    return msal_kwargs
+
+
+def _build_sdk_access_token_info(token_entry):
+    # MSAL token entry sample:
+    # {
+    #     'access_token': 'eyJ0eXAiOiJKV...',
+    #     'token_type': 'Bearer',
+    #     'expires_in': 1618,
+    #     'token_source': 'cache'
+    # }
+    from .constants import ACCESS_TOKEN, EXPIRES_IN
+    from .util import _now_timestamp
+    from azure.core.credentials import AccessTokenInfo
+
+    return AccessTokenInfo(token_entry[ACCESS_TOKEN], _now_timestamp() + token_entry[EXPIRES_IN])

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

@@ -0,0 +1,84 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+
+import unittest
+from collections import namedtuple
+from unittest import mock
+
+from ..credential_adaptor import CredentialAdaptor
+
+
+MOCK_ACCESS_TOKEN = "mock_access_token"
+MOCK_DATA = {"test_data": "value1"}
+MOCK_CLAIMS = {"test_claims": "value2"}
+
+class MsalCredentialStub:
+
+    def __init__(self, *args, **kwargs):
+        self.acquire_token_scopes = None
+        self.acquire_token_kwargs = None
+        super().__init__()
+
+    def acquire_token(self, scopes, **kwargs):
+        self.acquire_token_scopes = scopes
+        self.acquire_token_kwargs = kwargs
+        return {
+            'access_token': MOCK_ACCESS_TOKEN,
+            'token_type': 'Bearer',
+            'expires_in': 1800,
+            'token_source': 'cache'
+        }
+
+def _now_timestamp_mock():
+    # 2021-09-06 08:55:23
+    return 1630918523
+
+
+class TestCredentialAdaptor(unittest.TestCase):
+
+    @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)
+        access_token = sdk_cred.get_token('https://management.core.windows.net//.default')
+        assert msal_cred.acquire_token_scopes == ['https://management.core.windows.net//.default']
+
+        from ..util import AccessToken
+        assert isinstance(access_token, AccessToken)
+        assert access_token.token == MOCK_ACCESS_TOKEN
+        assert access_token.expires_on == 1630920323
+
+        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_kwargs['claims'] == MOCK_CLAIMS
+
+
+    @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)
+        access_token_info = sdk_cred.get_token_info('https://management.core.windows.net//.default')
+
+        from azure.core.credentials import AccessTokenInfo
+        assert isinstance(access_token_info, AccessTokenInfo)
+        assert access_token_info.token == MOCK_ACCESS_TOKEN
+        assert access_token_info.expires_on == 1630920323
+        assert access_token_info.token_type == 'Bearer'
+
+        assert msal_cred.acquire_token_scopes == ['https://management.core.windows.net//.default']
+
+        # Actually, TokenRequestOptions doesn't support 'data'.
+        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_kwargs['claims'] == MOCK_CLAIMS
+
+
+if __name__ == '__main__':
+    unittest.main()