Replace MsiCredential with CloudShellCredential (Azure#15880)

Author: chlowell

sdk/identity/azure-identity/azure/identity/_credentials/cloud_shell.py

@@ -0,0 +1,57 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import functools
+import os
+from typing import TYPE_CHECKING
+
+from azure.core.pipeline.transport import HttpRequest
+
+from .. import CredentialUnavailableError
+from .._constants import EnvironmentVariables
+from .._internal.managed_identity_client import ManagedIdentityClient
+from .._internal.get_token_mixin import GetTokenMixin
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from azure.core.credentials import AccessToken
+
+
+class CloudShellCredential(GetTokenMixin):
+    def __init__(self, **kwargs):
+        # type: (**Any) -> None
+        super(CloudShellCredential, self).__init__()
+        url = os.environ.get(EnvironmentVariables.MSI_ENDPOINT)
+        if url:
+            self._available = True
+            self._client = ManagedIdentityClient(
+                request_factory=functools.partial(_get_request, url),
+                base_headers={"Metadata": "true"},
+                _identity_config=kwargs.pop("identity_config", None),
+                **kwargs
+            )
+        else:
+            self._available = False
+
+    def get_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
+        if not self._available:
+            raise CredentialUnavailableError(
+                message="Cloud Shell managed identity configuration not found in environment"
+            )
+        return super(CloudShellCredential, self).get_token(*scopes, **kwargs)
+
+    def _acquire_token_silently(self, *scopes):
+        # type: (*str) -> Optional[AccessToken]
+        return self._client.get_cached_token(*scopes)
+
+    def _request_token(self, *scopes, **kwargs):
+        # type: (*str, **Any) -> AccessToken
+        return self._client.request_token(*scopes, **kwargs)
+
+
+def _get_request(url, scope, identity_config):
+    # type: (str, str, dict) -> HttpRequest
+    request = HttpRequest("POST", url, data=dict({"resource": scope}, **identity_config))
+    return request

sdk/identity/azure-identity/azure/identity/_credentials/managed_identity.py

@@ -61,8 +61,10 @@ def __init__(self, **kwargs):
 
                 self._credential = AppServiceCredential(**kwargs)
             else:
-                _LOGGER.info("%s will use MSI", self.__class__.__name__)
-                self._credential = MsiCredential(**kwargs)
+                _LOGGER.info("%s will use Cloud Shell managed identity", self.__class__.__name__)
+                from .cloud_shell import CloudShellCredential
+
+                self._credential = CloudShellCredential(**kwargs)
         elif os.environ.get(EnvironmentVariables.IDENTITY_ENDPOINT):
             if (
                 os.environ.get(EnvironmentVariables.IDENTITY_HEADER)
@@ -231,65 +233,3 @@ def _refresh_token(self, *scopes):
             # any other error is unexpected
             six.raise_from(ClientAuthenticationError(message=ex.message, response=ex.response), None)
         return token
-
-
-class MsiCredential(_ManagedIdentityBase):
-    """Authenticates via the MSI endpoint in an App Service or Cloud Shell environment.
-
-    :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
-        self._endpoint = os.environ.get(EnvironmentVariables.MSI_ENDPOINT)
-        if self._endpoint:
-            super(MsiCredential, self).__init__(endpoint=self._endpoint, client_cls=AuthnClient, **kwargs)
-
-    def get_token(self, *scopes, **kwargs):  # pylint:disable=unused-argument
-        # type: (*str, **Any) -> AccessToken
-        """Request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
-
-        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.identity.CredentialUnavailableError: the MSI endpoint is unavailable
-        """
-
-        if not self._endpoint:
-            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.")
-
-        token = self._client.get_cached_token(scopes)
-        if not token:
-            token = self._refresh_token(*scopes)
-        elif self._client.should_refresh(token):
-            try:
-                token = self._refresh_token(*scopes)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
-
-    def _refresh_token(self, *scopes):
-        resource = scopes[0]
-        if resource.endswith("/.default"):
-            resource = resource[: -len("/.default")]
-        secret = os.environ.get(EnvironmentVariables.MSI_SECRET)
-        if secret:
-            # MSI_ENDPOINT and MSI_SECRET set -> App Service
-            token = self._request_app_service_token(scopes=scopes, resource=resource, secret=secret)
-        else:
-            # only MSI_ENDPOINT set -> legacy-style MSI (Cloud Shell)
-            token = self._request_legacy_token(scopes=scopes, resource=resource)
-        return token
-
-    def _request_app_service_token(self, scopes, resource, secret):
-        params = dict({"api-version": "2017-09-01", "resource": resource}, **self._identity_config)
-        return self._client.request_token(scopes, method="GET", headers={"secret": secret}, params=params)
-
-    def _request_legacy_token(self, scopes, resource):
-        form_data = dict({"resource": resource}, **self._identity_config)
-        return self._client.request_token(scopes, method="POST", form_data=form_data)

sdk/identity/azure-identity/azure/identity/aio/_credentials/cloud_shell.py

@@ -0,0 +1,50 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import functools
+import os
+from typing import TYPE_CHECKING
+
+from .._internal import AsyncContextManager
+from .._internal.get_token_mixin import GetTokenMixin
+from .._internal.managed_identity_client import AsyncManagedIdentityClient
+from ... import CredentialUnavailableError
+from ..._constants import EnvironmentVariables
+from ..._credentials.cloud_shell import _get_request
+
+if TYPE_CHECKING:
+    from typing import Any, Optional
+    from azure.core.credentials import AccessToken
+
+
+class CloudShellCredential(AsyncContextManager, GetTokenMixin):
+    def __init__(self, **kwargs: "Any") -> None:
+        super(CloudShellCredential, self).__init__()
+        url = os.environ.get(EnvironmentVariables.MSI_ENDPOINT)
+        if url:
+            self._available = True
+            self._client = AsyncManagedIdentityClient(
+                request_factory=functools.partial(_get_request, url),
+                base_headers={"Metadata": "true"},
+                _identity_config=kwargs.pop("identity_config", None),
+                **kwargs,
+            )
+        else:
+            self._available = False
+
+    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        if not self._available:
+            raise CredentialUnavailableError(
+                message="Cloud Shell managed identity configuration not found in environment"
+            )
+        return await super(CloudShellCredential, self).get_token(*scopes, **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: str, **kwargs: "Any") -> "AccessToken":
+        return await self._client.request_token(*scopes, **kwargs)

sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py

@@ -49,8 +49,10 @@ def __init__(self, **kwargs: "Any") -> None:
 
                 self._credential = AppServiceCredential(**kwargs)
             else:
-                _LOGGER.info("%s will use MSI", self.__class__.__name__)
-                self._credential = MsiCredential(**kwargs)
+                _LOGGER.info("%s will use Cloud Shell managed identity", self.__class__.__name__)
+                from .cloud_shell import CloudShellCredential
+
+                self._credential = CloudShellCredential(**kwargs)
         elif os.environ.get(EnvironmentVariables.IDENTITY_ENDPOINT):
             if (
                 os.environ.get(EnvironmentVariables.IDENTITY_HEADER)
@@ -192,63 +194,3 @@ async def _refresh_token(self, *scopes):
             # any other error is unexpected
             raise ClientAuthenticationError(message=ex.message, response=ex.response) from None
         return token
-
-
-class MsiCredential(_AsyncManagedIdentityBase):
-    """Authenticates via the MSI endpoint in an App Service or Cloud Shell environment.
-
-    :keyword str client_id: ID of a user-assigned identity. Leave unspecified to use a system-assigned identity.
-    """
-
-    def __init__(self, **kwargs: "Any") -> None:
-        self._endpoint = os.environ.get(EnvironmentVariables.MSI_ENDPOINT)
-        if self._endpoint:
-            super().__init__(endpoint=self._endpoint, **kwargs)
-
-    async def get_token(self, *scopes: str, **kwargs: "Any") -> AccessToken:  # pylint:disable=unused-argument
-        """Asynchronously request an access token for `scopes`.
-
-        This method is called automatically by Azure SDK clients.
-
-        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
-        :rtype: :class:`azure.core.credentials.AccessToken`
-        :raises ~azure.identity.CredentialUnavailableError: the MSI endpoint is unavailable
-        """
-        if not self._endpoint:
-            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.")
-
-        token = self._client.get_cached_token(scopes)
-        if not token:
-            token = await self._refresh_token(*scopes)
-        elif self._client.should_refresh(token):
-            try:
-                token = await self._refresh_token(*scopes)
-            except Exception:  # pylint: disable=broad-except
-                pass
-        return token
-
-    async def _refresh_token(self, *scopes):
-        resource = scopes[0]
-        if resource.endswith("/.default"):
-            resource = resource[: -len("/.default")]
-
-        secret = os.environ.get(EnvironmentVariables.MSI_SECRET)
-        if secret:
-            # MSI_ENDPOINT and MSI_SECRET set -> App Service
-            token = await self._request_app_service_token(scopes=scopes, resource=resource, secret=secret)
-        else:
-            # only MSI_ENDPOINT set -> legacy-style MSI (Cloud Shell)
-            token = await self._request_legacy_token(scopes=scopes, resource=resource)
-        return token
-
-    async def _request_app_service_token(self, scopes, resource, secret):
-        params = {"api-version": "2017-09-01", "resource": resource, **self._identity_config}
-        return await self._client.request_token(scopes, method="GET", headers={"secret": secret}, params=params)
-
-    async def _request_legacy_token(self, scopes, resource):
-        form_data = {"resource": resource, **self._identity_config}
-        return await self._client.request_token(scopes, method="POST", form_data=form_data)