Identity msal mi mig (#36225)

* msal managed identity

* update

* update

* update

* updates

* updates

* updates

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* updates

* Update sdk/identity/azure-identity/CHANGELOG.md

Co-authored-by: Paul Van Eck <[email protected]>

* updates

* update

* update

---------

Co-authored-by: Paul Van Eck <[email protected]>

Author: xiangyan99

sdk/identity/azure-identity/CHANGELOG.md

@@ -1,17 +1,13 @@
 # Release History
 
-## 1.17.2 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
-
-### Bugs Fixed
+## 1.18.0b1 (2024-07-16)
 
 - Fixed the issue that `SharedTokenCacheCredential` was not picklable.
 
 ### Other Changes
 
+- The synchronous `ManagedIdentityCredential` was updated to use MSAL (Microsoft Authentication Library) for handling most of the underlying managed identity implementations.
+
 ## 1.17.1 (2024-06-21)
 
 ### Bugs Fixed

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

@@ -5,23 +5,15 @@
 import functools
 import os
 from typing import Optional, Dict, Any
-
 from azure.core.pipeline.transport import HttpRequest
 
 from .._constants import EnvironmentVariables
-from .._internal.managed_identity_base import ManagedIdentityBase
-from .._internal.managed_identity_client import ManagedIdentityClient
-
+from .._internal.msal_managed_identity_client import MsalManagedIdentityClient
 
-class AppServiceCredential(ManagedIdentityBase):
-    def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
-        client_args = _get_client_args(**kwargs)
-        if client_args:
-            return ManagedIdentityClient(**client_args)
-        return None
 
-    def get_unavailable_message(self) -> str:
-        return "App Service managed identity configuration not found in environment"
+class AppServiceCredential(MsalManagedIdentityClient):
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"App Service managed identity configuration not found in environment. {desc}"
 
 
 def _get_client_args(**kwargs: Any) -> Optional[Dict]:

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

@@ -2,47 +2,21 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-import functools
 import os
 import sys
-from typing import Any, Dict, Optional
+from typing import Dict
 
 from azure.core.exceptions import ClientAuthenticationError
 from azure.core.pipeline.transport import HttpRequest
 from azure.core.pipeline.policies import HTTPPolicy
 from azure.core.pipeline import PipelineRequest, PipelineResponse
 
-from .._constants import EnvironmentVariables
-from .._internal.managed_identity_base import ManagedIdentityBase
-from .._internal.managed_identity_client import ManagedIdentityClient
+from .._internal.msal_managed_identity_client import MsalManagedIdentityClient
 
 
-class AzureArcCredential(ManagedIdentityBase):
-    def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
-        url = os.environ.get(EnvironmentVariables.IDENTITY_ENDPOINT)
-        imds = os.environ.get(EnvironmentVariables.IMDS_ENDPOINT)
-        if url and imds:
-            return ManagedIdentityClient(
-                _per_retry_policies=[ArcChallengeAuthPolicy()],
-                request_factory=functools.partial(_get_request, url),
-                **kwargs,
-            )
-        return None
-
-    def __enter__(self) -> "AzureArcCredential":
-        if self._client:
-            self._client.__enter__()
-        return self
-
-    def __exit__(self, *args: Any) -> None:
-        if self._client:
-            self._client.__exit__(*args)
-
-    def close(self) -> None:
-        self.__exit__()
-
-    def get_unavailable_message(self) -> str:
-        return "Azure Arc managed identity configuration not found in environment"
+class AzureArcCredential(MsalManagedIdentityClient):
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"Azure Arc managed identity configuration not found in environment. {desc}"
 
 
 def _get_request(url: str, scope: str, identity_config: Dict) -> HttpRequest:
@@ -54,7 +28,7 @@ def _get_request(url: str, scope: str, identity_config: Dict) -> HttpRequest:
         )
 
     request = HttpRequest("GET", url)
-    request.format_parameters(dict({"api-version": "2019-11-01", "resource": scope}, **identity_config))
+    request.format_parameters(dict({"api-version": "2020-06-01", "resource": scope}, **identity_config))
     return request
 
 

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

@@ -9,19 +9,12 @@
 from azure.core.pipeline.transport import HttpRequest
 
 from .._constants import EnvironmentVariables
-from .._internal.managed_identity_base import ManagedIdentityBase
-from .._internal.managed_identity_client import ManagedIdentityClient
+from .._internal.msal_managed_identity_client import MsalManagedIdentityClient
 
 
-class AzureMLCredential(ManagedIdentityBase):
-    def get_client(self, **kwargs) -> Optional[ManagedIdentityClient]:
-        client_args = _get_client_args(**kwargs)
-        if client_args:
-            return ManagedIdentityClient(**client_args)
-        return None
-
-    def get_unavailable_message(self) -> str:
-        return "Azure ML managed identity configuration not found in environment"
+class AzureMLCredential(MsalManagedIdentityClient):
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"Azure ML managed identity configuration not found in environment. {desc}"
 
 
 def _get_client_args(**kwargs) -> Optional[Dict]:

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

@@ -22,8 +22,8 @@ def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
             )
         return None
 
-    def get_unavailable_message(self) -> str:
-        return "Cloud Shell managed identity configuration not found in environment"
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"Cloud Shell managed identity configuration not found in environment. {desc}"
 
 
 def _get_request(url: str, scope: str, identity_config: Dict) -> HttpRequest:

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

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 import os
+import json
 from typing import Any, Optional, Dict
 
 from azure.core.exceptions import ClientAuthenticationError, HttpResponseError
@@ -12,8 +13,8 @@
 from .. import CredentialUnavailableError
 from .._constants import EnvironmentVariables
 from .._internal import within_credential_chain
-from .._internal.get_token_mixin import GetTokenMixin
 from .._internal.managed_identity_client import ManagedIdentityClient
+from .._internal.msal_managed_identity_client import MsalManagedIdentityClient
 
 
 IMDS_AUTHORITY = "http://169.254.169.254"
@@ -55,16 +56,15 @@ def _check_forbidden_response(ex: HttpResponseError) -> None:
             raise CredentialUnavailableError(message=error_message) from ex
 
 
-class ImdsCredential(GetTokenMixin):
+class ImdsCredential(MsalManagedIdentityClient):
     def __init__(self, **kwargs: Any) -> None:
-        super(ImdsCredential, self).__init__()
+        super(ImdsCredential, self).__init__(**kwargs)
+        self._config = kwargs
 
-        self._client = ManagedIdentityClient(_get_request, **dict(PIPELINE_SETTINGS, **kwargs))
         if EnvironmentVariables.AZURE_POD_IDENTITY_AUTHORITY_HOST in os.environ:
             self._endpoint_available: Optional[bool] = True
         else:
             self._endpoint_available = None
-        self._user_assigned_identity = "client_id" in kwargs or "identity_config" in kwargs
 
     def __enter__(self) -> "ImdsCredential":
         self._client.__enter__()
@@ -76,16 +76,14 @@ def __exit__(self, *args):
     def close(self) -> None:
         self.__exit__()
 
-    def _acquire_token_silently(self, *scopes: str, **kwargs: Any) -> Optional[AccessToken]:
-        return self._client.get_cached_token(*scopes)
-
     def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
 
         if within_credential_chain.get() and not self._endpoint_available:
             # If within a chain (e.g. DefaultAzureCredential), we do a quick check to see if the IMDS endpoint
             # is available to avoid hanging for a long time if the endpoint isn't available.
             try:
-                self._client.request_token(*scopes, connection_timeout=1, retry_total=0)
+                client = ManagedIdentityClient(_get_request, **dict(PIPELINE_SETTINGS, **self._config))
+                client.request_token(*scopes, connection_timeout=1, retry_total=0)
                 self._endpoint_available = True
             except HttpResponseError as ex:
                 # IMDS responded
@@ -98,19 +96,18 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
                 raise CredentialUnavailableError(error_message) from ex
 
         try:
-            token = self._client.request_token(*scopes, headers={"Metadata": "true"})
+            token = super()._request_token(*scopes)
         except CredentialUnavailableError:
             # Response is not json, skip the IMDS credential
             raise
         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:
-                error_message = "ManagedIdentityCredential authentication unavailable. "
-                if self._user_assigned_identity:
-                    error_message += "The requested identity has not been assigned to this resource."
-                else:
-                    error_message += "No identity has been assigned to this resource."
+                error_message = (
+                    "ManagedIdentityCredential authentication unavailable. "
+                    "No identity has been assigned to this resource."
+                )
 
                 if ex.message:
                     error_message += f" Error: {ex.message}"
@@ -120,8 +117,13 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
             _check_forbidden_response(ex)
             # any other error is unexpected
             raise ClientAuthenticationError(message=ex.message, response=ex.response) from ex
+        except json.decoder.JSONDecodeError as ex:
+            raise CredentialUnavailableError(message="ManagedIdentityCredential authentication unavailable.") from ex
         except Exception as ex:  # pylint:disable=broad-except
             # if anything else was raised, assume the endpoint is unavailable
             error_message = "ManagedIdentityCredential authentication unavailable, no response from the IMDS endpoint."
             raise CredentialUnavailableError(error_message) from ex
         return token
+
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"ManagedIdentityCredential authentication unavailable, no response from the IMDS endpoint. {desc}"

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

@@ -9,19 +9,12 @@
 from azure.core.pipeline.transport import HttpRequest
 
 from .._constants import EnvironmentVariables
-from .._internal.managed_identity_base import ManagedIdentityBase
-from .._internal.managed_identity_client import ManagedIdentityClient
+from .._internal.msal_managed_identity_client import MsalManagedIdentityClient
 
 
-class ServiceFabricCredential(ManagedIdentityBase):
-    def get_client(self, **kwargs) -> Optional[ManagedIdentityClient]:
-        client_args = _get_client_args(**kwargs)
-        if client_args:
-            return ManagedIdentityClient(**client_args)
-        return None
-
-    def get_unavailable_message(self) -> str:
-        return "Service Fabric managed identity configuration not found in environment"
+class ServiceFabricCredential(MsalManagedIdentityClient):
+    def get_unavailable_message(self, desc: str = "") -> str:
+        return f"Service Fabric managed identity configuration not found in environment. {desc}"
 
 
 def _get_client_args(**kwargs: Any) -> Optional[Dict]:

sdk/identity/azure-identity/azure/identity/_internal/managed_identity_base.py

@@ -25,8 +25,7 @@ def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
         pass
 
     @abc.abstractmethod
-    def get_unavailable_message(self):
-        # type: () -> str
+    def get_unavailable_message(self, desc: str = "") -> str:
         pass
 
     def __enter__(self: T) -> T:

sdk/identity/azure-identity/azure/identity/_internal/msal_client.py

@@ -27,6 +27,7 @@ class MsalResponse:
 
     def __init__(self, response: PipelineResponse) -> None:
         self._response = response
+        self.headers = response.http_response.headers if response.http_response else {}
 
     @property
     def status_code(self) -> int: