wireguard_webadmin/containers/auth-gateway/auth_gateway/main.py

import logging
import time
from contextlib import asynccontextmanager
from pathlib import Path

from auth_gateway.config_loader import RuntimeConfigStore
from auth_gateway.limiter import get_real_client_ip, limiter
from auth_gateway.services.oidc_service import OIDCService
from auth_gateway.services.session_service import SessionService
from auth_gateway.settings import settings
from auth_gateway.storage.sqlite import SQLiteStorage
from auth_gateway.web.auth_routes import router as auth_router
from auth_gateway.web.login_routes import router as login_router
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from slowapi.errors import RateLimitExceeded

BASE_DIR = Path(__file__).resolve().parent
_access_logger = logging.getLogger("uvicorn.error")


@asynccontextmanager
async def lifespan(app: FastAPI):
    storage = SQLiteStorage(settings.database_path)
    app.state.settings = settings
    app.state.config_store = RuntimeConfigStore(settings.config_dir)
    app.state.session_service = SessionService(storage, settings.session_default_minutes, settings.oidc_state_ttl_minutes)
    app.state.oidc_service = OIDCService()
    app.state.templates = Jinja2Templates(directory=str(BASE_DIR / "templates"))
    yield


app = FastAPI(
    title="Auth Gateway",
    lifespan=lifespan,
    docs_url=None,
    redoc_url=None,
    openapi_url=None,
)
app.state.limiter = limiter


async def _rate_limit_handler(request: Request, exc: RateLimitExceeded) -> HTMLResponse:
    username = None
    try:
        form = await request.form()
        username = form.get("username")
    except Exception:
        pass
    client = get_real_client_ip(request)
    if username:
        _access_logger.warning("AUTH rate limit exceeded for '%s' on %s from %s", username, request.url.path, client)
    else:
        _access_logger.warning("AUTH rate limit exceeded on %s from %s", request.url.path, client)

    templates = request.app.state.templates
    external_path = request.app.state.settings.external_path
    return templates.TemplateResponse(
        "ratelimit.html",
        {"request": request, "external_path": external_path, "back_url": str(request.url.path)},
        status_code=429,
    )


app.add_exception_handler(RateLimitExceeded, _rate_limit_handler)


@app.middleware("http")
async def security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
    response.headers["Content-Security-Policy"] = (
        "default-src 'self'; script-src 'none'; style-src 'self'; "
        "img-src 'self' data:; frame-ancestors 'none'"
    )
    return response


@app.middleware("http")
async def access_log(request: Request, call_next):
    start = time.monotonic()
    response = await call_next(request)
    if request.url.path == "/auth/check" and response.status_code == 200:
        return response
    ms = (time.monotonic() - start) * 1000
    client = get_real_client_ip(request)
    _access_logger.info('%s - "%s %s" %d (%.0fms)', client, request.method, request.url.path, response.status_code, ms)
    return response


app.mount("/static", StaticFiles(directory=str(BASE_DIR / "static")), name="static")
app.include_router(auth_router)
app.include_router(login_router)


@app.get("/health")
async def healthcheck():
    return {"status": "ok"}