[Identity] Propagate additional token type values (#37579)

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

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

@@ -136,7 +136,10 @@ def _acquire_token_silently(self, *scopes: str, **kwargs: Any) -> Optional[Acces
                 if result and "access_token" in result and "expires_in" in result:
                     refresh_on = int(result["refresh_on"]) if "refresh_on" in result else None
                     return AccessTokenInfo(
-                        result["access_token"], now + int(result["expires_in"]), refresh_on=refresh_on
+                        result["access_token"],
+                        now + int(result["expires_in"]),
+                        token_type=result.get("token_type", "Bearer"),
+                        refresh_on=refresh_on,
                     )
 
         return None
@@ -157,4 +160,9 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
             pass  # non-fatal; we'll use the assertion again next time instead of a refresh token
 
         refresh_on = int(result["refresh_on"]) if "refresh_on" in result else None
-        return AccessTokenInfo(result["access_token"], request_time + int(result["expires_in"]), refresh_on=refresh_on)
+        return AccessTokenInfo(
+            result["access_token"],
+            request_time + int(result["expires_in"]),
+            token_type=result.get("token_type", "Bearer"),
+            refresh_on=refresh_on,
+        )

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

@@ -186,8 +186,15 @@ def _acquire_token_silent(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
             result = client_application.acquire_token_silent_with_error(
                 list(scopes), account=account, claims_challenge=kwargs.get("claims")
             )
+
             if result and "access_token" in result and "expires_in" in result:
-                return AccessTokenInfo(result["access_token"], now + int(result["expires_in"]))
+                refresh_on = int(result["refresh_on"]) if "refresh_on" in result else None
+                return AccessTokenInfo(
+                    result["access_token"],
+                    now + int(result["expires_in"]),
+                    token_type=result.get("token_type", "Bearer"),
+                    refresh_on=refresh_on,
+                )
 
         # if we get this far, the cache contained a matching account but MSAL failed to authenticate it silently
         if result:

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

@@ -95,7 +95,9 @@ def get_cached_access_token(self, scopes: Iterable[str], **kwargs: Any) -> Optio
             expires_on = int(token["expires_on"])
             if expires_on > int(time.time()):
                 refresh_on = int(token["refresh_on"]) if "refresh_on" in token else None
-                return AccessTokenInfo(token["secret"], expires_on, refresh_on=refresh_on)
+                return AccessTokenInfo(
+                    token["secret"], expires_on, token_type=token.get("token_type", "Bearer"), refresh_on=refresh_on
+                )
         return None
 
     def get_cached_refresh_tokens(self, scopes: Iterable[str], **kwargs) -> List[Dict]:
@@ -178,7 +180,9 @@ def _process_response(self, response: PipelineResponse, request_time: int, **kwa
             content["refresh_in"] = expires_in // 2
 
         refresh_on = request_time + int(content["refresh_in"]) if "refresh_in" in content else None
-        token = AccessTokenInfo(content["access_token"], expires_on, refresh_on=refresh_on)
+        token = AccessTokenInfo(
+            content["access_token"], expires_on, token_type=content.get("token_type", "Bearer"), refresh_on=refresh_on
+        )
 
         # caching is the final step because 'add' mutates 'content'
         cache.add(

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

@@ -34,6 +34,7 @@ def _acquire_token_silently(self, *scopes: str, **kwargs: Any) -> Optional[Acces
             return AccessTokenInfo(
                 result["access_token"],
                 request_time + int(result["expires_in"]),
+                token_type=result.get("token_type", "Bearer"),
                 refresh_on=refresh_on,
             )
         return None
@@ -51,5 +52,6 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
         return AccessTokenInfo(
             result["access_token"],
             request_time + int(result["expires_in"]),
+            token_type=result.get("token_type", "Bearer"),
             refresh_on=refresh_on,
         )

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

@@ -286,7 +286,10 @@ def _acquire_token_silent(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:
                 if result and "access_token" in result and "expires_in" in result:
                     refresh_on = int(result["refresh_on"]) if "refresh_on" in result else None
                     return AccessTokenInfo(
-                        result["access_token"], now + int(result["expires_in"]), refresh_on=refresh_on
+                        result["access_token"],
+                        now + int(result["expires_in"]),
+                        token_type=result.get("token_type", "Bearer"),
+                        refresh_on=refresh_on,
                     )
 
         # if we get this far, result is either None or the content of a Microsoft Entra ID error response

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

@@ -76,7 +76,12 @@ def _process_response(self, response: PipelineResponse, request_time: int) -> Ac
             content["refresh_in"] = expires_in // 2
 
         refresh_on = request_time + int(content["refresh_in"]) if "refresh_in" in content else None
-        token = AccessTokenInfo(content["access_token"], content["expires_on"], refresh_on=refresh_on)
+        token = AccessTokenInfo(
+            content["access_token"],
+            content["expires_on"],
+            token_type=content.get("token_type", "Bearer"),
+            refresh_on=refresh_on,
+        )
 
         # caching is the final step because TokenCache.add mutates its "event"
         self._cache.add(
@@ -93,7 +98,9 @@ def get_cached_token(self, *scopes: str) -> Optional[AccessTokenInfo]:
             expires_on = int(token["expires_on"])
             refresh_on = int(token["refresh_on"]) if "refresh_on" in token else None
             if expires_on > now and (not refresh_on or refresh_on > now):
-                return AccessTokenInfo(token["secret"], expires_on, refresh_on=refresh_on)
+                return AccessTokenInfo(
+                    token["secret"], expires_on, token_type=token.get("token_type", "Bearer"), refresh_on=refresh_on
+                )
 
         return None
 

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

@@ -53,7 +53,12 @@ def _request_token(self, *scopes: str, **kwargs: Any) -> AccessTokenInfo:  # pyl
         now = int(time.time())
         if result and "access_token" in result and "expires_in" in result:
             refresh_on = int(result["refresh_on"]) if "refresh_on" in result else None
-            return AccessTokenInfo(result["access_token"], now + int(result["expires_in"]), refresh_on=refresh_on)
+            return AccessTokenInfo(
+                result["access_token"],
+                now + int(result["expires_in"]),
+                token_type=result.get("token_type", "Bearer"),
+                refresh_on=refresh_on,
+            )
         if result and "error" in result:
             error_desc = cast(str, result["error"])
         error_message = self.get_unavailable_message(error_desc)

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

@@ -243,7 +243,9 @@ def _get_cached_access_token(
                 expires_on = int(token["expires_on"])
                 refresh_on = int(token["refresh_on"]) if "refresh_on" in token else None
                 if expires_on - 300 > int(time.time()):
-                    return AccessTokenInfo(token["secret"], expires_on, refresh_on=refresh_on)
+                    return AccessTokenInfo(
+                        token["secret"], expires_on, token_type=token.get("token_type", "Bearer"), refresh_on=refresh_on
+                    )
         except Exception as ex:  # pylint:disable=broad-except
             message = "Error accessing cached data: {}".format(ex)
             raise CredentialUnavailableError(message=message) from ex