Support ISO8601-ish format on expires_on in MI

Author: rayluo

msal/managed_identity.py

@@ -2,6 +2,7 @@
 # All rights reserved.
 #
 # This code is licensed under the MIT License.
+import datetime
 import json
 import logging
 import os
@@ -432,6 +433,29 @@ def _obtain_token(http_client, managed_identity, resource):
     return _obtain_token_on_azure_vm(http_client, managed_identity, resource)
 
 
+def _parse_expires_on(raw: str) -> int:
+    try:
+        return int(raw)  # It is typically an epoch time
+    except ValueError:
+        pass
+    try:
+        # '2024-10-18T19:51:37.0000000+00:00' was observed in
+        # https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/4963
+        # We support format(s) YYYY-MM-DD[*HH[:MM[:SS[.fff[fff]]]][+HH:MM[:SS[.ffffff]]]]
+        return int(datetime.datetime.fromisoformat(raw).timestamp())
+    except ValueError:
+        pass
+    for format in (
+        "%m/%d/%Y %H:%M:%S +00:00",  # Derived from https://github.com/Azure/azure-sdk-for-python/blob/azure-identity_1.21.0/sdk/identity/azure-identity/azure/identity/_credentials/azure_ml.py#L52
+        "%m/%d/%Y %H:%M:%S %p +00:00",  # Derived from https://github.com/Azure/azure-sdk-for-python/blob/azure-identity_1.21.0/sdk/identity/azure-identity/azure/identity/_credentials/azure_ml.py#L51
+    ):
+        try:
+            return int(time.mktime(time.strptime(raw, format)))
+        except ValueError:
+            pass
+    raise ManagedIdentityError(f"Cannot parse expires_on: {raw}")
+
+
 def _adjust_param(params, managed_identity, types_mapping=None):
     # Modify the params dict in place
     id_name = (types_mapping or ManagedIdentity._types_mapping).get(
@@ -504,7 +528,7 @@ def _obtain_token_on_app_service(
         if payload.get("access_token") and payload.get("expires_on"):
             return {  # Normalizing the payload into OAuth2 format
                 "access_token": payload["access_token"],
-                "expires_in": int(payload["expires_on"]) - int(time.time()),
+                "expires_in": _parse_expires_on(payload["expires_on"]) - int(time.time()),
                 "resource": payload.get("resource"),
                 "token_type": payload.get("token_type", "Bearer"),
                 }
@@ -538,7 +562,7 @@ def _obtain_token_on_machine_learning(
         if payload.get("access_token") and payload.get("expires_on"):
             return {  # Normalizing the payload into OAuth2 format
                 "access_token": payload["access_token"],
-                "expires_in": int(payload["expires_on"]) - int(time.time()),
+                "expires_in": _parse_expires_on(payload["expires_on"]) - int(time.time()),
                 "resource": payload.get("resource"),
                 "token_type": payload.get("token_type", "Bearer"),
                 }

tests/test_mi.py

@@ -25,6 +25,7 @@
     MACHINE_LEARNING,
     SERVICE_FABRIC,
     DEFAULT_TO_VM,
+    _parse_expires_on,
 )
 from msal.token_cache import is_subdict_of
 
@@ -49,6 +50,18 @@ def test_helper_class_should_be_interchangable_with_dict_which_could_be_loaded_f
             {"ManagedIdentityIdType": "SystemAssigned", "Id": None})
 
 
+class ExpiresOnTestCase(unittest.TestCase):
+    def test_expires_on_parsing(self):
+        for input, expected in {
+            "1234567890": 1234567890,
+            "1970-01-01T00:00:12.0000000+00:00": 12,
+            "2024-10-18T19:51:37.0000000+00:00": 1729281097,  # Copied from https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/4963
+            "1/16/2020 5:24:12 AM +00:00": 1579181052,  # Derived from https://github.com/Azure/azure-sdk-for-python/blob/azure-identity_1.21.0/sdk/identity/azure-identity/azure/identity/_credentials/azure_ml.py#L51
+            "06/20/2019 02:57:58 +00:00": 1561024678,  # Derived from https://github.com/Azure/azure-sdk-for-python/blob/azure-identity_1.21.0/sdk/identity/azure-identity/azure/identity/_credentials/azure_ml.py#L52
+        }.items():
+            self.assertEqual(_parse_expires_on(input), expected, f"Should parse {input} to {expected}")
+
+
 class ClientTestCase(unittest.TestCase):
     maxDiff = None