Implement a custom strategy for using the user_data method

Author: ArtyomVancyan

src/fastapi_oauth2/core.py

@@ -1,13 +1,18 @@
 import json
+import random
 import re
+import string
 from typing import Any
 from typing import Dict
 from typing import List
 from typing import Optional
 from urllib.parse import urljoin
 
 import httpx
+import requests
 from oauthlib.oauth2 import WebApplicationClient
+from social_core.backends.oauth import BaseOAuth2
+from social_core.strategy import BaseStrategy
 from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import RedirectResponse
@@ -21,31 +26,42 @@ class OAuth2LoginError(HTTPException):
     """
 
 
+class OAuth2Strategy(BaseStrategy):
+    def request_data(self, merge=True):
+        return {}
+
+    def absolute_uri(self, path=None):
+        return path
+
+    def get_setting(self, name):
+        return None
+
+    @staticmethod
+    def get_json(url, method='GET', *args, **kwargs):
+        return requests.request(method, url, *args, **kwargs)
+
+
 class OAuth2Core:
     """Base class (mixin) for all SSO providers"""
 
     client_id: str = None
     client_secret: str = None
     callback_url: Optional[str] = None
-    allow_http: bool = False
     scope: Optional[List[str]] = None
-    state: Optional[str] = None
+    backend: BaseOAuth2 = None
     _oauth_client: Optional[WebApplicationClient] = None
-    additional_headers: Optional[Dict[str, Any]] = None
 
     authorization_endpoint: str = None
     token_endpoint: str = None
-    userinfo_endpoint: str = None
 
     def __init__(self, client: OAuth2Client) -> None:
         self.client_id = client.client_id
         self.client_secret = client.client_secret
         self.scope = client.scope or self.scope
         self.provider = client.backend.name
+        self.backend = client.backend(OAuth2Strategy())
         self.authorization_endpoint = client.backend.AUTHORIZATION_URL
         self.token_endpoint = client.backend.ACCESS_TOKEN_URL
-        self.userinfo_endpoint = "https://api.github.com/user"
-        self.additional_headers = {"Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json"}
 
     @property
     def oauth_client(self) -> WebApplicationClient:
@@ -56,47 +72,24 @@ def oauth_client(self) -> WebApplicationClient:
     def get_redirect_uri(self, request: Request) -> str:
         return urljoin(str(request.base_url), "/oauth2/%s/token" % self.provider)
 
-    async def get_login_url(
-            self,
-            request: Request,
-            *,
-            params: Optional[Dict[str, Any]] = None,
-            state: Optional[str] = None,
-    ) -> Any:
-        self.state = state
-        params = params or {}
+    async def get_login_url(self, request: Request) -> Any:
         redirect_uri = self.get_redirect_uri(request)
+        state = "".join([random.choice(string.ascii_letters) for _ in range(32)])
         return self.oauth_client.prepare_request_uri(
-            self.authorization_endpoint, redirect_uri=redirect_uri, state=state, scope=self.scope, **params
+            self.authorization_endpoint, redirect_uri=redirect_uri, state=state, scope=self.scope
         )
 
-    async def login_redirect(
-            self,
-            request: Request,
-            *,
-            params: Optional[Dict[str, Any]] = None,
-            state: Optional[str] = None,
-    ) -> RedirectResponse:
-        login_uri = await self.get_login_url(request, params=params, state=state)
-        return RedirectResponse(login_uri, 303)
-
-    async def get_token_data(
-            self,
-            request: Request,
-            *,
-            params: Optional[Dict[str, Any]] = None,
-            headers: Optional[Dict[str, Any]] = None,
-    ) -> Optional[Dict[str, Any]]:
-        params = params or {}
-        additional_headers = headers or {}
-        additional_headers.update(self.additional_headers or {})
+    async def login_redirect(self, request: Request) -> RedirectResponse:
+        return RedirectResponse(await self.get_login_url(request), 303)
+
+    async def get_token_data(self, request: Request) -> Optional[Dict[str, Any]]:
         if not request.query_params.get("code"):
             raise OAuth2LoginError(400, "'code' parameter was not found in callback request")
-        if self.state != request.query_params.get("state"):
-            raise OAuth2LoginError(400, "'state' parameter does not match")
+        if not request.query_params.get("state"):
+            raise OAuth2LoginError(400, "'state' parameter was not found in callback request")
 
         url = request.url
-        scheme = "http" if self.allow_http else "https"
+        scheme = "http" if request.auth.http else "https"
         current_url = re.sub(r"^https?", scheme, str(url))
         redirect_uri = self.get_redirect_uri(request)
 
@@ -105,36 +98,30 @@ async def get_token_data(
             redirect_url=redirect_uri,
             authorization_response=current_url,
             code=request.query_params.get("code"),
-            **params,
+            state=request.query_params.get("state"),
         )
 
-        headers.update(additional_headers)
+        headers.update({
+            "Accept": "application/json",
+            "Content-Type": "application/x-www-form-urlencoded",
+        })
         auth = httpx.BasicAuth(self.client_id, self.client_secret)
         async with httpx.AsyncClient() as session:
             response = await session.post(token_url, headers=headers, content=content, auth=auth)
-            self.oauth_client.parse_request_body_response(json.dumps(response.json()))
-
-            url, headers, _ = self.oauth_client.add_token(self.userinfo_endpoint)
-            response = await session.get(url, headers=headers)
-            content = response.json()
-
-        return {**content, "scope": self.scope}
-
-    async def token_redirect(
-            self,
-            request: Request,
-            *,
-            params: Optional[Dict[str, Any]] = None,
-            headers: Optional[Dict[str, Any]] = None,
-    ) -> RedirectResponse:
-        token_data = await self.get_token_data(request, params=params, headers=headers)
+            token = self.oauth_client.parse_request_body_response(json.dumps(response.json()))
+            data = self.backend.user_data(token.get("access_token"))
+
+        return {**data, "scope": self.scope}
+
+    async def token_redirect(self, request: Request) -> RedirectResponse:
+        token_data = await self.get_token_data(request)
         access_token = request.auth.jwt_create(token_data)
         response = RedirectResponse(request.base_url)
         response.set_cookie(
             "Authorization",
             value=f"Bearer {access_token}",
-            httponly=self.allow_http,
             max_age=request.auth.expires,
             expires=request.auth.expires,
+            httponly=request.auth.http,
         )
         return response

src/fastapi_oauth2/middleware.py

@@ -23,6 +23,7 @@
 
 
 class Auth:
+    http: bool
     secret: str
     expires: int
     algorithm: str
@@ -32,6 +33,10 @@ class Auth:
     def __init__(self, scopes: Optional[List[str]] = None) -> None:
         self.scopes = scopes or []
 
+    @classmethod
+    def set_http(cls, http: bool) -> None:
+        cls.http = http
+
     @classmethod
     def set_secret(cls, secret: str) -> None:
         cls.secret = secret
@@ -72,10 +77,10 @@ def __init__(self, seq: Optional[dict] = None, **kwargs) -> None:
 
 class OAuth2Backend(AuthenticationBackend):
     def __init__(self, config: OAuth2Config) -> None:
+        Auth.set_http(config.allow_http)
         Auth.set_secret(config.jwt_secret)
         Auth.set_expires(config.jwt_expires)
         Auth.set_algorithm(config.jwt_algorithm)
-        OAuth2Core.allow_http = config.allow_http
         for client in config.clients:
             Auth.register_client(client)