Refactor OAuth flows to use OAuth2 as a base class.

httpx_auth/_oauth2/authentication_responses_server.py

@@ -166,7 +166,7 @@ def handle_timeout(self) -> None:
         raise TimeoutOccurred(self.timeout)
 
 
-def request_new_grant(grant_details: GrantDetails) -> (str, str):
+def request_new_grant(grant_details: GrantDetails) -> tuple[str, str]:
     """
     Ask for a new OAuth2 grant.
     :return: A tuple (state, grant)

httpx_auth/_oauth2/authorization_code.py

@@ -1,5 +1,5 @@
 from hashlib import sha512
-from typing import Generator, Iterable, Union
+from typing import Iterable, Union
 
 import httpx
 
@@ -8,14 +8,14 @@
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
     request_new_grant_with_post,
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
     _get_query_parameter,
 )
 
 
-class OAuth2AuthorizationCode(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2AuthorizationCode(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Authorization Code Grant
 
@@ -71,13 +71,11 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
 
         BrowserAuth.__init__(self, kwargs)
 
-        self.header_name = kwargs.pop("header_name", None) or "Authorization"
-        self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
-        if "{token}" not in self.header_value:
-            raise Exception("header_value parameter must contains {token}.")
+        header_name = kwargs.pop("header_name", None) or "Authorization"
+        header_value = kwargs.pop("header_value", None) or "Bearer {token}"
 
         self.token_field_name = kwargs.pop("token_field_name", None) or "access_token"
-        self.early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
+        early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
 
         username = kwargs.pop("username", None)
         password = kwargs.pop("password", None)
@@ -99,11 +97,11 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         authorization_url_without_nonce, nonce = _pop_parameter(
             authorization_url_without_nonce, "nonce"
         )
-        self.state = sha512(
+        state = sha512(
             authorization_url_without_nonce.encode("unicode_escape")
         ).hexdigest()
         custom_code_parameters = {
-            "state": self.state,
+            "state": state,
             "redirect_uri": self.redirect_uri,
         }
         if nonce:
@@ -129,17 +127,14 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         self.refresh_data = {"grant_type": "refresh_token"}
         self.refresh_data.update(kwargs)
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-            on_expired_token=self.refresh_token,
+        OAuth2BaseAuth.__init__(
+            self,
+            state,
+            early_expiry,
+            header_name,
+            header_value,
+            self.refresh_token,
         )
-        request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         # Request code

httpx_auth/_oauth2/authorization_code_pkce.py

@@ -1,7 +1,6 @@
 import base64
 import os
 from hashlib import sha256, sha512
-from typing import Generator
 
 import httpx
 
@@ -10,13 +9,13 @@
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
     request_new_grant_with_post,
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
 )
 
 
-class OAuth2AuthorizationCodePKCE(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2AuthorizationCodePKCE(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Proof Key for Code Exchange
 
@@ -72,13 +71,11 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
 
         self.client = kwargs.pop("client", None)
 
-        self.header_name = kwargs.pop("header_name", None) or "Authorization"
-        self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
-        if "{token}" not in self.header_value:
-            raise Exception("header_value parameter must contains {token}.")
+        header_name = kwargs.pop("header_name", None) or "Authorization"
+        header_value = kwargs.pop("header_value", None) or "Bearer {token}"
 
         self.token_field_name = kwargs.pop("token_field_name", None) or "access_token"
-        self.early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
+        early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
 
         # As described in https://tools.ietf.org/html/rfc6749#section-4.1.2
         code_field_name = kwargs.pop("code_field_name", "code")
@@ -98,11 +95,11 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         authorization_url_without_nonce, nonce = _pop_parameter(
             authorization_url_without_nonce, "nonce"
         )
-        self.state = sha512(
+        state = sha512(
             authorization_url_without_nonce.encode("unicode_escape")
         ).hexdigest()
         custom_code_parameters = {
-            "state": self.state,
+            "state": state,
             "redirect_uri": self.redirect_uri,
         }
         if nonce:
@@ -139,17 +136,9 @@ def __init__(self, authorization_url: str, token_url: str, **kwargs):
         self.refresh_data = {"grant_type": "refresh_token"}
         self.refresh_data.update(kwargs)
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
-            on_expired_token=self.refresh_token,
+        OAuth2BaseAuth.__init__(
+            self, state, early_expiry, header_name, header_value, self.refresh_token
         )
-        request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         # Request code

httpx_auth/_oauth2/client_credentials.py

@@ -1,16 +1,16 @@
 from hashlib import sha512
-from typing import Generator, Union, Iterable
+from typing import Union, Iterable
 
 import httpx
 from httpx_auth._authentication import SupportMultiAuth
 from httpx_auth._oauth2.common import (
-    OAuth2,
+    OAuth2BaseAuth,
     request_new_grant_with_post,
     _add_parameters,
 )
 
 
-class OAuth2ClientCredentials(httpx.Auth, SupportMultiAuth):
+class OAuth2ClientCredentials(OAuth2BaseAuth, SupportMultiAuth):
     """
     Client Credentials Grant
 
@@ -49,13 +49,11 @@ def __init__(self, token_url: str, client_id: str, client_secret: str, **kwargs)
         if not self.client_secret:
             raise Exception("client_secret is mandatory.")
 
-        self.header_name = kwargs.pop("header_name", None) or "Authorization"
-        self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
-        if "{token}" not in self.header_value:
-            raise Exception("header_value parameter must contains {token}.")
+        header_name = kwargs.pop("header_name", None) or "Authorization"
+        header_value = kwargs.pop("header_value", None) or "Bearer {token}"
 
         self.token_field_name = kwargs.pop("token_field_name", None) or "access_token"
-        self.early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
+        early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
 
         # Time is expressed in seconds
         self.timeout = int(kwargs.pop("timeout", None) or 60)
@@ -70,18 +68,14 @@ def __init__(self, token_url: str, client_id: str, client_secret: str, **kwargs)
         self.data.update(kwargs)
 
         all_parameters_in_url = _add_parameters(self.token_url, self.data)
-        self.state = sha512(all_parameters_in_url.encode("unicode_escape")).hexdigest()
-
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=self.request_new_token,
+        state = sha512(all_parameters_in_url.encode("unicode_escape")).hexdigest()
+
+        super().__init__(
+            state,
+            early_expiry,
+            header_name,
+            header_value,
         )
-        request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
 
     def request_new_token(self) -> tuple:
         client = self.client or httpx.Client()

httpx_auth/_oauth2/common.py

@@ -1,4 +1,5 @@
-from typing import Optional
+import abc
+from typing import Callable, Generator, Optional, Union
 from urllib.parse import parse_qs, urlsplit, urlunsplit, urlencode
 
 import httpx
@@ -86,3 +87,41 @@ def request_new_grant_with_post(
 class OAuth2:
     token_cache = TokenMemoryCache()
     display = DisplaySettings()
+
+
+class OAuth2BaseAuth(abc.ABC, httpx.Auth):
+    def __init__(
+        self,
+        state: str,
+        early_expiry: float,
+        header_name: str,
+        header_value: str,
+        refresh_token: Optional[Callable] = None,
+    ) -> None:
+        if "{token}" not in header_value:
+            raise Exception("header_value parameter must contains {token}.")
+
+        self.state = state
+        self.early_expiry = early_expiry
+        self.header_name = header_name
+        self.header_value = header_value
+        self.refresh_token = refresh_token
+
+    def auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response, None]:
+        token = OAuth2.token_cache.get_token(
+            self.state,
+            early_expiry=self.early_expiry,
+            on_missing_token=self.request_new_token,
+            on_expired_token=self.refresh_token,
+        )
+        self._update_user_request(request, token)
+        yield request
+
+    @abc.abstractmethod
+    def request_new_token(self) -> Union[tuple[str, str], tuple[str, str, int]]:
+        pass  # pragma: no cover
+
+    def _update_user_request(self, request: httpx.Request, token: str) -> None:
+        request.headers[self.header_name] = self.header_value.format(token=token)

httpx_auth/_oauth2/implicit.py

@@ -1,21 +1,20 @@
 import uuid
 from hashlib import sha512
-from typing import Generator
 
 import httpx
 
 from httpx_auth._authentication import SupportMultiAuth
 from httpx_auth._oauth2 import authentication_responses_server
 from httpx_auth._oauth2.browser import BrowserAuth
 from httpx_auth._oauth2.common import (
-    OAuth2,
+    OAuth2BaseAuth,
     _add_parameters,
     _pop_parameter,
     _get_query_parameter,
 )
 
 
-class OAuth2Implicit(httpx.Auth, SupportMultiAuth, BrowserAuth):
+class OAuth2Implicit(OAuth2BaseAuth, SupportMultiAuth, BrowserAuth):
     """
     Implicit Grant
 
@@ -62,10 +61,8 @@ def __init__(self, authorization_url: str, **kwargs):
 
         BrowserAuth.__init__(self, kwargs)
 
-        self.header_name = kwargs.pop("header_name", None) or "Authorization"
-        self.header_value = kwargs.pop("header_value", None) or "Bearer {token}"
-        if "{token}" not in self.header_value:
-            raise Exception("header_value parameter must contains {token}.")
+        header_name = kwargs.pop("header_name", None) or "Authorization"
+        header_value = kwargs.pop("header_value", None) or "Bearer {token}"
 
         response_type = _get_query_parameter(self.authorization_url, "response_type")
         if response_type:
@@ -82,18 +79,18 @@ def __init__(self, authorization_url: str, **kwargs):
                 "id_token" if "id_token" == response_type else "access_token"
             )
 
-        self.early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
+        early_expiry = float(kwargs.pop("early_expiry", None) or 30.0)
 
         authorization_url_without_nonce = _add_parameters(
             self.authorization_url, kwargs
         )
         authorization_url_without_nonce, nonce = _pop_parameter(
             authorization_url_without_nonce, "nonce"
         )
-        self.state = sha512(
+        state = sha512(
             authorization_url_without_nonce.encode("unicode_escape")
         ).hexdigest()
-        custom_parameters = {"state": self.state, "redirect_uri": self.redirect_uri}
+        custom_parameters = {"state": state, "redirect_uri": self.redirect_uri}
         if nonce:
             custom_parameters["nonce"] = nonce
         grant_url = _add_parameters(authorization_url_without_nonce, custom_parameters)
@@ -104,17 +101,16 @@ def __init__(self, authorization_url: str, **kwargs):
             self.redirect_uri_port,
         )
 
-    def auth_flow(
-        self, request: httpx.Request
-    ) -> Generator[httpx.Request, httpx.Response, None]:
-        token = OAuth2.token_cache.get_token(
-            self.state,
-            early_expiry=self.early_expiry,
-            on_missing_token=authentication_responses_server.request_new_grant,
-            grant_details=self.grant_details,
+        OAuth2BaseAuth.__init__(
+            self,
+            state,
+            early_expiry,
+            header_name,
+            header_value,
         )
-        request.headers[self.header_name] = self.header_value.format(token=token)
-        yield request
+
+    def request_new_token(self) -> tuple[str, str]:
+        return authentication_responses_server.request_new_grant(self.grant_details)
 
 
 class AzureActiveDirectoryImplicit(OAuth2Implicit):