Create the core middleware and integrate

Author: ArtyomVancyan

main.py

@@ -2,13 +2,10 @@
 
 from fastapi import FastAPI, Request, APIRouter
 from fastapi.responses import HTMLResponse
-from fastapi.security.utils import get_authorization_scheme_param
 from fastapi.templating import Jinja2Templates
-from starlette.authentication import AuthenticationBackend
-from starlette.middleware.authentication import AuthenticationMiddleware
 
-from demo.dependencies import get_current_user
 from demo.router import router as demo_router
+from fastapi_oauth2.middleware import OAuth2Middleware
 from fastapi_oauth2.router import router as oauth2_router
 
 router = APIRouter()
@@ -24,19 +21,13 @@ async def root(request: Request):
 app.include_router(router)
 app.include_router(demo_router)
 app.include_router(oauth2_router)
-
-
-class BearerTokenAuthBackend(AuthenticationBackend):
-    async def authenticate(self, request):
-        authorization = request.cookies.get("Authorization")
-        scheme, param = get_authorization_scheme_param(authorization)
-
-        if not scheme or not param:
-            return "", None
-
-        return authorization, await get_current_user(param)
-
-
-@app.on_event('startup')
-async def startup():
-    app.add_middleware(AuthenticationMiddleware, backend=BearerTokenAuthBackend())
+app.add_middleware(OAuth2Middleware, config={
+    "allow_http": True,
+    "providers": {
+        "github": {
+            "client_id": "eccd08d6736b7999a32a",
+            "client_secret": "642999c1c5f2b3df8b877afdc78252ef5b594d31",
+            "redirect_uri": "http://127.0.0.1:8000/",
+        },
+    }
+})

src/fastapi_oauth2/base.py

@@ -10,7 +10,7 @@
 from starlette.responses import RedirectResponse
 
 from .config import JWT_EXPIRES, OAUTH2_REDIRECT_URL
-from .utils import create_access_token
+from .utils import jwt_create
 
 
 class OAuth2LoginError(HTTPException):
@@ -68,27 +68,22 @@ def refresh_token(self) -> Optional[str]:
     async def get_login_url(
             self,
             *,
-            redirect_uri: Optional[str] = None,
             params: Optional[Dict[str, Any]] = None,
             state: Optional[str] = None,
     ) -> Any:
         self.state = state
         params = params or {}
-        redirect_uri = redirect_uri or self.callback_url
-        if redirect_uri is None:
-            raise ValueError("callback_url must be provided, either at construction or request time")
         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=self.callback_url, state=state, scope=self.scope, **params
         )
 
     async def login_redirect(
             self,
             *,
-            redirect_uri: Optional[str] = None,
             params: Optional[Dict[str, Any]] = None,
             state: Optional[str] = None,
     ) -> RedirectResponse:
-        login_uri = await self.get_login_url(redirect_uri=redirect_uri, params=params, state=state)
+        login_uri = await self.get_login_url(params=params, state=state)
         return RedirectResponse(login_uri, 303)
 
     async def get_token_data(
@@ -97,7 +92,6 @@ async def get_token_data(
             *,
             params: Optional[Dict[str, Any]] = None,
             headers: Optional[Dict[str, Any]] = None,
-            redirect_uri: Optional[str] = None,
     ) -> Optional[Dict[str, Any]]:
         params = params or {}
         additional_headers = headers or {}
@@ -116,7 +110,7 @@ async def get_token_data(
         token_url, headers, content = self.oauth_client.prepare_token_request(
             self.token_endpoint,
             authorization_response=current_url,
-            redirect_url=redirect_uri or self.callback_url or current_path,
+            redirect_url=self.callback_url or current_path,
             code=request.query_params.get("code"),
             **params,
         )
@@ -131,18 +125,17 @@ async def get_token_data(
             response = await session.get(url, headers=headers)
             content = response.json()
 
-        return content
+        return {**content, "scope": self.scope}
 
     async def token_redirect(
             self,
             request: Request,
             *,
             params: Optional[Dict[str, Any]] = None,
             headers: Optional[Dict[str, Any]] = None,
-            redirect_uri: Optional[str] = None,
     ) -> RedirectResponse:
-        token_data = await self.get_token_data(request, params=params, headers=headers, redirect_uri=redirect_uri)
-        access_token = create_access_token(token_data)
+        token_data = await self.get_token_data(request, params=params, headers=headers)
+        access_token = jwt_create(token_data)
         response = RedirectResponse(OAUTH2_REDIRECT_URL)
         response.set_cookie(
             "Authorization",

src/fastapi_oauth2/middleware.py

@@ -0,0 +1,38 @@
+from typing import Optional, Tuple, Union
+
+from fastapi.security.utils import get_authorization_scheme_param
+from starlette.authentication import AuthenticationBackend, AuthCredentials
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.requests import Request
+from starlette.types import Send, Receive, Scope, ASGIApp
+
+from .types import Config
+from .types import ConfigParams
+from .utils import jwt_decode
+
+
+class OAuth2Backend(AuthenticationBackend):
+    async def authenticate(self, request: Request) -> Optional[Tuple["AuthCredentials", Optional[dict]]]:
+        authorization = request.cookies.get("Authorization")
+        scheme, param = get_authorization_scheme_param(authorization)
+
+        if not scheme or not param:
+            return AuthCredentials(), None
+
+        access_token = jwt_decode(param)
+        scope = access_token.pop("scope")
+        return AuthCredentials(scope), access_token
+
+
+class OAuth2Middleware:
+    def __init__(self, app: ASGIApp, config: Union[Config, ConfigParams]) -> None:
+        if isinstance(config, Config):
+            self.config = config
+        elif isinstance(config, dict):
+            self.config = Config(**config)
+        else:
+            raise ValueError("config does not contain valid parameters")
+        self.auth_middleware = AuthenticationMiddleware(app, OAuth2Backend())
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        await self.auth_middleware(scope, receive, send)

src/fastapi_oauth2/types.py

@@ -0,0 +1,32 @@
+from enum import Enum
+from typing import Dict, TypedDict
+
+
+class OAuth2Provider(str, Enum):
+    github = "github"
+
+
+class OAuth2Client(Dict[str, str]):
+    client_id: str
+    client_secret: str
+    redirect_uri: str
+
+
+class ConfigParams(TypedDict):
+    allow_http: bool
+    jwt_secret: str
+    jwt_expires: int
+    jwt_algorithm: str
+    providers: Dict[OAuth2Provider, OAuth2Client]
+
+
+class Config:
+    allow_http: bool = False
+    jwt_secret: str = ""
+    jwt_expires: int = 900
+    jwt_algorithm: str = "HS256"
+    providers: Dict[OAuth2Provider, OAuth2Client] = {}
+
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)

src/fastapi_oauth2/utils.py

@@ -13,6 +13,6 @@ def jwt_decode(token: str) -> dict:
     return jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
 
 
-def create_access_token(token_data: dict) -> str:
+def jwt_create(token_data: dict) -> str:
     expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRES)
     return jwt_encode({**token_data, "exp": expire})