Refactor env vars and cookie setting

Author: ArtyomVancyan

.env

@@ -1,7 +1,7 @@
-CLIENT_ID=eccd08d6736b7999a32a
-CLIENT_SECRET=642999c1c5f2b3df8b877afdc78252ef5b594d31
-CALLBACK_URL=http://127.0.0.1:8000/auth/callback
-REDIRECT_URL=http://127.0.0.1:8000/
+OAUTH2_CLIENT_ID=eccd08d6736b7999a32a
+OAUTH2_CLIENT_SECRET=642999c1c5f2b3df8b877afdc78252ef5b594d31
+OAUTH2_CALLBACK_URL=http://127.0.0.1:8000/auth/callback
+OAUTH2_REDIRECT_URL=http://127.0.0.1:8000/
 
 JWT_SECRET=secret
 JWT_ALGORITHM=HS256

demo/dependencies.py

@@ -4,11 +4,11 @@
 from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
 from fastapi.security import OAuth2
 from fastapi.security.utils import get_authorization_scheme_param
-from jose import jwt, JWTError
+from jose import JWTError
 from starlette.requests import Request
 from starlette.status import HTTP_403_FORBIDDEN
 
-from fastapi_oauth2.config import JWT_SECRET, JWT_ALGORITHM
+from fastapi_oauth2.utils import jwt_decode
 
 
 class OAuth2PasswordBearerCookie(OAuth2):
@@ -44,7 +44,7 @@ async def __call__(self, request: Request) -> Optional[str]:
 
 async def get_current_user(token: str = Depends(oauth2_scheme)):
     try:
-        return jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        return jwt_decode(token)
     except JWTError:
         raise HTTPException(
             status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials"

fastapi_oauth2/base.py

@@ -9,14 +9,17 @@
 from starlette.requests import Request
 from starlette.responses import RedirectResponse
 
+from .config import JWT_EXPIRES, OAUTH2_REDIRECT_URL
+from .utils import create_access_token
 
-class SSOLoginError(HTTPException):
+
+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 SSOBase:
+class OAuth2Base:
     """Base class (mixin) for all SSO providers"""
 
     client_id: str = None
@@ -78,7 +81,7 @@ async def get_login_url(
             self.authorization_endpoint, redirect_uri=redirect_uri, state=state, scope=self.scope, **params
         )
 
-    async def get_login_redirect(
+    async def login_redirect(
             self,
             *,
             redirect_uri: Optional[str] = None,
@@ -88,7 +91,7 @@ async def get_login_redirect(
         login_uri = await self.get_login_url(redirect_uri=redirect_uri, params=params, state=state)
         return RedirectResponse(login_uri, 303)
 
-    async def verify_and_process(
+    async def get_token_data(
             self,
             request: Request,
             *,
@@ -100,9 +103,9 @@ async def verify_and_process(
         additional_headers = headers or {}
         additional_headers.update(self.additional_headers or {})
         if not request.query_params.get("code"):
-            raise SSOLoginError(400, "'code' parameter was not found in callback request")
+            raise OAuth2LoginError(400, "'code' parameter was not found in callback request")
         if self.state != request.query_params.get("state"):
-            raise SSOLoginError(400, "'state' parameter does not match")
+            raise OAuth2LoginError(400, "'state' parameter does not match")
 
         url = request.url
         scheme = "http" if self.allow_insecure_http else "https"
@@ -129,3 +132,23 @@ async def verify_and_process(
             content = response.json()
 
         return content
+
+    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)
+        response = RedirectResponse(OAUTH2_REDIRECT_URL)
+        response.set_cookie(
+            "Authorization",
+            value=f"Bearer {access_token}",
+            httponly=self.allow_insecure_http,
+            max_age=JWT_EXPIRES * 60,
+            expires=JWT_EXPIRES * 60,
+        )
+        return response

fastapi_oauth2/config.py

@@ -4,11 +4,11 @@
 
 load_dotenv()
 
-CLIENT_ID = os.getenv("CLIENT_ID")
-CLIENT_SECRET = os.getenv("CLIENT_SECRET")
-CALLBACK_URL = os.getenv("CALLBACK_URL")
-REDIRECT_URL = os.getenv("REDIRECT_URL")
+OAUTH2_CLIENT_ID = os.getenv("OAUTH2_CLIENT_ID")
+OAUTH2_CLIENT_SECRET = os.getenv("OAUTH2_CLIENT_SECRET")
+OAUTH2_CALLBACK_URL = os.getenv("OAUTH2_CALLBACK_URL")
+OAUTH2_REDIRECT_URL = os.getenv("OAUTH2_REDIRECT_URL")
 
 JWT_SECRET = os.getenv("JWT_SECRET")
 JWT_ALGORITHM = os.getenv("JWT_ALGORITHM")
-JWT_EXPIRES = int(os.getenv("JWT_EXPIRES"))
+JWT_EXPIRES = int(os.getenv("JWT_EXPIRES", "15"))

fastapi_oauth2/github.py

@@ -1,7 +1,7 @@
-from .base import SSOBase
+from .base import OAuth2Base
 
 
-class GitHubSSO(SSOBase):
+class GitHubOAuth2(OAuth2Base):
     """Class providing login via GitHub SSO"""
 
     scope = ["user:email"]

fastapi_oauth2/router.py

@@ -1,53 +1,36 @@
-from datetime import timedelta
-
 from fastapi import APIRouter
 from fastapi.responses import RedirectResponse
 from starlette.requests import Request
 
-from fastapi_oauth2.github import GitHubSSO
+from fastapi_oauth2.github import GitHubOAuth2
 from .config import (
-    CLIENT_ID,
-    CLIENT_SECRET,
-    CALLBACK_URL,
-    JWT_EXPIRES,
-    REDIRECT_URL,
+    OAUTH2_CLIENT_ID,
+    OAUTH2_CLIENT_SECRET,
+    OAUTH2_CALLBACK_URL,
+    OAUTH2_REDIRECT_URL,
 )
-from .utils import create_access_token
 
 router = APIRouter()
-sso = GitHubSSO(
-    client_id=CLIENT_ID,
-    client_secret=CLIENT_SECRET,
-    callback_url=CALLBACK_URL,
+oauth2 = GitHubOAuth2(
+    client_id=OAUTH2_CLIENT_ID,
+    client_secret=OAUTH2_CLIENT_SECRET,
+    callback_url=OAUTH2_CALLBACK_URL,
     allow_insecure_http=True,
 )
 
 
 @router.get("/auth/login")
 async def login():
-    return await sso.get_login_redirect()
+    return await oauth2.login_redirect()
 
 
 @router.get("/auth/callback")
 async def callback(request: Request):
-    user = await sso.verify_and_process(request)
-    expires_delta = timedelta(minutes=JWT_EXPIRES)
-    access_token = create_access_token(
-        data=dict(user), expires_delta=expires_delta
-    )
-    response = RedirectResponse(REDIRECT_URL)
-    response.set_cookie(
-        "Authorization",
-        value=f"Bearer {access_token}",
-        httponly=sso.allow_insecure_http,
-        max_age=JWT_EXPIRES * 60,
-        expires=JWT_EXPIRES * 60,
-    )
-    return response
+    return await oauth2.token_redirect(request)
 
 
 @router.get("/auth/logout")
 async def logout():
-    response = RedirectResponse(REDIRECT_URL)
+    response = RedirectResponse(OAUTH2_REDIRECT_URL)
     response.delete_cookie("Authorization")
     return response

fastapi_oauth2/utils.py

@@ -1,11 +1,18 @@
 from datetime import datetime, timedelta
-from typing import Optional
 
 from jose import jwt
 
-from .config import JWT_SECRET, JWT_ALGORITHM
+from .config import JWT_SECRET, JWT_ALGORITHM, JWT_EXPIRES
 
 
-def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
-    expire = datetime.utcnow() + expires_delta if expires_delta else timedelta(minutes=15)
-    return jwt.encode({**data, "exp": expire}, JWT_SECRET, algorithm=JWT_ALGORITHM)
+def jwt_encode(data: dict) -> str:
+    return jwt.encode(data, JWT_SECRET, algorithm=JWT_ALGORITHM)
+
+
+def jwt_decode(token: str) -> dict:
+    return jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+
+
+def create_access_token(token_data: dict) -> str:
+    expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRES)
+    return jwt_encode({**token_data, "exp": expire})