[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 <[email protected]>
Azure · Jan 23, 2024 · 501cb2c · 501cb2c
1 parent 370ecaa
commit 501cb2c
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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"""