ad/molodetz/main.py

# retoor <retoor@molodetz.nl>
import fcntl
import logging
import math
import time
import traceback
from collections import defaultdict
from contextlib import asynccontextmanager
from datetime import datetime, timezone

from fastapi import FastAPI, HTTPException, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, PlainTextResponse, RedirectResponse, Response
from fastapi.staticfiles import StaticFiles
from starlette.middleware.gzip import GZipMiddleware

from molodetz import config
from molodetz.attachments import INLINE_EXTENSIONS, extension_of
from molodetz.constants import (
    CRAWLER_MARKERS,
    CRAWLER_PRIVATE_PREFIXES,
    GZIP_SKIP_PREFIXES,
    MAINTENANCE_ALLOWED_PREFIXES,
    RATE_LIMIT_EXEMPT_PREFIXES,
    TERMS_EXEMPT_PREFIXES,
)
from molodetz.database import flush_visits, get_int_setting, get_setting, init_db, refresh_snapshot
from molodetz.emoji_builder import GENERATED_JS, write_js_module
from molodetz.friendly404 import friendly_message
from molodetz.html_reformat import reformat_html
from molodetz.responses import json_error, wants_json
from molodetz.routers import (
    auth,
    avatar,
    docs,
    flyers,
    health,
    home,
    invite,
    join,
    legacy,
    memes,
    notifications,
    people,
    posts,
    presence,
    privacy,
    profile,
    roll,
    seo,
    standard,
    terms,
)
from molodetz.routers import admin
from molodetz.bootstrap import bootstrap_content
from molodetz.services.base import manager
from molodetz.services.queue import background
from molodetz.services.registry import build_services
from molodetz.templating import templates
from molodetz.utils.audit import record
from molodetz.utils.auth import GuardRedirect, get_current_user, is_admin
from molodetz.utils.presence import touch
from molodetz.utils.request_meta import client_ip, safe_next

logging.basicConfig(level=config.LOG_LEVEL, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger("molodetz")

_visit_counts = defaultdict(int)
_visit_state = {"last_flush": time.monotonic()}
_lock_handles = {}


def _take_lock(path, blocking):
    handle = open(path, "a+")
    flags = fcntl.LOCK_EX if blocking else fcntl.LOCK_EX | fcntl.LOCK_NB
    try:
        fcntl.flock(handle, flags)
    except BlockingIOError:
        handle.close()
        return None
    return handle


def _flush_visit_counts():
    if not _visit_counts:
        return
    day = datetime.now(timezone.utc).date().isoformat()
    snapshot = dict(_visit_counts)
    _visit_counts.clear()
    _visit_state["last_flush"] = time.monotonic()
    background.submit(flush_visits, day, snapshot)


@asynccontextmanager
async def lifespan(app):
    config.ensure_data_dirs()
    init_handle = _take_lock(config.INIT_LOCK_FILE, blocking=True)
    try:
        init_db()
        bootstrap_content()
        if not GENERATED_JS.exists():
            write_js_module()
    finally:
        fcntl.flock(init_handle, fcntl.LOCK_UN)
        init_handle.close()
    for service in build_services():
        manager.register(service)
    if not config.DISABLE_SERVICES:
        background.start()
        service_handle = _take_lock(config.SERVICE_LOCK_FILE, blocking=False)
        if service_handle is not None:
            _lock_handles["service"] = service_handle
            manager.supervise()
            logger.info("service lock acquired; supervising %d services", len(manager.services))
        else:
            logger.info("service lock held elsewhere; skipping supervision")
    yield
    if _visit_counts:
        _flush_visit_counts()
    await manager.shutdown_all()
    await background.stop()
    for handle in _lock_handles.values():
        handle.close()
    _lock_handles.clear()


app = FastAPI(
    title=config.APP_TITLE,
    docs_url="/swagger",
    redoc_url=None,
    swagger_ui_oauth2_redirect_url=None,
    openapi_url="/openapi.json",
    lifespan=lifespan,
)

app.include_router(auth.router, prefix="/auth")
app.include_router(admin.router, prefix="/admin")
app.include_router(avatar.router, prefix="/avatar")
app.include_router(profile.router, prefix="/profile")
app.include_router(notifications.router, prefix="/notifications")
app.include_router(presence.router, prefix="/presence")
app.include_router(posts.router, prefix="/posts")
app.include_router(roll.router, prefix="/roll")
app.include_router(standard.router, prefix="/standard")
app.include_router(flyers.router, prefix="/flyers")
app.include_router(memes.router, prefix="/memes")
app.include_router(people.router, prefix="/people")
app.include_router(join.router, prefix="/join")
app.include_router(invite.router, prefix="/invite")
app.include_router(terms.router, prefix="/terms")
app.include_router(privacy.router, prefix="/privacy")
app.include_router(health.router, prefix="/health")
app.include_router(seo.router)
app.include_router(legacy.router)
app.include_router(docs.router)
app.include_router(home.router)


class UploadsStatic(StaticFiles):
    async def get_response(self, path, scope):
        response = await super().get_response(path, scope)
        disposition = "inline" if extension_of(path) in INLINE_EXTENSIONS else "attachment"
        response.headers["Content-Disposition"] = disposition
        response.headers["Cache-Control"] = "public, max-age=604800"
        return response


class VersionedStatic(StaticFiles):
    async def get_response(self, path, scope):
        response = await super().get_response(path, scope)
        if path.endswith("sw.js"):
            response.headers["Cache-Control"] = "no-cache"
        else:
            response.headers["Cache-Control"] = "public, max-age=31536000, immutable"
        return response


class FallbackStatic(StaticFiles):
    async def get_response(self, path, scope):
        response = await super().get_response(path, scope)
        response.headers["Cache-Control"] = "no-cache" if path.endswith("sw.js") else "public, max-age=3600"
        return response


app.mount("/static/uploads", UploadsStatic(directory=str(config.DATA_PATHS["uploads"]), check_dir=False), name="uploads")
app.mount(f"/static/v{config.STATIC_VERSION}", VersionedStatic(directory=str(config.STATIC_DIR)), name="static_versioned")
app.mount("/static", FallbackStatic(directory=str(config.STATIC_DIR)), name="static")


def _is_mutating(request):
    return request.method in ("POST", "PUT", "PATCH", "DELETE")


@app.middleware("http")
async def snapshot_middleware(request: Request, call_next):
    path = request.url.path
    if not path.startswith(("/static", "/avatar")):
        refresh_snapshot()
    return await call_next(request)


@app.middleware("http")
async def security_headers_middleware(request: Request, call_next):
    response = await call_next(request)
    response.headers.setdefault("X-Robots-Tag", "index, follow")
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
    response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
    response.headers["Content-Security-Policy"] = (
        "object-src 'none'; base-uri 'self'; frame-ancestors 'self'; form-action 'self'"
    )
    if request.url.path.startswith("/admin"):
        response.headers["Cache-Control"] = "no-store"
        response.headers["Pragma"] = "no-cache"
        response.headers["Expires"] = "0"
    return response


_rate_buckets = defaultdict(list)
_rate_state = {"limit": None, "window": None, "loaded": 0.0, "swept": time.monotonic()}


def _rate_settings():
    now = time.monotonic()
    if now - _rate_state["loaded"] > 2:
        _rate_state["limit"] = max(1, get_int_setting("rate_limit_per_minute", config.RATE_LIMIT))
        _rate_state["window"] = max(1, get_int_setting("rate_limit_window_seconds", 60))
        _rate_state["loaded"] = now
    per_worker = max(1, math.ceil(_rate_state["limit"] / max(1, config.WEB_WORKERS)))
    return per_worker, _rate_state["window"]


@app.middleware("http")
async def rate_limit_middleware(request: Request, call_next):
    if config.DISABLE_RATE_LIMIT or not _is_mutating(request):
        return await call_next(request)
    if RATE_LIMIT_EXEMPT_PREFIXES and request.url.path.startswith(RATE_LIMIT_EXEMPT_PREFIXES):
        return await call_next(request)
    limit, window = _rate_settings()
    now = time.monotonic()
    if now - _rate_state["swept"] > 60:
        for key in list(_rate_buckets):
            _rate_buckets[key] = [stamp for stamp in _rate_buckets[key] if now - stamp < window]
            if not _rate_buckets[key]:
                del _rate_buckets[key]
        _rate_state["swept"] = now
    key = client_ip(request)
    bucket = [stamp for stamp in _rate_buckets[key] if now - stamp < window]
    if len(bucket) >= max(1, limit):
        _rate_buckets[key] = bucket
        retry = max(1, int(window - (now - bucket[0])))
        record(request, "security.ratelimit.exceeded", result="denied", payload={"ip": key})
        if wants_json(request):
            response = json_error(429, "Too many requests", retry_after=retry)
        else:
            response = PlainTextResponse("Too many requests. Try again in a moment.", status_code=429)
        response.headers["Retry-After"] = str(retry)
        return response
    bucket.append(now)
    _rate_buckets[key] = bucket
    return await call_next(request)


@app.middleware("http")
async def maintenance_middleware(request: Request, call_next):
    if get_setting("maintenance_mode", "0") != "1":
        return await call_next(request)
    if request.url.path.startswith(MAINTENANCE_ALLOWED_PREFIXES) or is_admin(get_current_user(request)):
        return await call_next(request)
    message = get_setting("maintenance_message", "Maintenance")
    record(request, "security.maintenance.blocked", result="denied")
    if wants_json(request):
        return json_error(503, message)
    return templates.TemplateResponse(
        request,
        "error.html",
        {"status_code": 503, "message": message, "seo": {"title": "Maintenance", "robots": "noindex,nofollow"}},
        status_code=503,
    )


@app.middleware("http")
async def terms_middleware(request: Request, call_next):
    if not _is_mutating(request) or request.url.path.startswith(TERMS_EXEMPT_PREFIXES):
        return await call_next(request)
    user = get_current_user(request)
    version = get_setting("terms_version", "1")
    if user is not None and str(user.get("terms_version") or "") != str(version):
        target = f"/terms?next={request.url.path}"
        if wants_json(request):
            return json_error(403, "Accept the terms first", code="terms_required", redirect=target, terms_version=version)
        return RedirectResponse(target, status_code=303)
    return await call_next(request)


@app.middleware("http")
async def crawler_block_middleware(request: Request, call_next):
    if request.method == "GET" and request.url.path.startswith(CRAWLER_PRIVATE_PREFIXES):
        agent = (request.headers.get("user-agent") or "").lower()
        if any(marker in agent for marker in CRAWLER_MARKERS):
            record(request, "security.crawler.blocked", result="denied", payload={"agent": agent[:120]})
            if wants_json(request):
                return json_error(403, "Crawlers are not allowed here")
            return PlainTextResponse("Forbidden", status_code=403)
    return await call_next(request)


@app.middleware("http")
async def presence_middleware(request: Request, call_next):
    if not request.url.path.startswith(("/static", "/avatar")):
        user = get_current_user(request)
        if user:
            touch(user["uid"])
    return await call_next(request)


@app.middleware("http")
async def visits_middleware(request: Request, call_next):
    response = await call_next(request)
    path = request.url.path
    if not path.startswith(("/static", "/avatar")):
        _visit_counts[(path[:200], response.status_code)] += 1
        if time.monotonic() - _visit_state["last_flush"] > 30:
            _flush_visit_counts()
    return response


@app.middleware("http")
async def html_reformat_middleware(request: Request, call_next):
    response = await call_next(request)
    content_type = response.headers.get("content-type", "")
    if not content_type.startswith("text/html") or response.headers.get("content-encoding"):
        return response
    body = b""
    async for chunk in response.body_iterator:
        body += chunk
    try:
        body = reformat_html(body.decode("utf-8")).encode("utf-8")
    except Exception:
        logger.warning("html reformat failed; serving original body")
    headers = dict(response.headers)
    headers["content-length"] = str(len(body))
    return Response(content=body, status_code=response.status_code, headers=headers, media_type=None)


@app.middleware("http")
async def timing_middleware(request: Request, call_next):
    request.state.request_start = time.perf_counter()
    response = await call_next(request)
    response.headers["X-Response-Time"] = f"{(time.perf_counter() - request.state.request_start) * 1000:.1f}ms"
    return response


class SelectiveGZipMiddleware(GZipMiddleware):
    async def __call__(self, scope, receive, send):
        if scope["type"] == "http" and GZIP_SKIP_PREFIXES and scope["path"].startswith(GZIP_SKIP_PREFIXES):
            await self.app(scope, receive, send)
            return
        await super().__call__(scope, receive, send)


app.add_middleware(SelectiveGZipMiddleware, minimum_size=512, compresslevel=5)


def _error_page(request, status_code, message, headers=None):
    seo = {"title": f"{status_code}", "robots": "noindex,nofollow", "description": message}
    return templates.TemplateResponse(
        request,
        "error.html",
        {"status_code": status_code, "message": message, "seo": seo},
        status_code=status_code,
        headers=headers,
    )


@app.exception_handler(GuardRedirect)
async def guard_redirect_handler(request: Request, exc: GuardRedirect):
    return RedirectResponse(exc.url, status_code=exc.status_code)


@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    if wants_json(request) or request.headers.get("x-api-key") or request.headers.get("authorization"):
        return json_error(exc.status_code, str(exc.detail))
    if exc.status_code == 404:
        message = friendly_message() if get_setting("friendly_404", "1") == "1" else "Not found."
        return _error_page(request, 404, message)
    return _error_page(request, exc.status_code, str(exc.detail), headers=getattr(exc, "headers", None))


@app.exception_handler(404)
async def not_found_handler(request: Request, exc):
    return await http_exception_handler(request, HTTPException(404, "Not found"))


@app.exception_handler(500)
async def server_error_handler(request: Request, exc):
    logger.error("unhandled error: %s", "".join(traceback.format_exception(exc)))
    if wants_json(request):
        return json_error(500, "Internal error")
    return _error_page(request, 500, "Something broke on our side.")


AUTH_TEMPLATES = {"/auth/login": "auth/login.html"}


def _friendly_messages(errors):
    fields, messages = [], []
    for error in errors:
        loc = [str(part) for part in error.get("loc", ()) if part not in ("body", "query")]
        field = loc[-1] if loc else "body"
        message = str(error.get("msg", "Invalid input")).replace("Value error, ", "")
        if field == "password" and "at least" in message:
            message = "Password must be 6 to 128 characters."
        fields.append(field)
        messages.append(f"{field}: {message}" if field != "body" else message)
    return fields, messages


@app.exception_handler(RequestValidationError)
async def validation_handler(request: Request, exc: RequestValidationError):
    fields, messages = _friendly_messages(exc.errors())
    if wants_json(request):
        return JSONResponse({"error": "validation", "fields": fields, "messages": messages}, status_code=422)
    template = AUTH_TEMPLATES.get(request.url.path)
    if template:
        form = await request.form()
        echoed = {key: value for key, value in form.items() if isinstance(value, str) and key != "password"}
        return templates.TemplateResponse(
            request,
            template,
            {"errors": messages, "form": echoed, "next": echoed.get("next"), "seo": {"title": "Log in", "robots": "noindex,nofollow"}},
            status_code=400,
        )
    if request.url.path == "/join":
        form = await request.form()
        echoed = {key: value for key, value in form.items() if isinstance(value, str)}
        return templates.TemplateResponse(
            request,
            "join.html",
            {"errors": messages, "form": echoed, "submitted": False, "seo": {"title": "Join", "robots": "index,follow"}},
            status_code=400,
        )
    if request.url.path.startswith("/invite/"):
        from molodetz.database import get_invite_by_token

        form = await request.form()
        echoed = {key: value for key, value in form.items() if isinstance(value, str) and not key.startswith("password")}
        invite = get_invite_by_token(request.url.path.rsplit("/", 1)[-1])
        return templates.TemplateResponse(
            request,
            "invite.html",
            {
                "valid": True,
                "errors": messages,
                "form": echoed,
                "email_locked": bool(invite and invite.get("email")),
                "seo": {"title": "Accept invite", "robots": "noindex,nofollow"},
            },
            status_code=400,
        )
    return RedirectResponse(safe_next(request.headers.get("referer", "").replace(str(request.base_url).rstrip("/"), ""), "/"), status_code=303)