be-optimumoh/src/main.py

import logging
import asyncio
import time
from src.context import set_request_id, reset_request_id, get_request_id
from contextvars import ContextVar
from os import path
from typing import Final, Optional
from uuid import uuid1

from fastapi import Depends, FastAPI, HTTPException, status
from src.auth.service import JWTBearer
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import ValidationError
from slowapi import _rate_limit_exceeded_handler
from slowapi.errors import RateLimitExceeded
from slowapi.middleware import SlowAPIMiddleware
from sqlalchemy import inspect
from sqlalchemy.ext.asyncio import async_scoped_session
from sqlalchemy.orm import scoped_session
from starlette.middleware.base import (BaseHTTPMiddleware,
                                       RequestResponseEndpoint)
from starlette.middleware.gzip import GZipMiddleware
from starlette.requests import Request
from starlette.responses import FileResponse, Response, StreamingResponse
from starlette.routing import compile_path
from starlette.staticfiles import StaticFiles

import uuid as _uuid_module

from src.api import api_router
from src.database.core import async_session, engine, async_collector_session
from src.enums import ResponseStatus
from src.exceptions import handle_exception
from src.app_logging import configure_logging
from src.middleware import RequestValidationMiddleware, RequestTimeoutMiddleware
# from src.rate_limiter import limiter
from fastapi.exceptions import RequestValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException
from src.csrf_protect import csrf_protect

log = logging.getLogger(__name__)

# we configure the logging level and format
configure_logging()

REQUEST_TIMEOUT_SECONDS = 20

# we create the ASGI for the app
app = FastAPI(
    openapi_url="",
    title="Optimum OH API",
    description="Welcome to Optimum OH's API documentation!",
    version="0.1.0",
)

# Sync handler for RateLimitExceeded — SlowAPIMiddleware falls back to default
# when the registered handler is async, so this must be a plain sync function.
def _sync_rate_limit_handler(request: Request, exc: RateLimitExceeded) -> JSONResponse:
    error_id = str(_uuid_module.uuid4())
    request.state.log_status_code = 429
    request.state.log_error_id = error_id
    request.state.log_error_type = "Too Many Requests"
    request.state.log_error_message = "Please try again later"
    return JSONResponse(
        status_code=429,
        content={"data": None, "message": f"Please try again later. Error ID: {error_id}", "status": ResponseStatus.ERROR},
    )

# we define the exception handlers
app.add_exception_handler(Exception, handle_exception)
app.add_exception_handler(HTTPException, handle_exception)
app.add_exception_handler(StarletteHTTPException, handle_exception)
app.add_exception_handler(RequestValidationError, handle_exception)
app.add_exception_handler(RateLimitExceeded, _sync_rate_limit_handler)

# app.state.limiter = limiter
app.add_middleware(GZipMiddleware, minimum_size=1000)
app.add_middleware(SlowAPIMiddleware)
app.add_middleware(RequestTimeoutMiddleware, timeout=REQUEST_TIMEOUT_SECONDS)

# credentials: "include",


@app.route('/slow', methods=["GET"])
async def slow_endpoint(request: Request):
    await asyncio.sleep(REQUEST_TIMEOUT_SECONDS + 10)
    return JSONResponse(content={"message": "This should timeout"}, status_code=200)

@app.post('/testcsrf', dependencies=[Depends(JWTBearer()), Depends(csrf_protect)])
async def test_csrf(request: Request):
    return JSONResponse(
        content={"data": None, "message": "CSRF token is valid", "status": ResponseStatus.SUCCESS},
        status_code=200,
    )


def security_headers_middleware(app: FastAPI):
    is_production = False

    # CSP rules
    csp_policy = {
        "default-src": "'self'",
        "script-src": "'self' 'unsafe-inline' https://cdnjs.cloudflare.com https://cdn.jsdelivr.net",
        "style-src": "'self' 'unsafe-inline' https://fonts.googleapis.com https://cdn.jsdelivr.net",
        "img-src": "'self' data: https: blob:",
        "font-src": "'self' https://fonts.gstatic.com data:",
        "connect-src": "'self' https://api.your-domain.com wss://ws.your-domain.com",
        "frame-src": "'none'",
        "object-src": "'none'",
        "base-uri": "'self'",
        "form-action": "'self'",
    }

    # Feature / Permissions Policy
    feature_policy = {
        "geolocation": "'none'",
        "midi": "'none'",
        "notifications": "'none'",
        "push": "'none'",
        "sync-xhr": "'none'",
        "microphone": "'none'",
        "camera": "'none'",
        "magnetometer": "'none'",
        "gyroscope": "'none'",
        "speaker": "'none'",
        "vibrate": "'none'",
        "fullscreen": "'self'",
        "payment": "'none'",
    }

    csp_header_value = "; ".join(f"{k} {v}" for k, v in csp_policy.items())
    feature_header_value = "; ".join(f"{k}={v}" for k, v in feature_policy.items())

    # Middleware definition
    class SecurityHeadersMiddleware(BaseHTTPMiddleware):
        async def dispatch(self, request: Request, call_next):
            response: Response = await call_next(request)

            if is_production:
                response.headers["Strict-Transport-Security"] = "max-age=15724800; includeSubDomains; preload"
                response.headers["X-Frame-Options"] = "DENY"
                response.headers["X-Content-Type-Options"] = "nosniff"
                response.headers["X-XSS-Protection"] = "1; mode=block"
                response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
                response.headers["Content-Security-Policy"] = csp_header_value
                response.headers["Permissions-Policy"] = feature_header_value
            else:
                # Relaxed settings for development
                response.headers["Content-Security-Policy"] = "default-src 'self' 'unsafe-inline' 'unsafe-eval' *"
                # You can skip some headers here for local testing

            return response

    app.add_middleware(SecurityHeadersMiddleware)

security_headers_middleware(app)


# app.add_middleware(RequestValidationMiddleware)
# app.add_middleware(
#     CSRFProtectMiddleware,
#     excluded_paths=["/healthcheck"]
# )




@app.middleware("http")
async def db_session_middleware(request: Request, call_next):
    request_id = str(uuid1())

    # Per-request ID scopes the SQLAlchemy session.
    # see: https://github.com/tiangolo/fastapi/issues/726
    ctx_token = set_request_id(request_id)
    request.state.request_id = request_id

    try:
        session = async_scoped_session(async_session, scopefunc=get_request_id)
        request.state.db = session()

        collector_session = async_scoped_session(async_collector_session, scopefunc=get_request_id)
        request.state.collector_db = collector_session()

        start_time = time.time()
        try:
            response = await call_next(request)
        except Exception as exc:
            # Non-HTTP exceptions (SQLAlchemy, httpx, ValueError, etc.) are NOT caught by
            # ExceptionMiddleware because app.add_exception_handler(Exception, ...) registers
            # with ServerErrorMiddleware, not ExceptionMiddleware. If we let such exceptions
            # escape here, db_session_middleware exits via finally WITHOUT logging, then
            # ServerErrorMiddleware calls handle_exception (giving the 500 response) and
            # re-raises, causing uvicorn to emit an all-null JSON log entry.
            # Catching here ensures: proper error response, structured log, no re-raise.
            response = await handle_exception(request, exc)
        duration_ms = round((time.time() - start_time) * 1000, 2)

        # ── Determine status_code ─────────────────────────────────────────────
        # RequestTimeoutMiddleware flags the scope when a 408 bypasses exception handlers
        timed_out = request.scope.get("x_timed_out", False)
        if timed_out:
            status_code = 408
            error_id = request.scope.get("x_timeout_error_id", None)
            error_type = "Request Timeout"
            error_message = "Request timeout over 20 seconds"
        else:
            status_code = getattr(request.state, "log_status_code", None) or response.status_code
            error_id = getattr(request.state, "log_error_id", None)
            error_type = getattr(request.state, "log_error_type", None)
            error_message = getattr(request.state, "log_error_message", None)

        # ── Fallback error labels for responses that skip exception handler ───
        _FALLBACK_ERROR_TYPES = {
            400: "Bad Request",    401: "Unauthorized",   403: "Forbidden",
            404: "Not Found",      408: "Request Timeout", 409: "Conflict",
            422: "Unprocessable Entity", 429: "Too Many Requests",
            500: "Internal Server Error",
        }
        _FALLBACK_ERROR_MESSAGES = {
            400: "Bad request",           401: "Invalid authorization token",
            403: "Forbidden Access",      404: "Resource not found",
            408: "Request timeout",       409: "This record already exists",
            422: "Input validation failed", 429: "Please try again later",
            500: "An unexpected error occurred",
        }
        if status_code >= 400 and not error_type:
            error_type = _FALLBACK_ERROR_TYPES.get(status_code, f"HTTP {status_code}")
        if status_code >= 400 and not error_message:
            error_message = _FALLBACK_ERROR_MESSAGES.get(status_code, "Request failed")

        # ── User info ─────────────────────────────────────────────────────────
        from src.context import get_user_id, set_user_id
        user_id = get_user_id()
        user_obj = getattr(request.state, "user", None)
        if user_obj and not user_id:
            if isinstance(user_obj, dict):
                u_id = user_obj.get("user_id")
            else:
                u_id = getattr(user_obj, "user_id", None)
            if u_id:
                user_id = str(u_id)
                set_user_id(user_id)

        # ── Client IP ─────────────────────────────────────────────────────────
        headers = request.headers
        if "X-Real-IP" in headers:
            user_ip = headers["X-Real-IP"]
        elif "X-Forwarded-For" in headers:
            user_ip = headers["X-Forwarded-For"].split(",")[0].strip()
        else:
            user_ip = request.client.host if request.client else None

        # ── Full path with query string ───────────────────────────────────────
        path = request.url.path
        if request.url.query:
            path = f"{path}?{request.url.query}"

        # ── Emit single structured log ────────────────────────────────────────
        result = "Request completed" if status_code < 400 else "Request failed"
        if status_code >= 500:
            log_level = logging.ERROR
        elif status_code >= 400:
            log_level = logging.WARNING
        else:
            log_level = logging.INFO

        log.log(
            log_level,
            f"{request.method} {path} {status_code}",
            extra={
                "method": request.method,
                "path": path,
                "status_code": status_code,
                "duration_ms": duration_ms,
                "request_id": request_id,
                "user_id": user_id,
                "user_ip": user_ip,
                "error_id": error_id,
                "result": result,
                "error_type": error_type if status_code >= 400 else None,
                "error_message": error_message if status_code >= 400 else None,
            },
        )
    finally:
        await request.state.db.close()
        await request.state.collector_db.close()

    reset_request_id(ctx_token)
    return response


@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["Strict-Transport-Security"] = (
        "max-age=31536000 ; includeSubDomains"
    )
    return response


# class MetricsMiddleware(BaseHTTPMiddleware):
#     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
#         method = request.method
#         endpoint = request.url.path
#         tags = {"method": method, "endpoint": endpoint}

#         try:
#             start = time.perf_counter()
#             response = await call_next(request)
#             elapsed_time = time.perf_counter() - start
#             tags.update({"status_code": response.status_code})
#             metric_provider.counter("server.call.counter", tags=tags)
#             metric_provider.timer("server.call.elapsed", value=elapsed_time, tags=tags)
#             log.debug(f"server.call.elapsed.{endpoint}: {elapsed_time}")
#         except Exception as e:
#             metric_provider.counter("server.call.exception.counter", tags=tags)
#             raise e from None
#         return response


# app.add_middleware(ExceptionMiddleware)

app.include_router(api_router)