App/mcp/auth.py

from __future__ import annotations

import os
from typing import Callable, Optional

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.types import ASGIApp


def get_mcp_api_key() -> Optional[str]:
    value = os.getenv("UWEKEZAJI_MCP_API_KEY", "").strip()
    return value or None


def _extract_bearer_token(request: Request) -> Optional[str]:
    auth_header = request.headers.get("authorization", "")
    if auth_header.lower().startswith("bearer "):
        return auth_header[7:].strip()
    return request.headers.get("x-api-key")


class McpApiKeyMiddleware(BaseHTTPMiddleware):
    """Optional API-key gate for remotely hosted MCP endpoints."""

    def __init__(self, app: ASGIApp, api_key: str):
        super().__init__(app)
        self.api_key = api_key

    async def dispatch(self, request: Request, call_next: Callable):
        if request.url.path in {"/health", "/healthz"}:
            return await call_next(request)

        provided = _extract_bearer_token(request)
        if provided != self.api_key:
            return JSONResponse(
                status_code=401,
                content={
                    "error": "unauthorized",
                    "message": "Valid MCP API key required.",
                },
            )
        return await call_next(request)