1 files changed, 52 insertions, 25 deletions

diff --git a/authentication/authlib.py b/authentication/authlib.py
index 8472062..a37d68f 100644
--- a/authentication/authlib.py
+++ b/authentication/authlib.py
@@ -1,31 +1,37 @@
 import json
 from collections import defaultdict
+from typing import ClassVar
 from unittest.mock import patch
 
 from authlib.common.security import generate_token as generate_random_token
-from authlib.oauth2 import AuthorizationServer as BaseAuthorizationServer, OAuth2Request
-from authlib.oauth2.rfc6749 import AuthorizationCodeGrant as BaseAuthorizationCodeGrant, \
-    RefreshTokenGrant as BaseRefreshTokenGrant, OAuth2Payload
+from authlib.oauth2 import AuthorizationServer as BaseAuthorizationServer
+from authlib.oauth2 import OAuth2Request
+from authlib.oauth2.rfc6749 import AuthorizationCodeGrant as BaseAuthorizationCodeGrant
+from authlib.oauth2.rfc6749 import OAuth2Payload
+from authlib.oauth2.rfc6749 import RefreshTokenGrant as BaseRefreshTokenGrant
 from authlib.oauth2.rfc6750 import BearerTokenGenerator
 from authlib.oauth2.rfc7636 import CodeChallenge
+from fastapi import Request, Response
 from sqlalchemy.exc import NoResultFound
-
 from sqlmodel import select
-from fastapi import Request, Response
 
-from db.models import ClientApplication, AuthToken, AuthCode, User
+from db.models import AuthCode, AuthToken, ClientApplication, User
 from db.session import get_session_context
 
 
 class AuthorizationCodeGrant(BaseAuthorizationCodeGrant):
-    TOKEN_ENDPOINT_AUTH_METHODS = ['client_secret_basic', 'client_secret_post', 'none']
+    TOKEN_ENDPOINT_AUTH_METHODS: ClassVar[list[str]] = [
+        "client_secret_basic",
+        "client_secret_post",
+        "none",
+    ]
 
     def save_authorization_code(self, code: str, request: OAuth2Request):
         with get_session_context() as db_session:
             db_session.add(request.user)
 
-            code_challenge = request.payload.data.get('code_challenge')
-            code_challenge_method = request.payload.data.get('code_challenge_method')
+            code_challenge = request.payload.data.get("code_challenge")
+            code_challenge_method = request.payload.data.get("code_challenge_method")
 
             auth_code = AuthCode(
                 code=code,
@@ -41,10 +47,16 @@ class AuthorizationCodeGrant(BaseAuthorizationCodeGrant):
 
         return auth_code
 
-    def query_authorization_code(self, code: str, client: ClientApplication) -> AuthCode | None:
+    def query_authorization_code(
+        self, code: str, client: ClientApplication
+    ) -> AuthCode | None:
         try:
             with get_session_context() as db_session:
-                auth_code = db_session.exec(select(AuthCode).where(AuthCode.code == code, AuthCode.client_id == client.client_id)).one()
+                auth_code = db_session.exec(
+                    select(AuthCode).where(
+                        AuthCode.code == code, AuthCode.client_id == client.client_id
+                    )
+                ).one()
         except NoResultFound:
             return None
 
@@ -65,7 +77,9 @@ class RefreshTokenGrant(BaseRefreshTokenGrant):
     def authenticate_refresh_token(self, refresh_token) -> AuthToken | None:
         try:
             with get_session_context() as db_session:
-                token = db_session.exec(select(AuthToken).where(AuthToken.refresh_token == refresh_token)).one()
+                token = db_session.exec(
+                    select(AuthToken).where(AuthToken.refresh_token == refresh_token)
+                ).one()
         except NoResultFound:
             return None
 
@@ -82,6 +96,7 @@ class RefreshTokenGrant(BaseRefreshTokenGrant):
             credential.revoked = True
             db_session.add(credential)
 
+
 async def prepare_oauth_request(request: Request) -> Request:
     async with request.form() as form:
         request.state.form_data = dict(form)
@@ -96,7 +111,7 @@ class FastApiOAuth2Payload(OAuth2Payload):
     def data(self):
         return {
             **self._request.query_params,
-            **getattr(self._request.state, 'form_data', {}),
+            **getattr(self._request.state, "form_data", {}),
         }
 
     @property
@@ -106,9 +121,10 @@ class FastApiOAuth2Payload(OAuth2Payload):
             values[k].extend([self.data[k]])
         return values
 
+
 class FastApiOAuth2Request(OAuth2Request):
     def __init__(self, request: Request):
-        with patch('authlib.oauth2.rfc6749.errors.InsecureTransportError'):
+        with patch("authlib.oauth2.rfc6749.errors.InsecureTransportError"):
             super().__init__(
                 method=request.method,
                 uri=str(request.url),
@@ -128,7 +144,7 @@ class FastApiOAuth2Request(OAuth2Request):
 
     @property
     def form(self):
-        return getattr(self._request.state, 'form_data', {})
+        return getattr(self._request.state, "form_data", {})
 
 
 class AuthorizationServer(BaseAuthorizationServer):
@@ -136,22 +152,33 @@ class AuthorizationServer(BaseAuthorizationServer):
         super().__init__(*args, **kwargs)
         self.register_grant(AuthorizationCodeGrant, extensions=[CodeChallenge()])
         self.register_grant(RefreshTokenGrant)
-        self.register_token_generator("default", BearerTokenGenerator(
-            access_token_generator=lambda *args, **kwargs: generate_random_token(42),
-            refresh_token_generator=lambda *args, **kwargs: generate_random_token(48),
-        ))
+        self.register_token_generator(
+            "default",
+            BearerTokenGenerator(
+                access_token_generator=lambda *args, **kwargs: generate_random_token(
+                    42
+                ),
+                refresh_token_generator=lambda *args, **kwargs: generate_random_token(
+                    48
+                ),
+            ),
+        )
 
     def query_client(self, client_id: str) -> ClientApplication:
         with get_session_context() as db_session:
-            return db_session.exec(select(ClientApplication).where(ClientApplication.client_id == client_id)).one()
+            return db_session.exec(
+                select(ClientApplication).where(
+                    ClientApplication.client_id == client_id
+                )
+            ).one()
 
     def save_token(self, token: dict, request: OAuth2Request):
         with get_session_context() as db_session:
-            db_session.add(AuthToken(
-                user_id=request.user.id,
-                client_id=request.client.client_id,
-                **token
-            ))
+            db_session.add(
+                AuthToken(
+                    user_id=request.user.id, client_id=request.client.client_id, **token
+                )
+            )
 
     def create_oauth2_request(self, request: Request) -> OAuth2Request:
         return FastApiOAuth2Request(request)