Make OAuth2Core replace the GitHubOAuth2

Author: ArtyomVancyan

src/fastapi_oauth2/core.py

@@ -1,16 +1,24 @@
 import json
-import os
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from urllib.parse import urljoin
 
 import httpx
 from oauthlib.oauth2 import WebApplicationClient
+from starlette.exceptions import HTTPException
 from starlette.requests import Request
 from starlette.responses import RedirectResponse
 
-from .config import JWT_EXPIRES
-from .exceptions import OAuth2LoginError
-from .utils import jwt_create
+from .client import OAuth2Client
+
+
+class OAuth2LoginError(HTTPException):
+    """Raised when any login-related error occurs
+    (such as when user is not verified or if there was an attempt for fake login)
+    """
 
 
 class OAuth2Core:
@@ -19,7 +27,7 @@ class OAuth2Core:
     client_id: str = None
     client_secret: str = None
     callback_url: Optional[str] = None
-    allow_insecure_http: bool = False
+    allow_http: bool = False
     scope: Optional[List[str]] = None
     state: Optional[str] = None
     _oauth_client: Optional[WebApplicationClient] = None
@@ -29,55 +37,47 @@ class OAuth2Core:
     token_endpoint: str = None
     userinfo_endpoint: str = None
 
-    def __init__(
-            self,
-            client_id: str,
-            client_secret: str,
-            callback_url: Optional[str] = None,
-            allow_insecure_http: bool = False,
-            scope: Optional[List[str]] = None,
-    ):
-        self.client_id = client_id
-        self.client_secret = client_secret
-        self.callback_url = callback_url
-        self.allow_insecure_http = allow_insecure_http
-        if allow_insecure_http:
-            os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1"
-        self.scope = scope or self.scope
+    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.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:
         if self._oauth_client is None:
             self._oauth_client = WebApplicationClient(self.client_id)
         return self._oauth_client
 
-    @property
-    def access_token(self) -> Optional[str]:
-        return self.oauth_client.access_token
-
-    @property
-    def refresh_token(self) -> Optional[str]:
-        return self.oauth_client.refresh_token
+    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 {}
+        redirect_uri = self.get_redirect_uri(request)
         return self.oauth_client.prepare_request_uri(
-            self.authorization_endpoint, redirect_uri=self.callback_url, state=state, scope=self.scope, **params
+            self.authorization_endpoint, redirect_uri=redirect_uri, state=state, scope=self.scope, **params
         )
 
     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(params=params, state=state)
+        login_uri = await self.get_login_url(request, params=params, state=state)
         return RedirectResponse(login_uri, 303)
 
     async def get_token_data(
@@ -96,15 +96,14 @@ async def get_token_data(
             raise OAuth2LoginError(400, "'state' parameter does not match")
 
         url = request.url
-        scheme = "http" if self.allow_insecure_http else "https"
-        current_path = f"{scheme}://{url.netloc}{url.path}"
-        current_path = re.sub(r"^https?", scheme, current_path)
+        scheme = "http" if self.allow_http else "https"
         current_url = re.sub(r"^https?", scheme, str(url))
+        redirect_uri = self.get_redirect_uri(request)
 
         token_url, headers, content = self.oauth_client.prepare_token_request(
             self.token_endpoint,
+            redirect_url=redirect_uri,
             authorization_response=current_url,
-            redirect_url=self.callback_url or current_path,
             code=request.query_params.get("code"),
             **params,
         )
@@ -129,13 +128,13 @@ async def token_redirect(
             headers: Optional[Dict[str, Any]] = None,
     ) -> RedirectResponse:
         token_data = await self.get_token_data(request, params=params, headers=headers)
-        access_token = jwt_create(token_data)
+        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_insecure_http,
-            max_age=JWT_EXPIRES * 60,
-            expires=JWT_EXPIRES * 60,
+            httponly=self.allow_http,
+            max_age=request.auth.expires,
+            expires=request.auth.expires,
         )
         return response