From 501cb2c18cf2289ad903e16113f1739fd38981b0 Mon Sep 17 00:00:00 2001 From: Paul Van Eck Date: Tue, 23 Jan 2024 13:34:03 -0800 Subject: [PATCH] [Identity] Support expires_on in AzureCLICredential (#33947) Newer versions of Azure CLI now also return a Unix timestamp with the `expires_on` field when retrieving an access token. We should prefer using that. Signed-off-by: Paul Van Eck --- sdk/identity/azure-identity/CHANGELOG.md | 2 + .../azure/identity/_credentials/azure_cli.py | 14 +++--- .../tests/test_cli_credential.py | 42 ++++++++++++++++++ .../tests/test_cli_credential_async.py | 44 +++++++++++++++++++ 4 files changed, 94 insertions(+), 8 deletions(-) diff --git a/sdk/identity/azure-identity/CHANGELOG.md b/sdk/identity/azure-identity/CHANGELOG.md index 0b3182ef5914..ee312eba22d5 100644 --- a/sdk/identity/azure-identity/CHANGELOG.md +++ b/sdk/identity/azure-identity/CHANGELOG.md @@ -13,6 +13,8 @@ ### Other Changes +- `AzureCliCredential` utilizes the new `expires_on` property returned by `az` CLI versions >= 2.54.0 to determine token expiration. ([#33947](https://github.com/Azure/azure-sdk-for-python/issues/33947)) + ## 1.15.0 (2023-10-26) ### Features Added diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py b/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py index bde6eb3ecc49..efbf9d343f85 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py @@ -9,7 +9,6 @@ import shutil import subprocess import sys -import time from typing import List, Optional, Any, Dict from azure.core.credentials import AccessToken @@ -139,14 +138,13 @@ def parse_token(output) -> Optional[AccessToken]: """ try: token = json.loads(output) - dt = datetime.strptime(token["expiresOn"], "%Y-%m-%d %H:%M:%S.%f") - if hasattr(dt, "timestamp"): - # Python >= 3.3 - expires_on = dt.timestamp() - else: - # taken from Python 3.5's datetime.timestamp() - expires_on = time.mktime((dt.year, dt.month, dt.day, dt.hour, dt.minute, dt.second, -1, -1, -1)) + # Use "expires_on" if it's present, otherwise use "expiresOn". + if "expires_on" in token: + return AccessToken(token["accessToken"], int(token["expires_on"])) + + dt = datetime.strptime(token["expiresOn"], "%Y-%m-%d %H:%M:%S.%f") + expires_on = dt.timestamp() return AccessToken(token["accessToken"], int(expires_on)) except (KeyError, ValueError): return None diff --git a/sdk/identity/azure-identity/tests/test_cli_credential.py b/sdk/identity/azure-identity/tests/test_cli_credential.py index e2697764b03f..0b688df36faf 100644 --- a/sdk/identity/azure-identity/tests/test_cli_credential.py +++ b/sdk/identity/azure-identity/tests/test_cli_credential.py @@ -92,6 +92,48 @@ def test_get_token(): assert token.expires_on == expected_expires_on +def test_expires_on_used(): + """Test that 'expires_on' is preferred over 'expiresOn'.""" + expires_on = 1602015811 + successful_output = json.dumps( + { + "expiresOn": datetime.fromtimestamp(1555555555).strftime("%Y-%m-%d %H:%M:%S.%f"), + "expires_on": expires_on, + "accessToken": "access token", + "subscription": "some-guid", + "tenant": "some-guid", + "tokenType": "Bearer", + } + ) + + with mock.patch("shutil.which", return_value="az"): + with mock.patch(CHECK_OUTPUT, mock.Mock(return_value=successful_output)): + token = AzureCliCredential().get_token("scope") + + assert token.expires_on == expires_on + + +def test_expires_on_string(): + """Test that 'expires_on' still works if it's a string.""" + expires_on = 1602015811 + successful_output = json.dumps( + { + "expires_on": f"{expires_on}", + "accessToken": "access token", + "subscription": "some-guid", + "tenant": "some-guid", + "tokenType": "Bearer", + } + ) + + with mock.patch("shutil.which", return_value="az"): + with mock.patch(CHECK_OUTPUT, mock.Mock(return_value=successful_output)): + token = AzureCliCredential().get_token("scope") + + assert type(token.expires_on) == int + assert token.expires_on == expires_on + + def test_cli_not_installed(): """The credential should raise CredentialUnavailableError when the CLI isn't installed""" with mock.patch("shutil.which", return_value=None): diff --git a/sdk/identity/azure-identity/tests/test_cli_credential_async.py b/sdk/identity/azure-identity/tests/test_cli_credential_async.py index 28d3a72310eb..8920263d4b83 100644 --- a/sdk/identity/azure-identity/tests/test_cli_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_cli_credential_async.py @@ -119,6 +119,50 @@ async def test_get_token(): assert token.expires_on == expected_expires_on +async def test_expires_on_used(): + """Test that 'expires_on' is preferred over 'expiresOn'.""" + expires_on = 1602015811 + successful_output = json.dumps( + { + "expiresOn": datetime.fromtimestamp(1555555555).strftime("%Y-%m-%d %H:%M:%S.%f"), + "expires_on": expires_on, + "accessToken": "access token", + "subscription": "some-guid", + "tenant": "some-guid", + "tokenType": "Bearer", + } + ) + + with mock.patch("shutil.which", return_value="az"): + with mock.patch(SUBPROCESS_EXEC, mock_exec(successful_output)): + credential = AzureCliCredential() + token = await credential.get_token("scope") + + assert token.expires_on == expires_on + + +async def test_expires_on_string(): + """Test that 'expires_on' still works if it's a string.""" + expires_on = 1602015811 + successful_output = json.dumps( + { + "expires_on": f"{expires_on}", + "accessToken": "access token", + "subscription": "some-guid", + "tenant": "some-guid", + "tokenType": "Bearer", + } + ) + + with mock.patch("shutil.which", return_value="az"): + with mock.patch(SUBPROCESS_EXEC, mock_exec(successful_output)): + credential = AzureCliCredential() + token = await credential.get_token("scope") + + assert type(token.expires_on) == int + assert token.expires_on == expires_on + + async def test_cli_not_installed(): """The credential should raise CredentialUnavailableError when the CLI isn't installed"""