nocobase-mcp/nocobase_crm_mcp/server_runtime.py

from __future__ import annotations

import hmac
import os
from typing import Any

from starlette.applications import Starlette
from starlette.routing import Mount, Route

from .uploads import UploadEndpoint

DEFAULT_HTTP_HOST = "127.0.0.1"
DEFAULT_HTTP_PORT = 10087
DEFAULT_HTTP_PATH = "/mcp"
DEFAULT_TOKEN_ENV = "NOCOBASE_MCP_TOKEN"
ENABLE_UPLOADS_ENV = "NOCOBASE_MCP_ENABLE_UPLOADS"
UPLOAD_API_BASE_URL_ENV = "NOCOBASE_MCP_UPLOAD_API_BASE_URL"
UPLOAD_JWT_ENV = "NOCOBASE_MCP_UPLOAD_JWT"
UPLOAD_MAX_BYTES_ENV = "NOCOBASE_MCP_UPLOAD_MAX_BYTES"
API_BASE_URL_ENV = "NOCOBASE_MCP_API_BASE_URL"
API_TOKEN_ENV = "NOCOBASE_MCP_API_TOKEN"
DEFAULT_UPLOAD_MAX_BYTES = 25 * 1024 * 1024


def _bearer_token(authorization: str) -> str:
    scheme, _, token = authorization.partition(" ")
    if scheme.lower() != "bearer" or not token:
        return ""
    return token.strip()


def _authorization_header(scope: dict[str, Any]) -> str:
    headers = {key.lower(): value for key, value in scope.get("headers", [])}
    return headers.get(b"authorization", b"").decode("latin1")


class BearerTokenAuthMiddleware:
    def __init__(self, app: Any, token: str) -> None:
        if not token:
            raise ValueError("HTTP MCP token must not be empty")
        self.app = app
        self.token = token

    async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return
        if str(scope.get("method") or "").upper() == "OPTIONS":
            await send(
                {
                    "type": "http.response.start",
                    "status": 204,
                    "headers": [
                        (b"access-control-allow-origin", b"*"),
                        (b"access-control-allow-methods", b"POST, OPTIONS"),
                        (b"access-control-allow-headers", b"authorization, content-type, mcp-session-id"),
                    ],
                }
            )
            await send({"type": "http.response.body", "body": b""})
            return
        token = _bearer_token(_authorization_header(scope))
        if not hmac.compare_digest(token, self.token):
            await send(
                {
                    "type": "http.response.start",
                    "status": 401,
                    "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")],
                }
            )
            await send({"type": "http.response.body", "body": b'{"error":"unauthorized"}'})
            return
        await self.app(scope, receive, send)


def read_token_from_env(env_name: str) -> str:
    token = os.environ.get(env_name, "").strip()
    if not token:
        raise RuntimeError(f"{env_name} must be set for HTTP MCP transport")
    return token


def create_http_app(server: Any, token: str, upload_endpoint: UploadEndpoint | None = None) -> Any:
    if upload_endpoint is None:
        return BearerTokenAuthMiddleware(server.streamable_http_app(), token)

    async def upload_route(request: Any) -> Any:
        return await upload_endpoint(request)

    mcp_app = server.streamable_http_app()
    app = Starlette(
        routes=[Route("/uploads", endpoint=upload_route, methods=["POST"]), Mount("/", app=mcp_app)],
        lifespan=mcp_app.router.lifespan_context,
    )
    return BearerTokenAuthMiddleware(app, token)


def read_nocobase_environment() -> tuple[str | None, str | None, int]:
    canonical_url = os.environ.get(API_BASE_URL_ENV, "").strip().rstrip("/")
    canonical_token = os.environ.get(API_TOKEN_ENV, "").strip()
    legacy_url = os.environ.get(UPLOAD_API_BASE_URL_ENV, "").strip().rstrip("/")
    legacy_token = os.environ.get(UPLOAD_JWT_ENV, "").strip()

    canonical_any = bool(canonical_url or canonical_token)
    legacy_any = bool(legacy_url or legacy_token)
    if canonical_any and not (canonical_url and canonical_token):
        raise RuntimeError(f"{API_BASE_URL_ENV} and {API_TOKEN_ENV} must be set together")
    if legacy_any and not (legacy_url and legacy_token):
        raise RuntimeError(f"{UPLOAD_API_BASE_URL_ENV} and {UPLOAD_JWT_ENV} must be set together")
    if canonical_any and legacy_any and (canonical_url != legacy_url or canonical_token != legacy_token):
        raise RuntimeError("Canonical and legacy NocoBase credential environment variables conflict")
    if canonical_any:
        api_base_url, api_token = canonical_url, canonical_token
    elif legacy_any:
        api_base_url, api_token = legacy_url, legacy_token
    else:
        api_base_url, api_token = None, None

    raw_max_bytes = os.environ.get(UPLOAD_MAX_BYTES_ENV, str(DEFAULT_UPLOAD_MAX_BYTES)).strip()
    try:
        max_bytes = int(raw_max_bytes)
    except ValueError as error:
        raise RuntimeError(f"{UPLOAD_MAX_BYTES_ENV} must be a positive integer") from error
    if max_bytes <= 0:
        raise RuntimeError(f"{UPLOAD_MAX_BYTES_ENV} must be a positive integer")
    return api_base_url, api_token, max_bytes


_read_token_from_env = read_token_from_env
_read_nocobase_environment = read_nocobase_environment

__all__ = [
    "API_BASE_URL_ENV",
    "API_TOKEN_ENV",
    "BearerTokenAuthMiddleware",
    "DEFAULT_HTTP_HOST",
    "DEFAULT_HTTP_PATH",
    "DEFAULT_HTTP_PORT",
    "DEFAULT_TOKEN_ENV",
    "DEFAULT_UPLOAD_MAX_BYTES",
    "ENABLE_UPLOADS_ENV",
    "UPLOAD_API_BASE_URL_ENV",
    "UPLOAD_JWT_ENV",
    "UPLOAD_MAX_BYTES_ENV",
    "create_http_app",
    "read_nocobase_environment",
    "read_token_from_env",
    "_read_nocobase_environment",
    "_read_token_from_env",
]