|
2 | 2 | # All rights reserved. |
3 | 3 | # |
4 | 4 | # This code is licensed under the MIT License. |
| 5 | +import calendar |
| 6 | +import datetime |
5 | 7 | import json |
6 | 8 | import logging |
7 | 9 | import os |
| 10 | +import re |
8 | 11 | import socket |
9 | 12 | import sys |
10 | 13 | import time |
@@ -432,6 +435,35 @@ def _obtain_token(http_client, managed_identity, resource): |
432 | 435 | return _obtain_token_on_azure_vm(http_client, managed_identity, resource) |
433 | 436 |
|
434 | 437 |
|
| 438 | +def _parse_expires_on(raw: str) -> int: |
| 439 | + try: |
| 440 | + return int(raw) # It is typically an epoch time |
| 441 | + except ValueError: |
| 442 | + pass |
| 443 | + try: |
| 444 | + # '2024-10-18T19:51:37.0000000+00:00' was observed in |
| 445 | + # https://github.com/AzureAD/microsoft-authentication-library-for-dotnet/issues/4963 |
| 446 | + if sys.version_info < (3, 11): # Does not support 7-digit microseconds |
| 447 | + raw = re.sub( # Strip microseconds portion using regex |
| 448 | + r'(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(\.\d+)([+-]\d{2}:\d{2})', |
| 449 | + r'\1\3', |
| 450 | + raw) |
| 451 | + return int(datetime.datetime.fromisoformat(raw).timestamp()) |
| 452 | + except ValueError: |
| 453 | + pass |
| 454 | + for format in ( |
| 455 | + "%m/%d/%Y %H:%M:%S %z", # Support "06/20/2019 02:57:58 +00:00" |
| 456 | + # 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 |
| 457 | + "%m/%d/%Y %I:%M:%S %p %z", # Support "1/16/2020 12:0:12 AM +00:00" |
| 458 | + # 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 |
| 459 | + ): |
| 460 | + try: |
| 461 | + return calendar.timegm(time.strptime(raw, format)) |
| 462 | + except ValueError: |
| 463 | + pass |
| 464 | + raise ManagedIdentityError(f"Cannot parse expires_on: {raw}") |
| 465 | + |
| 466 | + |
435 | 467 | def _adjust_param(params, managed_identity, types_mapping=None): |
436 | 468 | # Modify the params dict in place |
437 | 469 | id_name = (types_mapping or ManagedIdentity._types_mapping).get( |
@@ -504,7 +536,7 @@ def _obtain_token_on_app_service( |
504 | 536 | if payload.get("access_token") and payload.get("expires_on"): |
505 | 537 | return { # Normalizing the payload into OAuth2 format |
506 | 538 | "access_token": payload["access_token"], |
507 | | - "expires_in": int(payload["expires_on"]) - int(time.time()), |
| 539 | + "expires_in": _parse_expires_on(payload["expires_on"]) - int(time.time()), |
508 | 540 | "resource": payload.get("resource"), |
509 | 541 | "token_type": payload.get("token_type", "Bearer"), |
510 | 542 | } |
@@ -538,7 +570,7 @@ def _obtain_token_on_machine_learning( |
538 | 570 | if payload.get("access_token") and payload.get("expires_on"): |
539 | 571 | return { # Normalizing the payload into OAuth2 format |
540 | 572 | "access_token": payload["access_token"], |
541 | | - "expires_in": int(payload["expires_on"]) - int(time.time()), |
| 573 | + "expires_in": _parse_expires_on(payload["expires_on"]) - int(time.time()), |
542 | 574 | "resource": payload.get("resource"), |
543 | 575 | "token_type": payload.get("token_type", "Bearer"), |
544 | 576 | } |
|
0 commit comments