{Core} Decouple get_raw_token from SDK token protocol (#31063)

Author: jiasli

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

@@ -365,7 +365,7 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
             if tenant:
                 raise CLIError("Tenant shouldn't be specified for Cloud Shell account")
             from .auth.msal_credentials import CloudShellCredential
-            sdk_cred = CredentialAdaptor(CloudShellCredential())
+            cred = CloudShellCredential()
 
         elif managed_identity_type:
             # managed identity
@@ -374,25 +374,27 @@ def get_raw_token(self, resource=None, scopes=None, subscription=None, tenant=No
             cred = ManagedIdentityAuth.credential_factory(managed_identity_type, managed_identity_id)
             if credential_out:
                 credential_out['credential'] = cred
-            sdk_cred = CredentialAdaptor(cred)
 
         else:
-            sdk_cred = CredentialAdaptor(self._create_credential(account, tenant_id=tenant))
+            cred = self._create_credential(account, tenant_id=tenant)
 
-        sdk_token = sdk_cred.get_token(*scopes)
+        msal_token = cred.acquire_token(scopes)
         # Convert epoch int 'expires_on' to datetime string 'expiresOn' for backward compatibility
         # WARNING: expiresOn is deprecated and will be removed in future release.
         import datetime
-        expiresOn = datetime.datetime.fromtimestamp(sdk_token.expires_on).strftime("%Y-%m-%d %H:%M:%S.%f")
+        from .auth.util import now_timestamp
+        from .auth.constants import EXPIRES_IN, ACCESS_TOKEN
+        expires_on = now_timestamp() + msal_token[EXPIRES_IN]
+        expiresOn = datetime.datetime.fromtimestamp(expires_on).strftime("%Y-%m-%d %H:%M:%S.%f")
 
         token_entry = {
-            'accessToken': sdk_token.token,
-            'expires_on': sdk_token.expires_on,  # epoch int, like 1605238724
+            'accessToken': msal_token[ACCESS_TOKEN],
+            'expires_on': expires_on,  # epoch int, like 1605238724
             'expiresOn': expiresOn  # datetime string, like "2020-11-12 13:50:47.114324"
         }
 
         # Build a tuple of (token_type, token, token_entry)
-        token_tuple = 'Bearer', sdk_token.token, token_entry
+        token_tuple = 'Bearer', msal_token[ACCESS_TOKEN], token_entry
 
         # Return a tuple of (token_tuple, subscription, tenant)
         return (token_tuple,

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

@@ -76,7 +76,7 @@ def _build_sdk_access_token_info(token_entry):
     #     'token_source': 'cache'
     # }
     from .constants import ACCESS_TOKEN, EXPIRES_IN
-    from .util import _now_timestamp
+    from .util import now_timestamp
     from azure.core.credentials import AccessTokenInfo
 
-    return AccessTokenInfo(token_entry[ACCESS_TOKEN], _now_timestamp() + token_entry[EXPIRES_IN])
+    return AccessTokenInfo(token_entry[ACCESS_TOKEN], now_timestamp() + token_entry[EXPIRES_IN])

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

@@ -151,7 +151,7 @@ def build_sdk_access_token(token_entry):
     # This can slow down commands that doesn't need azure.core, like `az account get-access-token`.
     # So We define our own AccessToken.
     from .constants import ACCESS_TOKEN, EXPIRES_IN
-    return AccessToken(token_entry[ACCESS_TOKEN], _now_timestamp() + token_entry[EXPIRES_IN])
+    return AccessToken(token_entry[ACCESS_TOKEN], now_timestamp() + token_entry[EXPIRES_IN])
 
 
 def decode_access_token(access_token):
@@ -177,6 +177,6 @@ def read_response_templates():
     return success_template, error_template
 
 
-def _now_timestamp():
+def now_timestamp():
     import time
     return int(time.time())

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

@@ -39,7 +39,7 @@ def _build_test_jwt(claims):
     return '.'.join(base64.urlsafe_b64encode(p.encode('utf-8')).decode('utf-8').replace('=', '') for p in parts)
 
 
-def _now_timestamp_mock():
+def now_timestamp_mock():
     # 2021-09-06 08:55:23
     return 1630918523
 
@@ -1013,7 +1013,7 @@ def test_get_login_credentials_mi_user_assigned_resource_id(self):
         assert cred._credential.object_id is None
         assert cred._credential.resource_id == self.test_mi_resource_id
 
-    @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)
     @mock.patch('azure.cli.core.auth.identity.Identity.get_user_credential')
     def test_get_raw_token(self, get_user_credential_mock):
         credential_mock_temp = MsalCredentialStub()
@@ -1061,7 +1061,7 @@ def test_get_raw_token(self, get_user_credential_mock):
         self.assertIsNone(sub)
         self.assertEqual(tenant, self.tenant_id)
 
-    @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)
     @mock.patch('azure.cli.core.auth.identity.Identity.get_service_principal_credential')
     def test_get_raw_token_for_sp(self, get_service_principal_credential_mock):
         credential_mock_temp = MsalCredentialStub()
@@ -1102,7 +1102,7 @@ def test_get_raw_token_for_sp(self, get_service_principal_credential_mock):
         self.assertIsNone(sub)
         self.assertEqual(tenant, self.tenant_id)
 
-    @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)
     @mock.patch('azure.cli.core.auth.msal_credentials.ManagedIdentityCredential', ManagedIdentityCredentialStub)
     def test_get_raw_token_mi_system_assigned(self):
         profile = Profile(cli_ctx=DummyCli(), storage={'subscriptions': None})
@@ -1136,7 +1136,7 @@ def test_get_raw_token_mi_system_assigned(self):
         with self.assertRaisesRegex(CLIError, "Tenant shouldn't be specified"):
             cred, subscription_id, _ = profile.get_raw_token(resource='http://test_resource', tenant=self.tenant_id)
 
-    @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)
     @mock.patch('azure.cli.core.auth.msal_credentials.ManagedIdentityCredential', ManagedIdentityCredentialStub)
     def test_get_raw_token_mi_user_assigned_client_id(self):
         profile = Profile(cli_ctx=DummyCli(), storage={'subscriptions': None})
@@ -1167,7 +1167,7 @@ def test_get_raw_token_mi_user_assigned_client_id(self):
         self.assertEqual(subscription_id, self.test_mi_subscription_id)
         self.assertEqual(tenant_id, self.test_mi_tenant)
 
-    @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)
     @mock.patch('azure.cli.core.auth.msal_credentials.ManagedIdentityCredential', ManagedIdentityCredentialStub)
     def test_get_raw_token_mi_user_assigned_object_id(self):
         profile = Profile(cli_ctx=DummyCli(), storage={'subscriptions': None})
@@ -1198,7 +1198,7 @@ def test_get_raw_token_mi_user_assigned_object_id(self):
         self.assertEqual(subscription_id, self.test_mi_subscription_id)
         self.assertEqual(tenant_id, self.test_mi_tenant)
 
-    @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)
     @mock.patch('azure.cli.core.auth.msal_credentials.ManagedIdentityCredential', ManagedIdentityCredentialStub)
     def test_get_raw_token_mi_user_assigned_resource_id(self):
         profile = Profile(cli_ctx=DummyCli(), storage={'subscriptions': None})
@@ -1229,7 +1229,7 @@ def test_get_raw_token_mi_user_assigned_resource_id(self):
         self.assertEqual(subscription_id, self.test_mi_subscription_id)
         self.assertEqual(tenant_id, self.test_mi_tenant)
 
-    @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)
     @mock.patch('azure.cli.core._profile.in_cloud_console', autospec=True)
     @mock.patch('azure.cli.core.auth.msal_credentials.CloudShellCredential', autospec=True)
     def test_get_raw_token_in_cloud_shell(self, cloud_shell_credential_mock, mock_in_cloud_console):