From 498c79f434856aa68e9a883247ea69a22256fce6 Mon Sep 17 00:00:00 2001 From: Tom van der Lee Date: Wed, 22 Jul 2026 14:51:44 +0200 Subject: Added auth layer --- authentication/authlib.py | 167 ++++++++++++++++++++++++++++++++ authentication/backend.py | 40 ++++++-- authentication/endpoints.py | 88 ++++++++++++----- authentication/mixins.py | 54 +++++++++++ authentication/templates/authorize.html | 5 - authentication/templates/login.html | 2 +- authentication/utils.py | 10 ++ conf/__init__.py | 3 + conf/local.py | 3 + db/models.py | 18 +++- ttun_server/__init__.py | 9 +- ttun_server/websockets.py | 3 + 12 files changed, 360 insertions(+), 42 deletions(-) create mode 100644 authentication/authlib.py create mode 100644 authentication/mixins.py delete mode 100644 authentication/templates/authorize.html 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\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 %} -

AUTH

-{% 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 %}
- + {% 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 + + diff --git a/conf/__init__.py b/conf/__init__.py index bcdb9c5..744747f 100644 --- a/conf/__init__.py +++ b/conf/__init__.py @@ -1,5 +1,7 @@ import os +IS_LOCAL = False + try: from .local import * except ImportError: @@ -7,3 +9,4 @@ except ImportError: DATABASE_URL = os.environ.get('DATABASE_URL') SECRET_KEY = os.environ.get('SECRET_KEY') +ENABLE_AUTH = os.environ.get('ENABLE_AUTH', 'false') == 'true' diff --git a/conf/local.py b/conf/local.py index 19b304e..23f4005 100644 --- a/conf/local.py +++ b/conf/local.py @@ -1,4 +1,7 @@ import os +IS_LOCAL = True + os.environ.setdefault('DATABASE_URL', 'postgresql:///ttun') os.environ.setdefault('SECRET_KEY', 'secret') +os.environ.setdefault('ENABLE_AUTH', 'true') diff --git a/db/models.py b/db/models.py index 7e0afb0..8a01c8c 100644 --- a/db/models.py +++ b/db/models.py @@ -1,10 +1,13 @@ +import re from datetime import datetime, UTC +from urllib.parse import urlsplit from uuid import uuid4, UUID, uuid7 from sqlalchemy import DateTime from sqlmodel import SQLModel, Field, Relationship -from authlib.integrations.sqla_oauth2 import OAuth2ClientMixin, OAuth2AuthorizationCodeMixin, OAuth2TokenMixin +from starlette.authentication import BaseUser +from authentication.mixins import OAuth2ClientMixin, OAuth2AuthorizationCodeMixin, OAuth2TokenMixin from authentication.utils import bcrypt_sha256_hash @@ -24,7 +27,7 @@ class BaseModel(SQLModel): ) -class User(BaseModel, table=True): +class User(BaseModel, BaseUser, table=True): username: str = Field(unique=True, index=True) password: str = Field(nullable=False) @@ -33,6 +36,9 @@ class User(BaseModel, table=True): def set_password(self, value): self.password = bcrypt_sha256_hash(value) + def is_authenticated(self) -> bool: + return True + class Session(BaseModel, table=True): user_id: UUID = Field(foreign_key='user.id', exclude=True) user: User = Relationship() @@ -40,8 +46,16 @@ class Session(BaseModel, table=True): class ClientApplication(BaseModel, OAuth2ClientMixin, table=True): name: str + def check_redirect_uri(self, redirect_uri: str) -> bool: + return any( + re.fullmatch(re.escape(str(pattern)).replace(r"\*", r"\d+"), redirect_uri) + for pattern in self.redirect_uris + ) + class AuthCode(BaseModel, OAuth2AuthorizationCodeMixin, table=True): user_id: UUID = Field(foreign_key="user.id", ondelete="CASCADE") + user: User = Relationship() class AuthToken(BaseModel, OAuth2TokenMixin, table=True): user_id: UUID = Field(foreign_key="user.id", ondelete="CASCADE") + user: User = Relationship() diff --git a/ttun_server/__init__.py b/ttun_server/__init__.py index 4af6ca0..f5d6b45 100644 --- a/ttun_server/__init__.py +++ b/ttun_server/__init__.py @@ -2,8 +2,10 @@ import logging import os from fastapi import FastAPI +from starlette.middleware.authentication import AuthenticationMiddleware from starlette.routing import Host, Route, WebSocketRoute, Mount +from authentication.backend import BearerTokenAuthBackend from authentication.endpoints import authentication, oauth from proxy import app as proxy_app from ttun_server.endpoints import health @@ -16,12 +18,17 @@ app = FastAPI( routes=[ Mount('/auth/', app=authentication), Mount('/oauth/', app=oauth), - WebSocketRoute('/tunnel/', endpoint=tunnel), + WebSocketRoute( + '/tunnel/', + endpoint=tunnel, + ), Route('/health/', endpoint=health), Host(f'{{subdomain}}.{os.environ['TUNNEL_DOMAIN']}', app=proxy_app) ] ) +app.add_middleware(AuthenticationMiddleware, backend=BearerTokenAuthBackend()) + try: from ._version import version __version__ = version diff --git a/ttun_server/websockets.py b/ttun_server/websockets.py index c791625..88d6cac 100644 --- a/ttun_server/websockets.py +++ b/ttun_server/websockets.py @@ -4,7 +4,9 @@ import os from uuid import uuid4 from fastapi import WebSocket, WebSocketDisconnect +from starlette.authentication import requires +import conf import ttun_server from proxy.queue import ProxyQueue from ttun_server.types import ( @@ -31,6 +33,7 @@ async def assert_compatible_version(websocket: WebSocket, config: Config) -> Non await websocket.close(4001, 'Your client is too new') +@requires(['authenticated', 'token'] if conf.ENABLE_AUTH else []) async def tunnel(websocket: WebSocket) -> None: request_tasks: dict[str, asyncio.Task] = {} proxy_queues: dict[str, ProxyQueue] = {} -- cgit v1.2.3