diff --git a/sdk/identity/azure-identity/azure/identity/_authn_client.py b/sdk/identity/azure-identity/azure/identity/_authn_client.py deleted file mode 100644 index 23901b0cae3d..000000000000 --- a/sdk/identity/azure-identity/azure/identity/_authn_client.py +++ /dev/null @@ -1,237 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -import abc -import calendar -import time - -from msal import TokenCache - -from azure.core.configuration import Configuration -from azure.core.credentials import AccessToken -from azure.core.exceptions import ClientAuthenticationError -from azure.core.pipeline import Pipeline -from azure.core.pipeline.policies import ( - ContentDecodePolicy, - DistributedTracingPolicy, - HttpLoggingPolicy, - NetworkTraceLoggingPolicy, - ProxyPolicy, - RetryPolicy, - UserAgentPolicy, -) -from azure.core.pipeline.transport import RequestsTransport, HttpRequest -from ._constants import DEVELOPER_SIGN_ON_CLIENT_ID, DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY -from ._internal import get_default_authority, normalize_authority -from ._internal.user_agent import USER_AGENT - -try: - ABC = abc.ABC -except AttributeError: # Python 2.7, abc exists, but not ABC - ABC = abc.ABCMeta("ABC", (object,), {"__slots__": ()}) # type: ignore - -try: - from typing import TYPE_CHECKING -except ImportError: - TYPE_CHECKING = False - -if TYPE_CHECKING: - # pylint:disable=unused-import,ungrouped-imports - from time import struct_time - from typing import Any, Dict, Iterable, List, Mapping, Optional, Union - from azure.core.pipeline import PipelineResponse - from azure.core.pipeline.transport import HttpTransport - from azure.core.pipeline.policies import HTTPPolicy, SansIOHTTPPolicy - - PolicyListType = List[Union[HTTPPolicy, SansIOHTTPPolicy]] - - -class AuthnClientBase(ABC): - """Sans I/O authentication client methods""" - - def __init__(self, endpoint=None, authority=None, tenant=None, **kwargs): # pylint:disable=unused-argument - # type: (Optional[str], Optional[str], Optional[str], **Any) -> None - super(AuthnClientBase, self).__init__() - if authority and endpoint: - raise ValueError( - "'authority' and 'endpoint' are mutually exclusive. 'authority' should be the authority of an AAD" - + " endpoint, whereas 'endpoint' should be the endpoint's full URL." - ) - - if endpoint: - self._auth_url = endpoint - else: - if not tenant: - raise ValueError("'tenant' is required") - authority = normalize_authority(authority) if authority else get_default_authority() - self._auth_url = "/".join((authority, tenant.strip("/"), "oauth2/v2.0/token")) - self._cache = kwargs.get("cache") or TokenCache() # type: TokenCache - self._token_refresh_retry_delay = DEFAULT_TOKEN_REFRESH_RETRY_DELAY - self._token_refresh_offset = DEFAULT_REFRESH_OFFSET - self._last_refresh_time = 0 - - @property - def auth_url(self): - return self._auth_url - - def get_cached_token(self, scopes): - # type: (Iterable[str]) -> Optional[AccessToken] - tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=list(scopes)) - for token in tokens: - expires_on = int(token["expires_on"]) - if expires_on > int(time.time()): - return AccessToken(token["secret"], expires_on) - return None - - def get_refresh_token_grant_request(self, refresh_token, scopes): - data = { - "grant_type": "refresh_token", - "refresh_token": refresh_token["secret"], - "scope": " ".join(scopes), - "client_id": DEVELOPER_SIGN_ON_CLIENT_ID, - } - return self._prepare_request(form_data=data) - - @abc.abstractmethod - def request_token(self, scopes, method, headers, form_data, params, **kwargs): - pass - - def _deserialize_and_cache_token(self, response, scopes, request_time): - # type: (PipelineResponse, Iterable[str], int) -> AccessToken - """Deserialize and cache an access token from an AAD response""" - - # ContentDecodePolicy sets this, and should have raised if it couldn't deserialize the response - payload = response.context[ContentDecodePolicy.CONTEXT_NAME] - - if not payload or "access_token" not in payload or not ("expires_in" in payload or "expires_on" in payload): - if payload and "access_token" in payload: - payload["access_token"] = "****" - raise ClientAuthenticationError( - message="Unexpected response '{}'".format(payload), response=response.http_response - ) - - token = payload["access_token"] - - # AccessToken wants expires_on as an int - expires_on = payload.get("expires_on") or int(payload["expires_in"]) + request_time # type: Union[str, int] - try: - expires_on = int(expires_on) - except ValueError: - # probably an App Service MSI 2017-09-01 response, convert it to epoch seconds - try: - t = self._parse_app_service_expires_on(expires_on) # type: ignore - expires_on = calendar.timegm(t) - except ValueError: - # have a token but don't know when it expires -> treat it as single-use - expires_on = request_time - - # now we have an int expires_on, ensure the cache entry gets it - payload["expires_on"] = expires_on - - self._cache.add({"response": payload, "scope": scopes}) - - return AccessToken(token, expires_on) - - @staticmethod - def _parse_app_service_expires_on(expires_on): - # type: (str) -> struct_time - """Parse an App Service MSI version 2017-09-01 expires_on value to struct_time. - - This version of the API returns expires_on as a UTC datetime string rather than epoch seconds. The string's - format depends on the OS. Responses on Windows include AM/PM, for example "1/16/2020 5:24:12 AM +00:00". - Responses on Linux do not, for example "06/20/2019 02:57:58 +00:00". - - :raises ValueError: ``expires_on`` didn't match an expected format - """ - - # parse the string minus the timezone offset - if expires_on.endswith(" +00:00"): - date_string = expires_on[: -len(" +00:00")] - for format_string in ("%m/%d/%Y %H:%M:%S", "%m/%d/%Y %I:%M:%S %p"): # (Linux, Windows) - try: - return time.strptime(date_string, format_string) - except ValueError: - pass - - raise ValueError("'{}' doesn't match the expected format".format(expires_on)) - - def _prepare_request( - self, - method="POST", # type: str - headers=None, # type: Optional[Mapping[str, str]] - form_data=None, # type: Optional[Mapping[str, str]] - params=None, # type: Optional[Dict[str, str]] - ): - # type: (...) -> HttpRequest - request = HttpRequest(method, self._auth_url, headers=headers) - if form_data: - request.headers["Content-Type"] = "application/x-www-form-urlencoded" - request.set_formdata_body(form_data) - if params: - request.format_parameters(params) - return request - - -class AuthnClient(AuthnClientBase): - """Synchronous authentication client. - - :param str auth_url: - :param config: Optional configuration for the HTTP pipeline. - :type config: :class:`azure.core.configuration` - :param policies: Optional policies for the HTTP pipeline. - :type policies: - :param transport: Optional HTTP transport. - :type transport: - """ - - # pylint:disable=missing-client-constructor-parameter-credential - def __init__( - self, - config=None, # type: Optional[Configuration] - policies=None, # type: Optional[PolicyListType] - transport=None, # type: Optional[HttpTransport] - **kwargs # type: Any - ): - # type: (...) -> None - config = config or self._create_config(**kwargs) - policies = policies or [ - ContentDecodePolicy(), - config.user_agent_policy, - config.proxy_policy, - config.retry_policy, - config.logging_policy, - DistributedTracingPolicy(**kwargs), - HttpLoggingPolicy(**kwargs), - ] - if not transport: - transport = RequestsTransport(**kwargs) - self._pipeline = Pipeline(transport=transport, policies=policies) # type: Pipeline - super(AuthnClient, self).__init__(**kwargs) - - def request_token( - self, - scopes, # type: Iterable[str] - method="POST", # type: str - headers=None, # type: Optional[Mapping[str, str]] - form_data=None, # type: Optional[Mapping[str, str]] - params=None, # type: Optional[Dict[str, str]] - **kwargs # type: Any - ): - # type: (...) -> AccessToken - request = self._prepare_request(method, headers=headers, form_data=form_data, params=params) - request_time = int(time.time()) - self._last_refresh_time = request_time # no matter succeed or not, update the last refresh time - response = self._pipeline.run(request, stream=False, **kwargs) - token = self._deserialize_and_cache_token(response=response, scopes=scopes, request_time=request_time) - return token - - @staticmethod - def _create_config(**kwargs): - # type: (Mapping[str, Any]) -> Configuration - config = Configuration(**kwargs) - config.logging_policy = NetworkTraceLoggingPolicy(**kwargs) - config.retry_policy = RetryPolicy(**kwargs) - config.proxy_policy = ProxyPolicy(**kwargs) - config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs) - return config diff --git a/sdk/identity/azure-identity/azure/identity/_constants.py b/sdk/identity/azure-identity/azure/identity/_constants.py index 4f729e678523..cb13afefacbe 100644 --- a/sdk/identity/azure-identity/azure/identity/_constants.py +++ b/sdk/identity/azure-identity/azure/identity/_constants.py @@ -43,10 +43,3 @@ class EnvironmentVariables: MSI_SECRET = "MSI_SECRET" AZURE_AUTHORITY_HOST = "AZURE_AUTHORITY_HOST" - - -class Endpoints: - # https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http - IMDS = "http://169.254.169.254/metadata/identity/oauth2/token" - - AAD_OAUTH2_V2_FORMAT = "https://" + KnownAuthorities.AZURE_PUBLIC_CLOUD + "/{}/oauth2/v2.0/token" diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/imds.py b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py new file mode 100644 index 000000000000..add9a1dd3c09 --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/_credentials/imds.py @@ -0,0 +1,94 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +import logging +from typing import TYPE_CHECKING + +import six + +from azure.core.exceptions import ClientAuthenticationError, HttpResponseError +from azure.core.pipeline.transport import HttpRequest + +from .. import CredentialUnavailableError +from .._internal.get_token_mixin import GetTokenMixin +from .._internal.managed_identity_client import ManagedIdentityClient + +if TYPE_CHECKING: + # pylint:disable=ungrouped-imports + from typing import Any, Optional + from azure.core.credentials import AccessToken + +_LOGGER = logging.getLogger(__name__) + +IMDS_URL = "http://169.254.169.254/metadata/identity/oauth2/token" + +PIPELINE_SETTINGS = { + "connection_timeout": 2, + "retry_backoff_factor": 4, + "retry_backoff_max": 60, + "retry_on_status_codes": [404, 429] + list(range(500, 600)), + "retry_status": 5, + "retry_total": 5, +} + + +def get_request(scope, identity_config): + request = HttpRequest("GET", IMDS_URL) + request.format_parameters(dict({"api-version": "2018-02-01", "resource": scope}, **identity_config)) + return request + + +class ImdsCredential(GetTokenMixin): + def __init__(self, **kwargs): + # type: (**Any) -> None + super(ImdsCredential, self).__init__() + + self._client = ManagedIdentityClient( + get_request, _identity_config=kwargs.pop("identity_config", None) or {}, **dict(PIPELINE_SETTINGS, **kwargs) + ) + self._endpoint_available = None # type: Optional[bool] + self._user_assigned_identity = "client_id" in kwargs or "identity_config" in kwargs + + def _acquire_token_silently(self, *scopes): + # type: (*str) -> Optional[AccessToken] + return self._client.get_cached_token(*scopes) + + def _request_token(self, *scopes, **kwargs): # pylint:disable=unused-argument + # type: (*str, **Any) -> AccessToken + if self._endpoint_available is None: + # Lacking another way to determine whether the IMDS endpoint is listening, + # we send a request it would immediately reject (because it lacks the Metadata header), + # setting a short timeout. + try: + self._client.request_token(*scopes, connection_timeout=0.3, retry_total=0) + self._endpoint_available = True + except HttpResponseError: + # received a response, choked on it + self._endpoint_available = True + except Exception: # pylint:disable=broad-except + # if anything else was raised, assume the endpoint is unavailable + self._endpoint_available = False + _LOGGER.info("No response from the IMDS endpoint.") + + if not self._endpoint_available: + message = "ManagedIdentityCredential authentication unavailable, no managed identity endpoint found." + raise CredentialUnavailableError(message=message) + + try: + token = self._client.request_token(*scopes, headers={"Metadata": "true"}) + except HttpResponseError as ex: + # 400 in response to a token request indicates managed identity is disabled, + # or the identity with the specified client_id is not available + if ex.status_code == 400: + self._endpoint_available = False + message = "ManagedIdentityCredential authentication unavailable. " + if self._user_assigned_identity: + message += "The requested identity has not been assigned to this resource." + else: + message += "No identity has been assigned to this resource." + six.raise_from(CredentialUnavailableError(message=message), ex) + + # any other error is unexpected + six.raise_from(ClientAuthenticationError(message=ex.message, response=ex.response), None) + return token diff --git a/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py b/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py index d2001a1002eb..ff4f6ae177db 100644 --- a/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py +++ b/sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py @@ -5,26 +5,9 @@ import logging import os -import six -from azure.core.configuration import Configuration -from azure.core.credentials import AccessToken -from azure.core.exceptions import ClientAuthenticationError, HttpResponseError -from azure.core.pipeline.policies import ( - ContentDecodePolicy, - DistributedTracingPolicy, - HeadersPolicy, - HttpLoggingPolicy, - NetworkTraceLoggingPolicy, - RetryPolicy, - UserAgentPolicy, -) - from .. import CredentialUnavailableError -from .._authn_client import AuthnClient -from .._constants import Endpoints, EnvironmentVariables +from .._constants import EnvironmentVariables from .._internal.decorators import log_get_token -from .._internal.get_token_mixin import GetTokenMixin -from .._internal.user_agent import USER_AGENT try: from typing import TYPE_CHECKING @@ -33,8 +16,8 @@ if TYPE_CHECKING: # pylint:disable=unused-import - from typing import Any, Optional, Type - from azure.core.credentials import TokenCredential + from typing import Any, Optional + from azure.core.credentials import AccessToken, TokenCredential _LOGGER = logging.getLogger(__name__) @@ -67,9 +50,8 @@ def __init__(self, **kwargs): self._credential = CloudShellCredential(**kwargs) elif os.environ.get(EnvironmentVariables.IDENTITY_ENDPOINT): - if ( - os.environ.get(EnvironmentVariables.IDENTITY_HEADER) - and os.environ.get(EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT) + if os.environ.get(EnvironmentVariables.IDENTITY_HEADER) and os.environ.get( + EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT ): _LOGGER.info("%s will use Service Fabric managed identity", self.__class__.__name__) from .service_fabric import ServiceFabricCredential @@ -81,6 +63,8 @@ def __init__(self, **kwargs): self._credential = AzureArcCredential(**kwargs) else: + from .imds import ImdsCredential + _LOGGER.info("%s will use IMDS", self.__class__.__name__) self._credential = ImdsCredential(**kwargs) @@ -99,122 +83,3 @@ def get_token(self, *scopes, **kwargs): if not self._credential: raise CredentialUnavailableError(message="No managed identity endpoint found.") return self._credential.get_token(*scopes, **kwargs) - - -class _ManagedIdentityBase(object): - def __init__(self, endpoint, client_cls, config=None, client_id=None, **kwargs): - # type: (str, Type, Optional[Configuration], Optional[str], **Any) -> None - self._identity_config = kwargs.pop("identity_config", None) or {} - if client_id: - if os.environ.get(EnvironmentVariables.MSI_ENDPOINT) and os.environ.get(EnvironmentVariables.MSI_SECRET): - # App Service: version 2017-09-1 accepts client ID as parameter "clientid" - if "clientid" not in self._identity_config: - self._identity_config["clientid"] = client_id - elif "client_id" not in self._identity_config: - self._identity_config["client_id"] = client_id - - config = config or self._create_config(**kwargs) - policies = [ - ContentDecodePolicy(), - config.headers_policy, - config.user_agent_policy, - config.retry_policy, - config.logging_policy, - DistributedTracingPolicy(**kwargs), - HttpLoggingPolicy(**kwargs), - ] - self._client = client_cls(endpoint=endpoint, config=config, policies=policies, **kwargs) - - @staticmethod - def _create_config(**kwargs): - # type: (**Any) -> Configuration - """Build a default configuration for the credential's HTTP pipeline.""" - - timeout = kwargs.pop("connection_timeout", 2) - config = Configuration(connection_timeout=timeout, **kwargs) - - # retry is the only IO policy, so its class is a kwarg to increase async code sharing - retry_policy = kwargs.pop("retry_policy", RetryPolicy) # type: ignore - args = kwargs.copy() # combine kwargs and default retry settings in a Python 2-compatible way - args.update(_ManagedIdentityBase._retry_settings) # type: ignore - config.retry_policy = retry_policy(**args) # type: ignore - - # Metadata header is required by IMDS and in Cloud Shell; App Service ignores it - config.headers_policy = HeadersPolicy(base_headers={"Metadata": "true"}, **kwargs) - config.logging_policy = NetworkTraceLoggingPolicy(**kwargs) - config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs) - - return config - - # given RetryPolicy's implementation, these settings most closely match the documented guidance for IMDS - # https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#retry-guidance - _retry_settings = { - "retry_total": 5, - "retry_status": 5, - "retry_backoff_factor": 4, - "retry_backoff_max": 60, - "retry_on_status_codes": [404, 429] + list(range(500, 600)), - } - - -class ImdsCredential(_ManagedIdentityBase, GetTokenMixin): - """Authenticates with a managed identity via the IMDS endpoint. - - :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity. - """ - - def __init__(self, **kwargs): - # type: (**Any) -> None - super(ImdsCredential, self).__init__(endpoint=Endpoints.IMDS, client_cls=AuthnClient, **kwargs) - self._endpoint_available = None # type: Optional[bool] - - def _acquire_token_silently(self, *scopes): - # type: (*str) -> Optional[AccessToken] - return self._client.get_cached_token(scopes) - - def _request_token(self, *scopes, **kwargs): # pylint:disable=unused-argument - # type: (*str, **Any) -> AccessToken - if self._endpoint_available is None: - # Lacking another way to determine whether the IMDS endpoint is listening, - # we send a request it would immediately reject (missing a required header), - # setting a short timeout. - try: - self._client.request_token(scopes, method="GET", connection_timeout=0.3, retry_total=0) - self._endpoint_available = True - except HttpResponseError: - # received a response, choked on it - self._endpoint_available = True - except Exception: # pylint:disable=broad-except - # if anything else was raised, assume the endpoint is unavailable - self._endpoint_available = False - _LOGGER.info("No response from the IMDS endpoint.") - - if not self._endpoint_available: - message = "ManagedIdentityCredential authentication unavailable, no managed identity endpoint found." - raise CredentialUnavailableError(message=message) - - if len(scopes) != 1: - raise ValueError("This credential requires exactly one scope per token request.") - - resource = scopes[0] - if resource.endswith("/.default"): - resource = resource[: -len("/.default")] - params = dict({"api-version": "2018-02-01", "resource": resource}, **self._identity_config) - - try: - token = self._client.request_token(scopes, method="GET", params=params) - except HttpResponseError as ex: - # 400 in response to a token request indicates managed identity is disabled, - # or the identity with the specified client_id is not available - if ex.status_code == 400: - self._endpoint_available = False - message = "ManagedIdentityCredential authentication unavailable. " - if self._identity_config: - message += "The requested identity has not been assigned to this resource." - else: - message += "No identity has been assigned to this resource." - six.raise_from(CredentialUnavailableError(message=message), ex) - - # any other error is unexpected - six.raise_from(ClientAuthenticationError(message=ex.message, response=ex.response), None) - return token diff --git a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py index e2da24a1de4a..700a3540c299 100644 --- a/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py +++ b/sdk/identity/azure-identity/azure/identity/_internal/managed_identity_client.py @@ -104,12 +104,12 @@ def _build_pipeline(self, config, policies=None, transport=None, **kwargs): class ManagedIdentityClient(ManagedIdentityClientBase): - def request_token(self, *scopes, **kwargs): # pylint:disable=unused-argument + def request_token(self, *scopes, **kwargs): # type: (*str, **Any) -> AccessToken resource = _scopes_to_resource(*scopes) request = self._request_factory(resource, self._identity_config) request_time = int(time.time()) - response = self._pipeline.run(request, retry_on_methods=[request.method]) + response = self._pipeline.run(request, retry_on_methods=[request.method], **kwargs) token = self._process_response(response, request_time) return token diff --git a/sdk/identity/azure-identity/azure/identity/aio/_authn_client.py b/sdk/identity/azure-identity/azure/identity/aio/_authn_client.py deleted file mode 100644 index 24e52e250d0a..000000000000 --- a/sdk/identity/azure-identity/azure/identity/aio/_authn_client.py +++ /dev/null @@ -1,92 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -import time -from typing import TYPE_CHECKING - -from azure.core.configuration import Configuration -from azure.core.credentials import AccessToken -from azure.core.pipeline import AsyncPipeline -from azure.core.pipeline.policies import ( - AsyncRetryPolicy, - ContentDecodePolicy, - DistributedTracingPolicy, - HttpLoggingPolicy, - NetworkTraceLoggingPolicy, - ProxyPolicy, - UserAgentPolicy, -) -from azure.core.pipeline.transport import AioHttpTransport - -from .._authn_client import AuthnClientBase -from .._internal.user_agent import USER_AGENT - -if TYPE_CHECKING: - from typing import Any, Dict, Iterable, List, Mapping, Optional, Union - from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy - from azure.core.pipeline.transport import AsyncHttpTransport - - PolicyListType = List[Union[AsyncHTTPPolicy, SansIOHTTPPolicy]] - - -class AsyncAuthnClient(AuthnClientBase): # pylint:disable=async-client-bad-name - """Async authentication client""" - - # pylint:disable=missing-client-constructor-parameter-credential - def __init__( - self, - config: "Optional[Configuration]" = None, - policies: "Optional[PolicyListType]" = None, - transport: "Optional[AsyncHttpTransport]" = None, - **kwargs: "Any" - ) -> None: - config = config or self._create_config(**kwargs) - policies = policies or [ - ContentDecodePolicy(), - config.user_agent_policy, - config.proxy_policy, - config.retry_policy, - config.logging_policy, - DistributedTracingPolicy(**kwargs), - HttpLoggingPolicy(**kwargs), - ] - if not transport: - transport = AioHttpTransport(**kwargs) - self._pipeline = AsyncPipeline(transport=transport, policies=policies) # type: AsyncPipeline - super().__init__(**kwargs) - - async def __aenter__(self): - await self._pipeline.__aenter__() - return self - - async def __aexit__(self, *args): - await self.close() - - async def close(self) -> None: - await self._pipeline.__aexit__() - - async def request_token( # pylint:disable=invalid-overridden-method - self, - scopes: "Iterable[str]", - method: str = "POST", - headers: "Optional[Mapping[str, str]]" = None, - form_data: "Optional[Mapping[str, str]]" = None, - params: "Optional[Dict[str, str]]" = None, - **kwargs: "Any" - ) -> AccessToken: - request = self._prepare_request(method, headers=headers, form_data=form_data, params=params) - request_time = int(time.time()) - self._last_refresh_time = request_time # no matter succeed or not, update the last refresh time - response = await self._pipeline.run(request, stream=False, **kwargs) - token = self._deserialize_and_cache_token(response=response, scopes=scopes, request_time=request_time) - return token - - @staticmethod - def _create_config(**kwargs: "Any") -> Configuration: - config = Configuration(**kwargs) - config.logging_policy = NetworkTraceLoggingPolicy(**kwargs) - config.retry_policy = AsyncRetryPolicy(**kwargs) - config.proxy_policy = ProxyPolicy(**kwargs) - config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs) - return config diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py new file mode 100644 index 000000000000..9f39a347587e --- /dev/null +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/imds.py @@ -0,0 +1,75 @@ +# ------------------------------------ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +# ------------------------------------ +import logging +from typing import TYPE_CHECKING + +from azure.core.exceptions import ClientAuthenticationError, HttpResponseError + +from ... import CredentialUnavailableError +from .._internal import AsyncContextManager +from .._internal.get_token_mixin import GetTokenMixin +from .._internal.managed_identity_client import AsyncManagedIdentityClient +from ..._credentials.imds import get_request, PIPELINE_SETTINGS + +if TYPE_CHECKING: + from typing import Any, Optional + from azure.core.credentials import AccessToken + +_LOGGER = logging.getLogger(__name__) + + +class ImdsCredential(AsyncContextManager, GetTokenMixin): + def __init__(self, **kwargs: "Any") -> None: + super().__init__() + + self._client = AsyncManagedIdentityClient( + get_request, _identity_config=kwargs.pop("identity_config", None) or {}, **PIPELINE_SETTINGS, **kwargs + ) + self._endpoint_available = None # type: Optional[bool] + self._user_assigned_identity = "client_id" in kwargs or "identity_config" in kwargs + + async def close(self) -> None: + await self._client.close() + + async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]": + return self._client.get_cached_token(*scopes) + + async def _request_token(self, *scopes, **kwargs: "Any") -> "AccessToken": # pylint:disable=unused-argument + if self._endpoint_available is None: + # Lacking another way to determine whether the IMDS endpoint is listening, + # we send a request it would immediately reject (because it lacks the Metadata header), + # setting a short timeout. + try: + await self._client.request_token(*scopes, connection_timeout=0.3, retry_total=0) + self._endpoint_available = True + except HttpResponseError: + # received a response, choked on it + self._endpoint_available = True + except Exception: # pylint:disable=broad-except + # if anything else was raised, assume the endpoint is unavailable + self._endpoint_available = False + _LOGGER.info("No response from the IMDS endpoint.") + + if not self._endpoint_available: + message = "ManagedIdentityCredential authentication unavailable, no managed identity endpoint found." + raise CredentialUnavailableError(message=message) + + try: + token = await self._client.request_token(*scopes, headers={"Metadata": "true"}) + except HttpResponseError as ex: + # 400 in response to a token request indicates managed identity is disabled, + # or the identity with the specified client_id is not available + if ex.status_code == 400: + self._endpoint_available = False + message = "ManagedIdentityCredential authentication unavailable. " + if self._user_assigned_identity: + message += "The requested identity has not been assigned to this resource." + else: + message += "No identity has been assigned to this resource." + raise CredentialUnavailableError(message=message) from ex + + # any other error is unexpected + raise ClientAuthenticationError(message=ex.message, response=ex.response) from None + return token diff --git a/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py b/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py index d2d64f80ee6e..cf11359489cb 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py @@ -7,20 +7,14 @@ from typing import TYPE_CHECKING from azure.core.credentials import AccessToken -from azure.core.exceptions import ClientAuthenticationError, HttpResponseError -from azure.core.pipeline.policies import AsyncRetryPolicy -from .._authn_client import AsyncAuthnClient from .._internal import AsyncContextManager from .._internal.decorators import log_get_token_async -from .._internal.get_token_mixin import GetTokenMixin from ... import CredentialUnavailableError -from ..._constants import Endpoints, EnvironmentVariables -from ..._credentials.managed_identity import _ManagedIdentityBase +from ..._constants import EnvironmentVariables if TYPE_CHECKING: from typing import Any, Optional - from azure.core.configuration import Configuration from azure.core.credentials_async import AsyncTokenCredential _LOGGER = logging.getLogger(__name__) @@ -54,9 +48,8 @@ def __init__(self, **kwargs: "Any") -> None: self._credential = CloudShellCredential(**kwargs) elif os.environ.get(EnvironmentVariables.IDENTITY_ENDPOINT): - if ( - os.environ.get(EnvironmentVariables.IDENTITY_HEADER) - and os.environ.get(EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT) + if os.environ.get(EnvironmentVariables.IDENTITY_HEADER) and os.environ.get( + EnvironmentVariables.IDENTITY_SERVER_THUMBPRINT ): _LOGGER.info("%s will use Service Fabric managed identity", self.__class__.__name__) from .service_fabric import ServiceFabricCredential @@ -68,6 +61,8 @@ def __init__(self, **kwargs: "Any") -> None: self._credential = AzureArcCredential(**kwargs) else: + from .imds import ImdsCredential + _LOGGER.info("%s will use IMDS", self.__class__.__name__) self._credential = ImdsCredential(**kwargs) @@ -94,82 +89,3 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken": if not self._credential: raise CredentialUnavailableError(message="No managed identity endpoint found.") return await self._credential.get_token(*scopes, **kwargs) - - -class _AsyncManagedIdentityBase(_ManagedIdentityBase, AsyncContextManager): - def __init__(self, endpoint: str, **kwargs: "Any") -> None: - super().__init__(endpoint=endpoint, client_cls=AsyncAuthnClient, **kwargs) - - async def __aenter__(self): - await self._client.__aenter__() - return self - - async def close(self): - """Close the credential's transport session.""" - - await self._client.__aexit__() - - @staticmethod - def _create_config(**kwargs: "Any") -> "Configuration": - """Build a default configuration for the credential's HTTP pipeline.""" - return _ManagedIdentityBase._create_config(retry_policy=AsyncRetryPolicy, **kwargs) - - -class ImdsCredential(_AsyncManagedIdentityBase, GetTokenMixin): - """Asynchronously authenticates with a managed identity via the IMDS endpoint. - - :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity. - """ - - def __init__(self, **kwargs: "Any") -> None: - super().__init__(endpoint=Endpoints.IMDS, **kwargs) - self._endpoint_available = None # type: Optional[bool] - - async def _acquire_token_silently(self, *scopes: str) -> "Optional[AccessToken]": - return self._client.get_cached_token(scopes) - - async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken": # pylint:disable=unused-argument - if self._endpoint_available is None: - # Lacking another way to determine whether the IMDS endpoint is listening, - # we send a request it would immediately reject (missing a required header), - # setting a short timeout. - try: - await self._client.request_token(scopes, method="GET", connection_timeout=0.3, retry_total=0) - self._endpoint_available = True - except HttpResponseError: - # received a response, choked on it - self._endpoint_available = True - except Exception: # pylint:disable=broad-except - # if anything else was raised, assume the endpoint is unavailable - self._endpoint_available = False - _LOGGER.info("No response from the IMDS endpoint.") - - if not self._endpoint_available: - message = "ManagedIdentityCredential authentication unavailable, no managed identity endpoint found." - raise CredentialUnavailableError(message=message) - - if len(scopes) != 1: - raise ValueError("This credential requires exactly one scope per token request.") - - resource = scopes[0] - if resource.endswith("/.default"): - resource = resource[: -len("/.default")] - params = {"api-version": "2018-02-01", "resource": resource, **self._identity_config} - - try: - token = await self._client.request_token(scopes, method="GET", params=params) - except HttpResponseError as ex: - # 400 in response to a token request indicates managed identity is disabled, - # or the identity with the specified client_id is not available - if ex.status_code == 400: - self._endpoint_available = False - message = "ManagedIdentityCredential authentication unavailable. " - if self._identity_config: - message += "The requested identity has not been assigned to this resource." - else: - message += "No identity has been assigned to this resource." - raise CredentialUnavailableError(message=message) from ex - - # any other error is unexpected - raise ClientAuthenticationError(message=ex.message, response=ex.response) from None - return token diff --git a/sdk/identity/azure-identity/azure/identity/aio/_internal/managed_identity_client.py b/sdk/identity/azure-identity/azure/identity/aio/_internal/managed_identity_client.py index 4dfb0cc1e7db..27b3426e8be5 100644 --- a/sdk/identity/azure-identity/azure/identity/aio/_internal/managed_identity_client.py +++ b/sdk/identity/azure-identity/azure/identity/aio/_internal/managed_identity_client.py @@ -32,11 +32,11 @@ async def close(self) -> None: await self._pipeline.__aexit__() async def request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken": - # pylint:disable=invalid-overridden-method,unused-argument + # pylint:disable=invalid-overridden-method resource = _scopes_to_resource(*scopes) request = self._request_factory(resource, self._identity_config) request_time = int(time.time()) - response = await self._pipeline.run(request, retry_on_methods=[request.method]) + response = await self._pipeline.run(request, retry_on_methods=[request.method], **kwargs) token = self._process_response(response, request_time) return token diff --git a/sdk/identity/azure-identity/tests/test_authn_client.py b/sdk/identity/azure-identity/tests/test_authn_client.py deleted file mode 100644 index ecab2c8a462d..000000000000 --- a/sdk/identity/azure-identity/tests/test_authn_client.py +++ /dev/null @@ -1,235 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -"""These tests use the synchronous AuthnClient as a driver to test functionality -of the sans I/O AuthnClientBase shared with AsyncAuthnClient.""" -import json -import time - -try: - from unittest.mock import Mock, patch -except ImportError: # python < 3.3 - from mock import Mock, patch # type: ignore - -from azure.core.credentials import AccessToken -from azure.identity._authn_client import AuthnClient -from azure.identity._constants import EnvironmentVariables, DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY -import pytest -from six.moves.urllib_parse import urlparse -from helpers import mock_response - - -def test_deserialization_expires_integers(): - now = 6 - expires_in = 59 - now - expires_on = now + expires_in - access_token = "***" - expected_access_token = AccessToken(access_token, expires_on) - scope = "scope" - - # response with expires_on only - token_payload = {"access_token": access_token, "expires_on": expires_on, "token_type": "Bearer", "resource": scope} - mock_send = Mock(return_value=mock_response(json_payload=token_payload)) - token = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)).request_token(scope) - assert token == expected_access_token - - # response with expires_in as well - token_payload = { - "access_token": access_token, - "expires_in": expires_in, - "token_type": "Bearer", - "ext_expires_in": expires_in, - } - with patch(AuthnClient.__module__ + ".time.time") as mock_time: - mock_time.return_value = now - mock_send = Mock(return_value=mock_response(json_payload=token_payload)) - token = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)).request_token(scope) - assert token == expected_access_token - - -def test_deserialization_app_service_msi(): - now = 6 - expires_in = 59 - now - expires_on = now + expires_in - access_token = "***" - expected_access_token = AccessToken(access_token, expires_on) - scope = "scope" - - # response with expires_on only and it's a datetime string (App Service MSI, Linux) - token_payload = { - "access_token": access_token, - "expires_on": "01/01/1970 00:00:{} +00:00".format(now + expires_in), - "token_type": "Bearer", - "resource": scope, - } - mock_send = Mock(return_value=mock_response(json_payload=token_payload)) - token = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)).request_token(scope) - assert token == expected_access_token - - -def test_deserialization_expires_strings(): - now = 6 - expires_in = 59 - now - expires_on = now + expires_in - access_token = "***" - expected_access_token = AccessToken(access_token, expires_on) - scope = "scope" - - # response with string expires_in and expires_on (IMDS, Cloud Shell) - token_payload = { - "access_token": access_token, - "expires_in": str(expires_in), - "ext_expires_in": str(expires_in), - "expires_on": str(expires_on), - "token_type": "Bearer", - "resource": scope, - } - mock_send = Mock(return_value=mock_response(json_payload=token_payload)) - token = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)).request_token(scope) - assert token == expected_access_token - - -def test_caching_when_only_expires_in_set(): - """the cache should function when auth responses don't include an explicit expires_on""" - - access_token = "token" - now = 42 - expires_in = 1800 - expires_on = now + expires_in - expected_token = AccessToken(access_token, expires_on) - - mock_send = Mock( - return_value=mock_response( - json_payload={"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer"} - ) - ) - - client = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)) - with patch("azure.identity._authn_client.time.time") as mock_time: - mock_time.return_value = 42 - token = client.request_token(["scope"]) - assert token.token == expected_token.token - assert token.expires_on == expected_token.expires_on - - cached_token = client.get_cached_token(["scope"]) - assert cached_token == expected_token - - -def test_expires_in_strings(): - expected_token = "token" - - mock_send = Mock( - return_value=mock_response( - json_payload={ - "access_token": expected_token, - "expires_in": "42", - "ext_expires_in": "42", - "token_type": "Bearer", - } - ) - ) - - now = int(time.time()) - with patch("azure.identity._authn_client.time.time") as mock_time: - mock_time.return_value = now - token = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)).request_token("scope") - assert token.token == expected_token - assert token.expires_on == now + 42 - - -def test_cache_expiry(): - access_token = "token" - now = 42 - expires_in = 1800 - expires_on = now + expires_in - expected_token = AccessToken(access_token, expires_on) - token_payload = {"access_token": access_token, "expires_in": expires_in, "token_type": "Bearer"} - mock_send = Mock(return_value=mock_response(json_payload=token_payload)) - - client = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)) - with patch("azure.identity._authn_client.time.time") as mock_time: - # populate the cache with a valid token - mock_time.return_value = now - token = client.request_token("scope") - assert token.token == expected_token.token - assert token.expires_on == expected_token.expires_on - - cached_token = client.get_cached_token("scope") - assert cached_token == expected_token - - # advance time past the cached token's expires_on - mock_time.return_value = expires_on + 3600 - cached_token = client.get_cached_token("scope") - assert not cached_token - - # request a new token - new_token = "new token" - token_payload["access_token"] = new_token - token = client.request_token("scope") - assert token.token == new_token - - # it should be cached - cached_token = client.get_cached_token("scope") - assert cached_token.token == new_token - - -def test_cache_scopes(): - scope_a = "scope_a" - scope_b = "scope_b" - scope_ab = scope_a + " " + scope_b - expected_tokens = { - scope_a: {"access_token": scope_a, "expires_in": 1 << 31, "ext_expires_in": 1 << 31, "token_type": "Bearer"}, - scope_b: {"access_token": scope_b, "expires_in": 1 << 31, "ext_expires_in": 1 << 31, "token_type": "Bearer"}, - scope_ab: {"access_token": scope_ab, "expires_in": 1 << 31, "ext_expires_in": 1 << 31, "token_type": "Bearer"}, - } - - def mock_send(request, **kwargs): - token = expected_tokens[request.data["resource"]] - return mock_response(json_payload=token) - - client = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)) - - # if the cache has a token for a & b, it should hit for a, b, a & b - token = client.request_token([scope_a, scope_b], form_data={"resource": scope_ab}) - assert token.token == scope_ab - for scope in (scope_a, scope_b): - assert client.get_cached_token([scope]).token == scope_ab - assert client.get_cached_token([scope_a, scope_b]).token == scope_ab - - # if the cache has only tokens for a and b alone, a & b should miss - client = AuthnClient(endpoint="http://foo", transport=Mock(send=mock_send)) - for scope in (scope_a, scope_b): - token = client.request_token([scope], form_data={"resource": scope}) - assert token.token == scope - assert client.get_cached_token([scope]).token == scope - assert not client.get_cached_token([scope_a, scope_b]) - - -@pytest.mark.parametrize("authority", ("localhost", "https://localhost")) -def test_request_url(authority): - tenant_id = "expected-tenant" - parsed_authority = urlparse(authority) - expected_netloc = parsed_authority.netloc or authority # "localhost" parses to netloc "", path "localhost" - - def validate_url(url): - actual = urlparse(url) - assert actual.scheme == "https" - assert actual.netloc == expected_netloc - assert actual.path.startswith("/" + tenant_id) - - def mock_send(request, **kwargs): - validate_url(request.url) - return mock_response(json_payload={"token_type": "Bearer", "expires_in": 42, "access_token": "***"}) - - client = AuthnClient(tenant=tenant_id, transport=Mock(send=mock_send), authority=authority) - client.request_token(("scope",)) - request = client.get_refresh_token_grant_request({"secret": "***"}, "scope") - validate_url(request.url) - - # authority can be configured via environment variable - with patch.dict("os.environ", {EnvironmentVariables.AZURE_AUTHORITY_HOST: authority}, clear=True): - client = AuthnClient(tenant=tenant_id, transport=Mock(send=mock_send)) - client.request_token(("scope",)) - request = client.get_refresh_token_grant_request({"secret": "***"}, "scope") - validate_url(request.url) diff --git a/sdk/identity/azure-identity/tests/test_authn_client_async.py b/sdk/identity/azure-identity/tests/test_authn_client_async.py deleted file mode 100644 index 6b6178163580..000000000000 --- a/sdk/identity/azure-identity/tests/test_authn_client_async.py +++ /dev/null @@ -1,38 +0,0 @@ -# ------------------------------------ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. -# ------------------------------------ -import time -from unittest.mock import Mock, patch -from urllib.parse import urlparse - -import pytest -from azure.core.credentials import AccessToken -from azure.identity._constants import EnvironmentVariables, DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY -from azure.identity.aio._authn_client import AsyncAuthnClient - -from helpers import mock_response -from helpers_async import wrap_in_future - - -@pytest.mark.asyncio -@pytest.mark.parametrize("authority", ("localhost", "https://localhost")) -async def test_request_url(authority): - tenant_id = "expected-tenant" - parsed_authority = urlparse(authority) - expected_netloc = parsed_authority.netloc or authority # "localhost" parses to netloc "", path "localhost" - - def mock_send(request, **kwargs): - actual = urlparse(request.url) - assert actual.scheme == "https" - assert actual.netloc == expected_netloc - assert actual.path.startswith("/" + tenant_id) - return mock_response(json_payload={"token_type": "Bearer", "expires_in": 42, "access_token": "*"}) - - client = AsyncAuthnClient(tenant=tenant_id, transport=Mock(send=wrap_in_future(mock_send)), authority=authority) - await client.request_token(("scope",)) - - # authority can be configured via environment variable - with patch.dict("os.environ", {EnvironmentVariables.AZURE_AUTHORITY_HOST: authority}, clear=True): - client = AsyncAuthnClient(tenant=tenant_id, transport=Mock(send=wrap_in_future(mock_send))) - await client.request_token(("scope",)) diff --git a/sdk/identity/azure-identity/tests/test_imds_credential.py b/sdk/identity/azure-identity/tests/test_imds_credential.py index 95f53088b11a..808a75c0e70d 100644 --- a/sdk/identity/azure-identity/tests/test_imds_credential.py +++ b/sdk/identity/azure-identity/tests/test_imds_credential.py @@ -4,14 +4,14 @@ # ------------------------------------ import json import time -from azure.core.credentials import AccessToken +from azure.core.credentials import AccessToken from azure.core.exceptions import ClientAuthenticationError + from azure.identity import CredentialUnavailableError -from azure.identity._constants import Endpoints -from azure.identity._credentials.managed_identity import ImdsCredential -import pytest +from azure.identity._credentials.imds import ImdsCredential, IMDS_URL, PIPELINE_SETTINGS from azure.identity._internal.user_agent import USER_AGENT +import pytest from helpers import mock, mock_response, Request, validating_transport @@ -81,7 +81,7 @@ def test_retries(): ) mock_send = mock.Mock(return_value=mock_response) - total_retries = ImdsCredential._create_config().retry_policy.total_retries + total_retries = PIPELINE_SETTINGS["retry_total"] for status_code in (404, 429, 500): mock_send.reset_mock() @@ -145,9 +145,9 @@ def test_identity_config(): scope = "scope" transport = validating_transport( requests=[ - Request(base_url=Endpoints.IMDS), + Request(base_url=IMDS_URL), Request( - base_url=Endpoints.IMDS, + base_url=IMDS_URL, method="GET", required_headers={"Metadata": "true", "User-Agent": USER_AGENT}, required_params={"api-version": "2018-02-01", "resource": scope, param_name: param_value}, diff --git a/sdk/identity/azure-identity/tests/test_imds_credential_async.py b/sdk/identity/azure-identity/tests/test_imds_credential_async.py index a4d056f399fc..7f80f667e8d1 100644 --- a/sdk/identity/azure-identity/tests/test_imds_credential_async.py +++ b/sdk/identity/azure-identity/tests/test_imds_credential_async.py @@ -9,9 +9,9 @@ from azure.core.credentials import AccessToken from azure.core.exceptions import ClientAuthenticationError from azure.identity import CredentialUnavailableError -from azure.identity._constants import Endpoints from azure.identity._internal.user_agent import USER_AGENT -from azure.identity.aio._credentials.managed_identity import ImdsCredential +from azure.identity.aio._credentials.imds import ImdsCredential, PIPELINE_SETTINGS +from azure.identity._credentials.imds import IMDS_URL import pytest from helpers import mock_response, Request @@ -148,7 +148,7 @@ async def test_retries(): ) mock_send = mock.Mock(return_value=mock_response) - total_retries = ImdsCredential._create_config().retry_policy.total_retries + total_retries = PIPELINE_SETTINGS["retry_total"] for status_code in (404, 429, 500): mock_send.reset_mock() @@ -174,9 +174,9 @@ async def test_identity_config(): transport = async_validating_transport( requests=[ - Request(base_url=Endpoints.IMDS), + Request(base_url=IMDS_URL), Request( - base_url=Endpoints.IMDS, + base_url=IMDS_URL, method="GET", required_headers={"Metadata": "true", "User-Agent": USER_AGENT}, required_params={"api-version": "2018-02-01", "resource": scope, param_name: param_value}, diff --git a/sdk/identity/azure-identity/tests/test_managed_identity.py b/sdk/identity/azure-identity/tests/test_managed_identity.py index e64e9389b893..daec08b10764 100644 --- a/sdk/identity/azure-identity/tests/test_managed_identity.py +++ b/sdk/identity/azure-identity/tests/test_managed_identity.py @@ -12,10 +12,10 @@ from azure.core.credentials import AccessToken from azure.core.exceptions import ClientAuthenticationError, ServiceRequestError -from azure.core.pipeline.policies import RetryPolicy from azure.core.pipeline.transport import HttpRequest from azure.identity import ManagedIdentityCredential -from azure.identity._constants import Endpoints, EnvironmentVariables +from azure.identity._constants import EnvironmentVariables +from azure.identity._credentials.imds import IMDS_URL from azure.identity._internal.managed_identity_client import ManagedIdentityClient from azure.identity._internal.user_agent import USER_AGENT import pytest @@ -438,9 +438,9 @@ def test_imds(): scope = "scope" transport = validating_transport( requests=[ - Request(url=Endpoints.IMDS), # first request should be availability probe => match only the URL + Request(base_url=IMDS_URL), # first request should be availability probe => match only the URL Request( - base_url=Endpoints.IMDS, + base_url=IMDS_URL, method="GET", required_headers={"Metadata": "true", "User-Agent": USER_AGENT}, required_params={"api-version": "2018-02-01", "resource": scope}, @@ -532,7 +532,7 @@ def test_imds_user_assigned_identity(): access_token = "****" expires_on = 42 expected_token = AccessToken(access_token, expires_on) - endpoint = Endpoints.IMDS + endpoint = IMDS_URL scope = "scope" client_id = "some-guid" transport = validating_transport( diff --git a/sdk/identity/azure-identity/tests/test_managed_identity_async.py b/sdk/identity/azure-identity/tests/test_managed_identity_async.py index ca1a6f21622b..a93774a24421 100644 --- a/sdk/identity/azure-identity/tests/test_managed_identity_async.py +++ b/sdk/identity/azure-identity/tests/test_managed_identity_async.py @@ -11,7 +11,8 @@ from azure.core.pipeline.transport import HttpRequest from azure.identity.aio import ManagedIdentityCredential from azure.identity.aio._internal.managed_identity_client import AsyncManagedIdentityClient -from azure.identity._constants import Endpoints, EnvironmentVariables +from azure.identity._credentials.imds import IMDS_URL +from azure.identity._constants import EnvironmentVariables from azure.identity._internal.user_agent import USER_AGENT import pytest @@ -460,9 +461,9 @@ async def test_imds(): scope = "scope" transport = async_validating_transport( requests=[ - Request(url=Endpoints.IMDS), # first request should be availability probe => match only the URL + Request(base_url=IMDS_URL), # first request should be availability probe => match only the URL Request( - base_url=Endpoints.IMDS, + base_url=IMDS_URL, method="GET", required_headers={"Metadata": "true", "User-Agent": USER_AGENT}, required_params={"api-version": "2018-02-01", "resource": scope}, @@ -496,14 +497,13 @@ async def test_imds_user_assigned_identity(): access_token = "****" expires_on = 42 expected_token = AccessToken(access_token, expires_on) - url = Endpoints.IMDS scope = "scope" client_id = "some-guid" transport = async_validating_transport( requests=[ - Request(base_url=url), # first request should be availability probe => match only the URL + Request(base_url=IMDS_URL), # first request should be availability probe => match only the URL Request( - base_url=url, + base_url=IMDS_URL, method="GET", required_headers={"Metadata": "true", "User-Agent": USER_AGENT}, required_params={"api-version": "2018-02-01", "client_id": client_id, "resource": scope},