infrasynth-backend-kit/infrasynth/tenancy/middleware.py

"""``TenantMiddleware`` — resolves and binds the request tenant (``TENANCY.md`` §3.2)."""

from __future__ import annotations

from typing import Any

from django.apps import apps

from infrasynth.shared.settings_utils import get_setting

from .context import reset_current_tenant, set_current_tenant
from .models import Tenant, TenantMembership

__all__ = ["TenantMiddleware"]

_DEFAULT_ALLOWLIST = (
    "/api/v1/auth/login/",
    "/api/v1/auth/refresh/",
    "/api/v1/auth/select-workspace/",
    "/api/v1/auth/altcha/",
    "/api/v1/auth/2fa/",
    "/api/v1/billing/webhook/",
    "/api/v1/schema/",
    "/api/v1/tenancy/accept-invitation/",
    "/healthz",
    "/readyz",
)


class TenantMiddleware:
    """Binds ``current_tenant`` from the token claim or a tenant-scoped API key.

    Runs after authentication. Rejects tenant endpoints with no resolved tenant
    (except the allowlist), and rejects a request whose membership was revoked
    immediately rather than waiting for token expiry.
    """

    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        if not get_setting("INFRASYNTH_TENANCY", "ENABLED", True):
            return self.get_response(request)

        tenant, error_code = self._resolve(request)
        token = set_current_tenant(tenant)
        request.tenant = tenant
        try:
            if error_code:
                return self._reject(error_code, request)
            if self._should_reject(request, tenant):
                return self._reject("AUTH_TENANT_REQUIRED", request)
            return self.get_response(request)
        finally:
            reset_current_tenant(token)

    # --- resolution ---------------------------------------------------------

    def _resolve(self, request) -> tuple[Tenant | None, str | None]:
        claim = get_setting("INFRASYNTH_TENANCY", "TENANT_CLAIM", "tenant")
        tenant_id = None

        auth = getattr(request, "auth", None)
        if auth is not None and hasattr(auth, "get"):
            tenant_id = auth.get(claim)

        if tenant_id is None:
            tenant_id = self._tenant_from_api_key(request)

        user = getattr(request, "user", None)
        is_authenticated = bool(user and getattr(user, "is_authenticated", False))

        if tenant_id is not None:
            tenant = Tenant.objects.filter(pk=tenant_id).first()
            if tenant is None:
                return None, "AUTH_TENANT_NOT_FOUND"
            if (
                is_authenticated
                and not TenantMembership.objects.filter(tenant=tenant, user=user, is_active=True).exists()
            ):
                return None, "AUTH_MEMBERSHIP_REVOKED"
            return tenant, None

        # Fallback: a user with exactly one active membership is auto-selected
        # (mirrors the login auto-select in TENANCY.md §3.1).
        if is_authenticated:
            memberships = list(
                TenantMembership.objects.filter(user=user, is_active=True, tenant__status__in=["active", "trialing"])
                .select_related("tenant")
                .order_by("joined_at")
            )
            if len(memberships) == 1:
                return memberships[0].tenant, None
        return None, None

    def _tenant_from_api_key(self, request) -> Any:
        raw_key = request.META.get("HTTP_X_API_KEY")
        if not raw_key or "." not in raw_key:
            return None
        prefix = raw_key.split(".", 1)[0]
        try:
            api_key_model = apps.get_model("infrasynth_security", "APIKey")
        except LookupError:
            return None
        api_key = api_key_model.all_objects.filter(prefix=prefix, is_active=True).first()
        if api_key is None:
            return None
        return api_key.tenant_id

    # --- rejection ----------------------------------------------------------

    def _should_reject(self, request, tenant: Tenant | None) -> bool:
        if tenant is not None:
            return False
        if not get_setting("INFRASYNTH_TENANCY", "REQUIRE_TENANT_BY_DEFAULT", True):
            return False
        path = getattr(request, "path", "") or ""
        if not path.startswith("/api/"):
            return False
        allowlist = get_setting("INFRASYNTH_TENANCY", "TENANT_ALLOWLIST_PATHS", None) or _DEFAULT_ALLOWLIST
        if any(path.startswith(prefix) for prefix in allowlist):
            return False
        user = getattr(request, "user", None)
        is_authenticated = bool(user and getattr(user, "is_authenticated", False))
        has_api_key = bool(request.META.get("HTTP_X_API_KEY"))
        # Unauthenticated requests are left to the auth classes (401), not 403'd here.
        return is_authenticated or has_api_key

    def _reject(self, code: str, request):
        from infrasynth.api.exceptions import error_response

        messages = {
            "AUTH_TENANT_REQUIRED": "A workspace context is required for this endpoint.",
            "AUTH_MEMBERSHIP_REVOKED": "Your membership in this workspace is no longer active.",
            "AUTH_TENANT_NOT_FOUND": "The workspace could not be resolved.",
        }
        return error_response(code, messages.get(code, "Tenant resolution failed."), status_code=403, request=request)