Skip to content

{Core} Decouple get_raw_token from SDK token protocol #31063

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jiasli merged 1 commit into from
Jun 23, 2025
Merged

{Core} Decouple get_raw_token from SDK token protocol #31063

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 10 additions & 8 deletions src/azure-cli-core/azure/cli/core/_profile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 = {
Copy link
Member Author

@jiasli jiasli Jun 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

token_entry is actually a remnant of ADAL, but replacing it with msal_token will be a breaking change.

msal_token contains:

{
    "access_token": "...",
    "token_type": "Bearer",
    "expires_in": 4678,
    "token_source": "cache"
}

'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,
Expand Down
4 changes: 2 additions & 2 deletions src/azure-cli-core/azure/cli/core/auth/credential_adaptor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
4 changes: 2 additions & 2 deletions src/azure-cli-core/azure/cli/core/auth/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -177,6 +177,6 @@ def read_response_templates():
return success_template, error_template


def _now_timestamp():
def now_timestamp():
Copy link

Copilot AI Jun 10, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Insert an additional blank line above this function so there are two blank lines between top-level definitions, per the style guide.

Copilot uses AI. Check for mistakes.
Copy link
Member Author

@jiasli jiasli Jun 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are already 2 blank lines.

import time
return int(time.time())
16 changes: 8 additions & 8 deletions src/azure-cli-core/azure/cli/core/tests/test_profile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -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):
Expand Down
Loading