diff --git a/src/azure-cli-core/azure/cli/core/_profile.py b/src/azure-cli-core/azure/cli/core/_profile.py index 6cff8040708..0c4583c078b 100644 --- a/src/azure-cli-core/azure/cli/core/_profile.py +++ b/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, diff --git a/src/azure-cli-core/azure/cli/core/auth/credential_adaptor.py b/src/azure-cli-core/azure/cli/core/auth/credential_adaptor.py index 69365cb45a8..1a2ac83f24f 100644 --- a/src/azure-cli-core/azure/cli/core/auth/credential_adaptor.py +++ b/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]) diff --git a/src/azure-cli-core/azure/cli/core/auth/util.py b/src/azure-cli-core/azure/cli/core/auth/util.py index ecbe67627d5..1bbd328808f 100644 --- a/src/azure-cli-core/azure/cli/core/auth/util.py +++ b/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()) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_profile.py b/src/azure-cli-core/azure/cli/core/tests/test_profile.py index 2cb6c04bb01..c5e7d70a815 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_profile.py +++ b/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):