1 files changed, 64 insertions, 24 deletions

diff --git a/authentication/endpoints.py b/authentication/endpoints.py
index ebdc7a1..f94d55b 100644
--- a/authentication/endpoints.py
+++ b/authentication/endpoints.py
@@ -1,31 +1,40 @@
+import json
+from base64 import b64encode
+from functools import partial
 from typing import Annotated
 from uuid import uuid7
 
-from fastapi import FastAPI, Request, Form, Query, HTTPException
-from pydantic import BaseModel, EmailStr, ValidationError
+from alembic.testing import requirements
+from fastapi import FastAPI, Request, Form, Query
+from pydantic import BaseModel, ValidationError
+from sqlmodel import select
+from starlette.authentication import requires
 from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.sessions import SessionMiddleware
 from starlette.responses import HTMLResponse, RedirectResponse
 from starlette.templating import Jinja2Templates
 
 import conf
-from authentication.backend import AuthBackend
-from authentication.utils import bcrypt_sha256_verify
+from authentication.authlib import AuthorizationServer, AuthorizationCodeGrant, 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.session import get_session_context
-from proxy.utils import get_path_with_query_string
 
 authentication = FastAPI()
 authentication.add_middleware(SessionMiddleware, secret_key=conf.SECRET_KEY)
 
 oauth = FastAPI()
-oauth.add_middleware(AuthenticationMiddleware, backend=AuthBackend())
+oauth.add_middleware(AuthenticationMiddleware, backend=SessionAuthBackend())
 oauth.add_middleware(SessionMiddleware, secret_key=conf.SECRET_KEY)
 
+auth_server = AuthorizationServer()
+code_grant = partial(AuthorizationCodeGrant, server=auth_server)
+
 templates = Jinja2Templates(directory='authentication/templates')
 
 class LoginForm(BaseModel):
-    username: EmailStr
+    username: str
     password: str
 
 async def base_login(request: Request, form: LoginForm | None = None, errors: dict[str, str] | None = None):
@@ -35,11 +44,11 @@ async def base_login(request: Request, form: LoginForm | None = None, errors: di
     })
 
 @authentication.get('/login/', response_class=HTMLResponse)
-async def login(request: Request):
+async def get_login(request: Request, next: Annotated[str, Query()] = None):
     return await base_login(request)
 
 @authentication.post('/login/', response_class=HTMLResponse)
-async def login(request: Request, username: Annotated[str, Form()] = '', password: Annotated[str, Form()] = '', redirect: Annotated[str, Query()] = None):
+async def post_login(request: Request, username: Annotated[str, Form()] = '', password: Annotated[str, Form()] = '', next: Annotated[str, Query()] = None):
     form = None
     errors = {}
 
@@ -58,30 +67,61 @@ async def login(request: Request, username: Annotated[str, Form()] = '', passwor
             session = Session(id=uuid7(), user=user)
             db_session.add(session)
 
-        if redirect is not None:
-            return RedirectResponse(redirect)
+        request.session.update({'id': str(session.id)})
+
+        if next is not None:
+            return RedirectResponse(next)
 
     except ValidationError as exc:
         errors.update({ e['loc'][0]: e['msg'] for e in exc.errors() })
 
     return await base_login(request, form, errors)
 
+async def base_authorize(request: Request, user: User | None = None):
+    if not conf.IS_LOCAL:
+        with get_session_context() as db_session:
+            db_session.query(Session).filter(Session.id == request.session.get('id')).delete()
 
-async def base_authorize(request: Request, client_id: str, redirect_uri: str):
-    with get_session_context() as session:
-        app = session.query(ClientApplication).where(ClientApplication.client_id == client_id).first()
+    return auth_server.create_authorization_response(request, request.user, grant=None)
 
-    if app is None:
-        raise HTTPException(status_code=404, detail='Invalid client id')
+@oauth.get('/authorize/')
+@requires(['authenticated', 'session'], redirect='get_login')
+async def get_authorize(request: Request):
+    return await base_authorize(request, request.user)
 
-    if not request.user.is_authenticated:
-        return RedirectResponse(f'/auth/login/?redirect={get_path_with_query_string(request)}')
+@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/')
+async def post_token(request: Request):
+    await prepare_oauth_request(request)
+    return auth_server.create_token_response(request)
 
-@oauth.get('/authorize/')
-async def authorize(request: Request, client_id: Annotated[str, Query()], redirect_uri: Annotated[str, Query()]):
-    return await base_authorize(request, client_id, redirect_uri)
+@authentication.get('/connect/')
+async def connect(request: Request):
+    origin = request.url.scheme + '://' + request.url.netloc
 
-@oauth.post('/authorize/')
-async def authorize(request: Request, client_id: Annotated[str, Query()], redirect_uri: Annotated[str, Query()]):
-    return await base_authorize(request, client_id, redirect_uri)
+    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('/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,
+    }