7 files changed, 327 insertions, 39 deletions

diff --git a/authentication/authlib.py b/authentication/authlib.py
new file mode 100644
index 0000000..8472062
--- /dev/null
+++ b/authentication/authlib.py
@@ -0,0 +1,167 @@
+import json
+from collections import defaultdict
+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.rfc6750 import BearerTokenGenerator
+from authlib.oauth2.rfc7636 import CodeChallenge
+from sqlalchemy.exc import NoResultFound
+
+from sqlmodel import select
+from fastapi import Request, Response
+
+from db.models import ClientApplication, AuthToken, AuthCode, User
+from db.session import get_session_context
+
+
+class AuthorizationCodeGrant(BaseAuthorizationCodeGrant):
+    TOKEN_ENDPOINT_AUTH_METHODS = ['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')
+
+            auth_code = AuthCode(
+                code=code,
+                client_id=request.client.client_id,
+                redirect_uri=request.payload.redirect_uri,
+                response_type=request.payload.response_type,
+                scope=request.payload.scope,
+                user_id=request.user.id,
+                code_challenge=code_challenge,
+                code_challenge_method=code_challenge_method,
+            )
+            db_session.add(auth_code)
+
+        return auth_code
+
+    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()
+        except NoResultFound:
+            return None
+
+        if auth_code.is_expired():
+            return None
+
+        return auth_code
+
+    def delete_authorization_code(self, code: AuthCode):
+        with get_session_context() as db_session:
+            db_session.delete(code)
+
+    def authenticate_user(self, authentication_code: AuthCode) -> User:
+        return authentication_code.user
+
+
+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()
+        except NoResultFound:
+            return None
+
+        if not token.is_refresh_token_active():
+            return None
+
+        return token
+
+    def authenticate_user(self, credential: AuthToken):
+        return credential.user
+
+    def revoke_old_credential(self, credential: AuthToken):
+        with get_session_context() as db_session:
+            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)
+    return request
+
+
+class FastApiOAuth2Payload(OAuth2Payload):
+    def __init__(self, request: Request):
+        self._request = request
+
+    @property
+    def data(self):
+        return {
+            **self._request.query_params,
+            **getattr(self._request.state, 'form_data', {}),
+        }
+
+    @property
+    def datalist(self):
+        values = defaultdict(list)
+        for k in self.data:
+            values[k].extend([self.data[k]])
+        return values
+
+class FastApiOAuth2Request(OAuth2Request):
+    def __init__(self, request: Request):
+        with patch('authlib.oauth2.rfc6749.errors.InsecureTransportError'):
+            super().__init__(
+                method=request.method,
+                uri=str(request.url),
+                headers=request.headers,
+            )
+        self.method = request.method
+        self.uri = str(request.url)
+        self.headers = request.headers
+        self.payload = FastApiOAuth2Payload(request)
+        self.user = request.user
+
+        self._request = request
+
+    @property
+    def args(self):
+        return self._request.query_params
+
+    @property
+    def form(self):
+        return getattr(self._request.state, 'form_data', {})
+
+
+class AuthorizationServer(BaseAuthorizationServer):
+    def __init__(self, *args, **kwargs):
+        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),
+        ))
+
+    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()
+
+    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
+            ))
+
+    def create_oauth2_request(self, request: Request) -> OAuth2Request:
+        return FastApiOAuth2Request(request)
+
+    def handle_response(self, status, body, headers) -> Response:
+        return Response(
+            status_code=status,
+            content=json.dumps(body) if isinstance(body, dict) else body,
+            headers=dict(headers),
+        )
+
+    def send_signal(self, name, *args, **kwargs):
+        pass
diff --git a/authentication/backend.py b/authentication/backend.py
index 29408e3..9434906 100644
--- a/authentication/backend.py
+++ b/authentication/backend.py
@@ -1,4 +1,4 @@
-from datetime import datetime, UTC
+import re
 from typing import Optional
 
 from sqlalchemy import update
@@ -6,11 +6,11 @@ from sqlmodel import select
 from starlette.authentication import AuthenticationBackend, AuthCredentials, BaseUser, UnauthenticatedUser
 from starlette.requests import HTTPConnection
 
-from db.models import Session, User
+from db.models import Session, User, AuthToken
 from db.session import get_session_context
 
 
-class AuthBackend(AuthenticationBackend):
+class SessionAuthBackend(AuthenticationBackend):
     async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
         if "id" not in request.session:
             return None
@@ -33,16 +33,38 @@ class AuthBackend(AuthenticationBackend):
                 )
                 db_session.commit()
 
+
             return (
                 AuthCredentials(
-                    [
-                        "authenticated",
-                        *[
-                            connection.type for connection in session.user.app_connections
-                        ]
-                    ]
+                    ["authenticated", "session"]
                     if session is not None
                     else []
                 ),
                 session.user if session is not None else UnauthenticatedUser()
             )
+
+
+class BearerTokenAuthBackend(AuthenticationBackend):
+    regex = re.compile(r"^[Bb]earer\s(?P<token>\S+)$")
+    async def authenticate(self, request: HTTPConnection) -> Optional[tuple[AuthCredentials, BaseUser]]:
+        token = None
+        if "Authorization" in request.headers:
+            match = self.regex.match(request.headers["Authorization"])
+
+            query = (
+                select(AuthToken)
+                .join(User)
+                .where(AuthToken.access_token == match.group('token'))
+            )
+
+            with get_session_context() as db_session:
+                token: AuthToken = db_session.exec(query).first()
+
+        return (
+            AuthCredentials(
+                ["authenticated", "token"]
+                if token is not None
+                else []
+            ),
+            token.user if token is not None else UnauthenticatedUser()
+        )
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,
+    }
diff --git a/authentication/mixins.py b/authentication/mixins.py
new file mode 100644
index 0000000..3d6a91f
--- /dev/null
+++ b/authentication/mixins.py
@@ -0,0 +1,54 @@
+import json
+import time
+
+from authlib.integrations.sqla_oauth2 import OAuth2ClientMixin as BaseOAuth2ClientMixin, OAuth2AuthorizationCodeMixin as BaseOAuth2AuthorizationCodeMixin, OAuth2TokenMixin as BaseOAuth2TokenMixin
+from sqlalchemy.sql.sqltypes import Text
+from sqlmodel import Field
+
+
+class OAuth2ClientMixin(BaseOAuth2ClientMixin):
+    client_id: str = Field(max_length=48, index=True, nullable=True)
+    client_secret: str = Field(max_length=120, nullable=True)
+    client_id_issued_at: int = Field(nullable=False, default=0)
+    client_secret_expires_at: int = Field(nullable=False, default=0)
+    client_metadata_field: str = Field(sa_type=Text, sa_column_kwargs={
+        'name': 'client_metadata'
+    }, nullable=True)
+
+    _client_metadata = None
+
+    @property
+    def client_metadata(self) -> dict:
+        if self.client_metadata_field:
+            return json.loads(self.client_metadata_field)
+        return {}
+
+    def set_client_metadata(self, metadata: dict):
+        self.client_metadata_field = json.dumps(metadata)
+
+
+class OAuth2AuthorizationCodeMixin(BaseOAuth2AuthorizationCodeMixin):
+    code: str = Field(max_length=120, unique=True, nullable=False)
+    client_id: str = Field(max_length=48, nullable=True)
+    redirect_uri: str = Field(sa_type=Text, default="", nullable=True)
+    response_type: str = Field(sa_type=Text, default="", nullable=True)
+    scope: str = Field(sa_type=Text, default="", nullable=True)
+    nonce: str = Field(sa_type=Text, nullable=True)
+    auth_time: int = Field(nullable=False, default_factory=lambda: int(time.time()))
+    acr: str = Field(sa_type=Text, nullable=True)
+    amr: str = Field(sa_type=Text, nullable=True)
+
+    code_challenge: str = Field(sa_type=Text, nullable=True)
+    code_challenge_method: str = Field(max_length=48, nullable=True)
+
+
+class OAuth2TokenMixin(BaseOAuth2TokenMixin):
+    client_id: str = Field(max_length=48, nullable=True)
+    token_type: str = Field(max_length=40, nullable=True)
+    access_token: str = Field(max_length=255, unique=True, nullable=False)
+    refresh_token: str = Field(max_length=255, index=True, nullable=True)
+    scope: str = Field(sa_type=Text, default="", nullable=True)
+    issued_at: int = Field(nullable=False, default_factory=lambda: int(time.time()))
+    access_token_revoked_at: int = Field(nullable=False, default=0)
+    refresh_token_revoked_at: int = Field(nullable=False, default=0)
+    expires_in: int = Field(nullable=False, default=0)
diff --git a/authentication/templates/authorize.html b/authentication/templates/authorize.html
deleted file mode 100644
index 1ed6506..0000000
--- a/authentication/templates/authorize.html
+++ /dev/null
@@ -1,5 +0,0 @@
-{% extends "./base.html" %}
-
-{% block content %}
-    <h1>AUTH</h1>
-{%  endblock %}
diff --git a/authentication/templates/login.html b/authentication/templates/login.html
index fcba5cd..2db2512 100644
--- a/authentication/templates/login.html
+++ b/authentication/templates/login.html
@@ -3,7 +3,7 @@
 {% block content %}
     <form method="post">
 
-        <input type="email" name="username" {% if form %}value="{{ form.username }}"{% endif %} placeholder="username" />
+        <input type="text" name="username" {% if form %}value="{{ form.username }}"{% endif %} placeholder="username" />
         {% if errors and errors.username %}
             {{ errors.username }}
         {% endif %}
diff --git a/authentication/utils.py b/authentication/utils.py
index f383969..4f19bc3 100644
--- a/authentication/utils.py
+++ b/authentication/utils.py
@@ -1,6 +1,9 @@
 import base64
 import hashlib
+from urllib.parse import urlencode
+
 import bcrypt
+from fastapi import FastAPI
 
 
 def bcrypt_sha256_hash(password: str) -> str:
@@ -13,3 +16,10 @@ def bcrypt_sha256_verify(password: str, hashed: str) -> bool:
     digest = hashlib.sha256(password.encode("utf-8")).digest()
     encoded = base64.b64encode(digest)
     return bcrypt.checkpw(encoded, hashed.encode('utf-8'))
+
+
+def get_url_path_with_query(app: FastAPI, name: str, params: dict | None = None, /, **path_params: str):
+    base = app.url_path_for(name, **path_params)
+    return base + '?' + urlencode(params) if params is not None else base
+
+