agent-platform/services/api-gateway/app/infrastructure/request_context.py

from dataclasses import dataclass
from datetime import datetime
from time import perf_counter
from uuid import uuid4

import httpx
from fastapi import Request, Response
from starlette.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from core_shared.security import build_internal_service_headers

from app.bootstrap.settings import ApiGatewaySettings
from app.domain.repositories import ApiKeyRepository
from app.infrastructure.api_keys import hash_api_key
from app.infrastructure.rate_limit import (
    apply_gateway_rate_limit_headers,
    enforce_gateway_rate_limit,
)

REQUEST_ID_HEADER = "x-request-id"
TENANT_ID_HEADER = "x-tenant-id"
DEFAULT_TENANT_ID = "public"


@dataclass
class GatewayRequestContext:
    request_id: str
    tenant_id: str
    started_perf_counter: float
    api_key_id: str | None = None
    user_id: str | None = None
    target_service: str | None = None
    target_url: str | None = None


class GatewayRequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        request_id = request.headers.get(REQUEST_ID_HEADER) or str(uuid4())
        tenant_id = resolve_tenant_id(request)
        request.state.gateway_context = GatewayRequestContext(
            request_id=request_id,
            tenant_id=tenant_id,
            started_perf_counter=perf_counter(),
        )
        auth_response = authenticate_gateway_request(request)
        if auth_response is not None:
            from app.infrastructure.audit import persist_gateway_audit

            persist_gateway_audit(
                request=request,
                session_factory=request.app.state.session_factory,
                status_code=auth_response.status_code,
                error_message=None,
            )
            context = get_gateway_request_context(request)
            auth_response.headers[REQUEST_ID_HEADER] = request_id
            auth_response.headers[TENANT_ID_HEADER] = context.tenant_id
            return auth_response

        settings = request.app.state.settings
        rate_limit_response = enforce_gateway_rate_limit(
            request=request,
            settings=settings,
            limiter=request.app.state.rate_limiter,
        )
        if rate_limit_response is not None:
            from app.infrastructure.audit import persist_gateway_audit

            persist_gateway_audit(
                request=request,
                session_factory=request.app.state.session_factory,
                status_code=rate_limit_response.status_code,
                error_message="rate limit exceeded",
            )
            context = get_gateway_request_context(request)
            rate_limit_response.headers[REQUEST_ID_HEADER] = request_id
            rate_limit_response.headers[TENANT_ID_HEADER] = context.tenant_id
            return rate_limit_response

        try:
            response = await call_next(request)
        except Exception as exc:
            from app.infrastructure.audit import persist_gateway_audit

            persist_gateway_audit(
                request=request,
                session_factory=request.app.state.session_factory,
                status_code=500,
                error_message=str(exc),
            )
            raise

        from app.infrastructure.audit import persist_gateway_audit

        persist_gateway_audit(
            request=request,
            session_factory=request.app.state.session_factory,
            status_code=response.status_code,
        )
        context = get_gateway_request_context(request)
        response.headers[REQUEST_ID_HEADER] = request_id
        response.headers[TENANT_ID_HEADER] = context.tenant_id
        apply_gateway_rate_limit_headers(
            response=response,
            request=request,
            settings=settings,
            limiter=request.app.state.rate_limiter,
        )
        return response


def resolve_tenant_id(request: Request) -> str:
    header_tenant_id = request.headers.get(TENANT_ID_HEADER)
    if header_tenant_id:
        return header_tenant_id

    query_tenant_id = request.query_params.get("tenant_id")
    if query_tenant_id:
        return query_tenant_id

    return DEFAULT_TENANT_ID


def get_gateway_request_context(request: Request) -> GatewayRequestContext:
    context = getattr(request.state, "gateway_context", None)
    if isinstance(context, GatewayRequestContext):
        return context
    return GatewayRequestContext(
        request_id=str(uuid4()),
        tenant_id=DEFAULT_TENANT_ID,
        started_perf_counter=perf_counter(),
    )


def authenticate_gateway_request(request: Request) -> Response | None:
    settings = ApiGatewaySettings()
    if not settings.auth_required:
        return None
    if not request.url.path.startswith("/gateway/"):
        return None
    if request.url.path in {"/gateway/services/health"}:
        return None

    if is_initial_api_key_bootstrap_request(request):
        return None

    api_key = request.headers.get(settings.api_key_header_name)
    if not api_key:
        return JSONResponse(
            status_code=401,
            content={"detail": "missing api key"},
        )

    db = request.app.state.session_factory()
    try:
        entity = ApiKeyRepository(db).get_active_by_hash(key_hash=hash_api_key(api_key))
        if entity is None:
            return JSONResponse(
                status_code=401,
                content={"detail": "invalid api key"},
            )
        if entity.expires_time is not None and entity.expires_time <= datetime.utcnow():
            return JSONResponse(
                status_code=401,
                content={"detail": "api key expired"},
            )

        context = get_gateway_request_context(request)
        requested_tenant_id = resolve_tenant_id(request)
        if requested_tenant_id not in {DEFAULT_TENANT_ID, entity.tenant_id}:
            return JSONResponse(
                status_code=403,
                content={"detail": "api key tenant mismatch"},
            )
        context.tenant_id = entity.tenant_id
        context.api_key_id = entity.id
        context.user_id = request.headers.get(settings.user_id_header_name)
        permission = derive_gateway_permission(request)
        if permission is not None and not api_key_scope_allows(
            scopes=entity.scopes,
            permission=permission,
        ):
            return JSONResponse(
                status_code=403,
                content={"detail": "api key scope denied", "permission": permission},
            )
        if settings.authz_required:
            if context.user_id is None:
                return JSONResponse(
                    status_code=401,
                    content={"detail": "missing user id"},
                )
            authz_response = check_auth_service_permission(
                settings=settings,
                tenant_id=entity.tenant_id,
                user_id=context.user_id,
                permission=permission or "gateway:access",
            )
            if authz_response is not None:
                return authz_response
        ApiKeyRepository(db).touch_last_used_time(api_key_id=entity.id)
    finally:
        db.close()

    return None


def is_initial_api_key_bootstrap_request(request: Request) -> bool:
    if request.method.upper() != "POST" or request.url.path != "/gateway/api-keys":
        return False

    db = request.app.state.session_factory()
    try:
        return not ApiKeyRepository(db).has_any()
    finally:
        db.close()


def derive_gateway_permission(request: Request) -> str | None:
    if not request.url.path.startswith("/gateway/"):
        return None
    path_parts = [part for part in request.url.path.split("/") if part]
    if len(path_parts) < 2:
        return "gateway:access"
    if path_parts[1] in {"services", "api-keys", "audits"}:
        resource = path_parts[1]
    else:
        resource = path_parts[1].replace("_", "-")
    action = "read" if request.method.upper() in {"GET", "HEAD", "OPTIONS"} else "write"
    return f"gateway:{resource}:{action}"


def api_key_scope_allows(*, scopes: str | None, permission: str) -> bool:
    if scopes is None or not scopes.strip():
        return True
    scope_values = parse_scope_values(scopes)
    if "*" in scope_values or permission in scope_values:
        return True
    resource_prefix = permission.rsplit(":", 1)[0]
    return f"{resource_prefix}:*" in scope_values


def parse_scope_values(scopes: str) -> set[str]:
    normalized = scopes.replace(",", " ").replace("\n", " ")
    return {item.strip() for item in normalized.split(" ") if item.strip()}


def check_auth_service_permission(
    *,
    settings: ApiGatewaySettings,
    tenant_id: str,
    user_id: str,
    permission: str,
) -> Response | None:
    try:
        with httpx.Client(timeout=settings.authz_timeout_seconds) as client:
            response = client.post(
                f"{settings.auth_service_url.rstrip('/')}/auth/permissions/check",
                headers=build_internal_service_headers(settings),
                json={
                    "tenant_id": tenant_id,
                    "user_id": user_id,
                    "permission": permission,
                },
            )
            response.raise_for_status()
            payload = response.json()
    except (httpx.HTTPError, ValueError) as exc:
        return JSONResponse(
            status_code=503,
            content={"detail": "auth service permission check failed", "error": str(exc)},
        )

    allowed = payload.get("allowed")
    if allowed is True:
        return None
    reason = payload.get("reason")
    return JSONResponse(
        status_code=403,
        content={
            "detail": "permission denied",
            "permission": permission,
            "reason": reason if isinstance(reason, str) else "denied",
        },
    )