Skip to content

Commit

Permalink
auth azure: support access token
Browse files Browse the repository at this point in the history
  • Loading branch information
LiliDeng committed Dec 25, 2024
1 parent f2014f6 commit af30a30
Showing 1 changed file with 38 additions and 3 deletions.
41 changes: 38 additions & 3 deletions lisa/sut_orchestrator/azure/platform_.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import os
import re
import sys
import time
from copy import deepcopy
from dataclasses import InitVar, dataclass, field
from datetime import datetime
Expand All @@ -17,6 +18,7 @@
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union, cast

import requests
from azure.core.credentials import AccessToken, TokenCredential
from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
from azure.identity import DefaultAzureCredential
from azure.mgmt.compute.models import (
Expand Down Expand Up @@ -400,6 +402,29 @@ def cloud(self, value: Optional[CloudSchema]) -> None:
self.cloud_raw = value.to_dict() # type: ignore


class StaticAccessTokenCredential(TokenCredential):
def __init__(self, token: str, expires_on: int) -> None:
"""
Initialize StaticAccessTokenCredential with the provided token and expiry time.
:param token: The Azure access token as a string.
:param expires_on: The expiry time of the token as an integer (Unix timestamp).
"""
self._token = token
self._expires_on = expires_on

def get_token(self, *scopes: str, **kwargs) -> AccessToken:
"""
Get the access token for the specified scopes.
:param scopes: The OAuth 2.0 scopes the token applies to.
:param kwargs: Additional keyword arguments that may be required by the SDK.
:return: An AccessToken instance containing the token and its expiry time.
"""
# You can choose to print or log the scopes and kwargs for debugging if needed
return AccessToken(self._token, self._expires_on)


class AzurePlatform(Platform):
_diagnostic_storage_container_pattern = re.compile(
r"(https:\/\/)(?P<storage_name>.*)([.].*){4}\/(?P<container_name>.*)\/",
Expand Down Expand Up @@ -937,9 +962,19 @@ def _initialize_credential(self) -> None:
if azure_runbook.service_principal_key:
os.environ["AZURE_CLIENT_SECRET"] = azure_runbook.service_principal_key

credential = DefaultAzureCredential(
authority=self.cloud.endpoints.active_directory,
)
if "AZURE_ACCESS_TOKEN" in os.environ:
token = os.environ["AZURE_ACCESS_TOKEN"]
else:
token = None

if token:
credential = StaticAccessTokenCredential(
token, int(time.time()) + 3600 * 24
)
else:
credential = DefaultAzureCredential(
authority=self.cloud.endpoints.active_directory,
)

with SubscriptionClient(
credential,
Expand Down

0 comments on commit af30a30

Please sign in to comment.