1 files changed, 111 insertions, 66 deletions

diff --git a/authentication/endpoints.py b/authentication/endpoints.py
index 3af49fc..17d49d7 100644
--- a/authentication/endpoints.py
+++ b/authentication/endpoints.py
@@ -4,8 +4,8 @@ from functools import partial
 from typing import Annotated, Self
 from uuid import uuid7
 
-from fastapi import FastAPI, Request, Form, Query
-from pydantic import BaseModel, ValidationError, model_validator, EmailStr
+from fastapi import FastAPI, Form, Query, Request
+from pydantic import BaseModel, EmailStr, ValidationError, model_validator
 from sqlalchemy import func
 from sqlmodel import select
 from starlette.authentication import requires
@@ -15,10 +15,14 @@ from starlette.responses import HTMLResponse, RedirectResponse
 from starlette.templating import Jinja2Templates
 
 import conf
-from authentication.authlib import AuthorizationServer, AuthorizationCodeGrant, prepare_oauth_request
+from authentication.authlib import (
+    AuthorizationCodeGrant,
+    AuthorizationServer,
+    prepare_oauth_request,
+)
 from authentication.backend import SessionAuthBackend
 from authentication.utils import bcrypt_sha256_verify, get_url_path_with_query
-from db.models import User, Session, ClientApplication
+from db.models import ClientApplication, Session, User
 from db.session import get_session_context
 
 authentication = FastAPI()
@@ -31,10 +35,7 @@ oauth.add_middleware(SessionMiddleware, secret_key=conf.SECRET_KEY)
 auth_server = AuthorizationServer()
 code_grant = partial(AuthorizationCodeGrant, server=auth_server)
 
-templates = Jinja2Templates(directory=[
-    'templates',
-    'authentication/templates'
-])
+templates = Jinja2Templates(directory=["templates", "authentication/templates"])
 
 
 class LoginForm(BaseModel):
@@ -42,57 +43,75 @@ class LoginForm(BaseModel):
     password: str
 
 
-async def base_login(request: Request, form: LoginForm | None = None, errors: dict[str, str] | None = None):
-    return templates.TemplateResponse(request, 'login.html', context={
-        'form': form,
-        'errors': errors,
-    })
+async def base_login(
+    request: Request,
+    form: LoginForm | None = None,
+    errors: dict[str, str] | None = None,
+):
+    return templates.TemplateResponse(
+        request,
+        "login.html",
+        context={
+            "form": form,
+            "errors": errors,
+        },
+    )
 
 
-@authentication.get('/login/', response_class=HTMLResponse)
-async def get_login(request: Request, next: Annotated[str, Query()] = None):
+@authentication.get("/login/", response_class=HTMLResponse)
+async def get_login(request: Request, next: Annotated[str | None, Query()] = None):
     return await base_login(request)
 
 
-@authentication.post('/login/', response_class=HTMLResponse)
-async def post_login(request: Request, username: Annotated[str, Form()] = '', password: Annotated[str, Form()] = '',
-                     next: Annotated[str, Query()] = None):
+@authentication.post("/login/", response_class=HTMLResponse)
+async def post_login(
+    request: Request,
+    username: Annotated[str, Form()] = "",
+    password: Annotated[str, Form()] = "",
+    next: Annotated[str | None, Query()] = None,
+):
     form = None
     errors = {}
 
     try:
-        form = LoginForm.model_validate({
-            'username': username,
-            'password': password
-        })
+        form = LoginForm.model_validate({"username": username, "password": password})
 
         with get_session_context() as db_session:
             user = db_session.query(User).filter(User.username == form.username).first()
 
-        if not bcrypt_sha256_verify(password, user.password if user is not None else '') or user is None:
-            errors.update({'password': 'Invalid username or password'})
+        if (
+            not bcrypt_sha256_verify(
+                password, user.password if user is not None else ""
+            )
+            or user is None
+        ):
+            errors.update({"password": "Invalid username or password"})
         else:
             with get_session_context() as db_session:
                 session = Session(id=uuid7(), user=user)
                 db_session.add(session)
 
-            request.session.update({'id': str(session.id)})
+            request.session.update({"id": str(session.id)})
 
             if next is not None:
                 return RedirectResponse(next, status_code=303)
     except ValidationError as exc:
-        errors.update({e['loc'][0]: e['msg'] for e in exc.errors()})
+        errors.update({e["loc"][0]: e["msg"] for e in exc.errors()})
 
     return await base_login(request, form, errors)
 
 
-@authentication.get('/logout/')
+@authentication.get("/logout/")
 def logout(request: Request):
     with get_session_context() as db_session:
-        db_session.query(Session).filter(Session.id == request.session.get('id')).delete()
+        db_session.query(Session).filter(
+            Session.id == request.session.get("id")
+        ).delete()
 
     request.session.clear()
-    return RedirectResponse(request.headers.get('Referer', '/management/dashboard/'), status_code=303)
+    return RedirectResponse(
+        request.headers.get("Referer", "/management/dashboard/"), status_code=303
+    )
 
 
 class RegisterForm(BaseModel):
@@ -100,10 +119,10 @@ class RegisterForm(BaseModel):
     password: str
     verify_password: str
 
-    @model_validator(mode='after')
+    @model_validator(mode="after")
     def check_passwords_match(self) -> Self:
         if self.password != self.verify_password:
-            raise ValueError('Passwords do not match')
+            raise ValueError("Passwords do not match")
         return self
 
 
@@ -113,31 +132,42 @@ def is_first_run() -> bool:
         return user_count.first() == 0
 
 
-async def base_register(request: Request, form: RegisterForm | None = None, errors: dict[str, str] | None = None):
-    return templates.TemplateResponse(request, 'register.html', context={
-        'form': form,
-        'errors': errors,
-        'is_first_run': is_first_run(),
-    })
+async def base_register(
+    request: Request,
+    form: RegisterForm | None = None,
+    errors: dict[str, str] | None = None,
+):
+    return templates.TemplateResponse(
+        request,
+        "register.html",
+        context={
+            "form": form,
+            "errors": errors,
+            "is_first_run": is_first_run(),
+        },
+    )
 
 
-@authentication.get('/register/', response_class=HTMLResponse)
-async def get_register(request: Request, next: Annotated[str, Query()] = None):
+@authentication.get("/register/", response_class=HTMLResponse)
+async def get_register(request: Request, next: Annotated[str | None, Query()] = None):
     return await base_register(request)
 
 
-@authentication.post('/register/', response_class=HTMLResponse)
-async def post_register(request: Request, email: Annotated[str, Form()] = '', password: Annotated[str, Form()] = '',
-                        verify_password: Annotated[str, Form()] = '', next: Annotated[str, Query()] = None):
+@authentication.post("/register/", response_class=HTMLResponse)
+async def post_register(
+    request: Request,
+    email: Annotated[str, Form()] = "",
+    password: Annotated[str, Form()] = "",
+    verify_password: Annotated[str, Form()] = "",
+    next: Annotated[str | None, Query()] = None,
+):
     form = None
     errors = {}
 
     try:
-        form = RegisterForm.model_validate({
-            'email': email,
-            'password': password,
-            'verify_password': verify_password
-        })
+        form = RegisterForm.model_validate(
+            {"email": email, "password": password, "verify_password": verify_password}
+        )
 
         with get_session_context() as db_session:
             user = User(username=email, email=email)
@@ -147,13 +177,18 @@ async def post_register(request: Request, email: Annotated[str, Form()] = '', pa
             session = Session(id=uuid7(), user=user)
             db_session.add(session)
 
-        request.session.update({'id': str(session.id)})
+        request.session.update({"id": str(session.id)})
 
         if next is not None:
             return RedirectResponse(next, status_code=303)
 
     except ValidationError as exc:
-        errors.update({(e['loc'][0] if e['loc'] else 'verify_password'): e['msg'] for e in exc.errors()})
+        errors.update(
+            {
+                (e["loc"][0] if e["loc"] else "verify_password"): e["msg"]
+                for e in exc.errors()
+            }
+        )
 
     return await base_register(request, form, errors)
 
@@ -162,49 +197,59 @@ async def base_authorize(request: Request, user: User | None = None):
     return auth_server.create_authorization_response(request, request.user, grant=None)
 
 
-@oauth.get('/authorize/')
-@requires(['authenticated', 'session'], redirect='get_login')
+@oauth.get("/authorize/")
+@requires(["authenticated", "session"], redirect="get_login")
 async def get_authorize(request: Request):
     return await base_authorize(request, request.user)
 
 
-@oauth.post('/authorize/')
-@requires(['authenticated', 'session'], redirect='get_login')
+@oauth.post("/authorize/")
+@requires(["authenticated", "session"], redirect="get_login")
 async def post_authorize(request: Request):
     return await base_authorize(request, request.user)
 
 
-@oauth.post('/token/')
+@oauth.post("/token/")
 async def post_token(request: Request):
     await prepare_oauth_request(request)
     return auth_server.create_token_response(request)
 
 
-@authentication.get('/connect/')
+@authentication.get("/connect/")
 async def connect(request: Request):
-    origin = request.url.scheme + '://' + request.url.netloc
+    origin = request.url.scheme + "://" + request.url.netloc
 
     with get_session_context() as session:
         app = session.exec(select(ClientApplication)).first()
 
     return RedirectResponse(
-        '/oauth' + get_url_path_with_query(oauth, 'get_authorize', {
-            'client_id': app.client_id,
-            'redirect_uri': f'{origin}/oauth' + get_url_path_with_query(oauth, 'callback'),
-            'response_type': 'code',
-            'state': b64encode(json.dumps({
-                'client_id': app.client_id,
-            }).encode()).decode(),
-        }),
+        "/oauth"
+        + get_url_path_with_query(
+            oauth,
+            "get_authorize",
+            {
+                "client_id": app.client_id,
+                "redirect_uri": f"{origin}/oauth"
+                + get_url_path_with_query(oauth, "callback"),
+                "response_type": "code",
+                "state": b64encode(
+                    json.dumps(
+                        {
+                            "client_id": app.client_id,
+                        }
+                    ).encode()
+                ).decode(),
+            },
+        ),
         status_code=303,
     )
 
 
-@oauth.get('/config/')
+@oauth.get("/config/")
 async def get_config(request: Request):
     with get_session_context() as session:
         app = session.exec(select(ClientApplication)).first()
 
     return {
-        'client_id': app.client_id,
+        "client_id": app.client_id,
     }