nat-engine 1__py3-none-any.whl

mannf/product/auth/saml.py

@@ -0,0 +1,212 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+
+"""SAML 2.0 / SSO callback handling for NAT enterprise auth.
+
+Integrates with standard SAML 2.0 IdPs (Okta, Azure AD, Google Workspace).
+Maps SAML assertion attributes to NAT tenant + user records.
+
+Configuration (env vars):
+    NAT_SAML_ENABLED — set to "true" to enable SAML auth
+    NAT_SAML_TENANT_ATTRIBUTE — SAML attribute that carries the tenant ID (default: "tenant_id")
+    NAT_SAML_ROLE_ATTRIBUTE — SAML attribute that carries the role (default: "nat_role")
+    NAT_SAML_EMAIL_ATTRIBUTE — SAML attribute for email (default: "email")
+    NAT_SAML_NAME_ATTRIBUTE — SAML attribute for display name (default: "name")
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import os
+import uuid
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# SAML 2.0 namespaces
+_SAML_NS = {
+    "saml": "urn:oasis:names:tc:SAML:2.0:assertion",
+    "samlp": "urn:oasis:names:tc:SAML:2.0:protocol",
+}
+
+VALID_ROLES = {"admin", "operator", "viewer"}
+
+
+class SAMLNotConfiguredError(RuntimeError):
+    """Raised when SAML auth is requested but not configured."""
+
+
+class SAMLAssertionError(ValueError):
+    """Raised when a SAML assertion is invalid or cannot be processed."""
+
+
+def is_saml_enabled() -> bool:
+    """Return True if SAML auth is enabled via NAT_SAML_ENABLED env var."""
+    return os.environ.get("NAT_SAML_ENABLED", "").lower() in ("true", "1", "yes")
+
+
+def parse_saml_response(saml_response_b64: str) -> dict[str, Any]:
+    """Parse a base64-encoded SAML response and extract assertion attributes.
+
+    Args:
+        saml_response_b64: The base64-encoded SAMLResponse form parameter.
+
+    Returns:
+        Dict with keys: name_id, email, name, tenant_id, role, attributes (raw dict).
+
+    Raises:
+        SAMLAssertionError: on parse failure or missing required attributes.
+    """
+    try:
+        xml_bytes = base64.b64decode(saml_response_b64)
+        try:
+            import defusedxml.ElementTree as ET  # noqa: PLC0415
+        except ImportError:
+            from xml.etree import ElementTree as ET  # noqa: PLC0415
+        root = ET.fromstring(xml_bytes)
+    except Exception as exc:
+        raise SAMLAssertionError(f"Failed to parse SAML response: {exc}") from exc
+
+    # Extract NameID
+    name_id_el = root.find(".//{urn:oasis:names:tc:SAML:2.0:assertion}NameID")
+    if name_id_el is None or not name_id_el.text:
+        raise SAMLAssertionError("SAML assertion missing NameID")
+    name_id = name_id_el.text.strip()
+
+    # Extract all AttributeValue elements keyed by attribute name
+    attributes: dict[str, str] = {}
+    for attr_el in root.findall(".//{urn:oasis:names:tc:SAML:2.0:assertion}Attribute"):
+        attr_name = attr_el.get("Name", "")
+        val_el = attr_el.find("{urn:oasis:names:tc:SAML:2.0:assertion}AttributeValue")
+        if val_el is not None and val_el.text:
+            attributes[attr_name] = val_el.text.strip()
+
+    email_attr = os.environ.get("NAT_SAML_EMAIL_ATTRIBUTE", "email")
+    name_attr = os.environ.get("NAT_SAML_NAME_ATTRIBUTE", "name")
+    tenant_attr = os.environ.get("NAT_SAML_TENANT_ATTRIBUTE", "tenant_id")
+    role_attr = os.environ.get("NAT_SAML_ROLE_ATTRIBUTE", "nat_role")
+
+    email = attributes.get(email_attr) or name_id
+    name = attributes.get(name_attr) or email
+    tenant_id = attributes.get(tenant_attr)
+    role = attributes.get(role_attr, "viewer")
+
+    if role not in VALID_ROLES:
+        role = "viewer"
+
+    return {
+        "name_id": name_id,
+        "email": email,
+        "name": name,
+        "tenant_id": tenant_id,
+        "role": role,
+        "attributes": attributes,
+    }
+
+
+async def process_saml_callback(saml_response_b64: str) -> dict[str, Any]:
+    """Process a SAML callback — resolve or create user and return auth context.
+
+    1. Parse the SAML assertion.
+    2. Look up the tenant by ID from the assertion.
+    3. Find-or-create the user record (matched on sso_subject = NameID).
+    4. Ensure a UserRole record exists for this user with the asserted role.
+    5. Return an auth context dict suitable for setting on request.state.user.
+
+    Args:
+        saml_response_b64: Base64-encoded SAMLResponse.
+
+    Returns:
+        Auth context dict: user_id, email, name, role, tenant_id, sso_subject.
+
+    Raises:
+        SAMLNotConfiguredError: if SAML is not enabled.
+        SAMLAssertionError: on parse/validation failure.
+        ValueError: if the asserted tenant_id is unknown.
+    """
+    if not is_saml_enabled():
+        raise SAMLNotConfiguredError("SAML auth is not enabled (set NAT_SAML_ENABLED=true)")
+
+    attrs = parse_saml_response(saml_response_b64)
+    tenant_id_str = attrs.get("tenant_id")
+    if not tenant_id_str:
+        raise SAMLAssertionError("SAML assertion did not include a tenant_id attribute")
+
+    try:
+        tenant_uuid = uuid.UUID(tenant_id_str)
+    except ValueError as exc:
+        raise SAMLAssertionError(f"SAML tenant_id is not a valid UUID: {tenant_id_str!r}") from exc
+
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine  # noqa: PLC0415
+        from mannf.product.models import Tenant, User, UserRole  # noqa: PLC0415
+        from sqlalchemy import select  # noqa: PLC0415
+    except ImportError as exc:
+        raise RuntimeError("Database models not available") from exc
+
+    if _build_engine() is None or _async_session_factory is None:
+        raise RuntimeError("Database not configured — set DATABASE_URL and restart.")
+
+    async with _async_session_factory() as session:
+        # Verify tenant exists and is active
+        tenant = await session.get(Tenant, tenant_uuid)
+        if tenant is None or not tenant.is_active:
+            raise ValueError(f"Unknown or inactive tenant: {tenant_id_str!r}")
+
+        # Find or create user by sso_subject (NameID)
+        stmt = select(User).where(
+            User.sso_subject == attrs["name_id"],
+            User.tenant_id == tenant_uuid,
+        )
+        result = await session.execute(stmt)
+        user_obj = result.scalars().first()
+
+        if user_obj is None:
+            user_obj = User(
+                id=uuid.uuid4(),
+                tenant_id=tenant_uuid,
+                email=attrs["email"],
+                name=attrs["name"],
+                sso_subject=attrs["name_id"],
+                is_active=True,
+            )
+            session.add(user_obj)
+            await session.flush()
+
+        if not user_obj.is_active:
+            raise ValueError(f"User {attrs['email']!r} is disabled")
+
+        # Ensure tenant-level role exists
+        role_stmt = select(UserRole).where(
+            UserRole.user_id == user_obj.id,
+            UserRole.tenant_id == tenant_uuid,
+            UserRole.workspace_id.is_(None),
+        )
+        role_result = await session.execute(role_stmt)
+        role_obj = role_result.scalars().first()
+
+        if role_obj is None:
+            role_obj = UserRole(
+                id=uuid.uuid4(),
+                user_id=user_obj.id,
+                tenant_id=tenant_uuid,
+                workspace_id=None,
+                role=attrs["role"],
+            )
+            session.add(role_obj)
+        else:
+            role_obj.role = attrs["role"]
+
+        await session.commit()
+        await session.refresh(user_obj)
+
+    return {
+        "id": str(user_obj.id),
+        "email": user_obj.email,
+        "name": user_obj.name,
+        "role": attrs["role"],
+        "tenant_id": tenant_id_str,
+        "sso_subject": attrs["name_id"],
+    }

mannf/product/billing/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+
+"""NAT SaaS billing package."""

mannf/product/billing/audit.py

@@ -0,0 +1,160 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+
+"""Billing event audit log.
+
+Provides :func:`log_billing_event` which persists a :class:`BillingEvent`
+record to the database.  All operations degrade gracefully when
+``DATABASE_URL`` is not configured.
+
+Event types (not exhaustive):
+- ``tenant.created``
+- ``tenant.updated``
+- ``tenant.deactivated``
+- ``tenant.synced``
+- ``plan.upgraded``
+- ``plan.changed``
+- ``subscription.cancelled``
+- ``payment.failed``
+- ``payment.succeeded``
+- ``charge.refunded``
+- ``api_key.rotated``
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+
+async def log_billing_event(
+    event_type: str,
+    *,
+    tenant_id: Optional[str] = None,
+    source: str = "system",
+    details: Optional[dict[str, Any]] = None,
+) -> None:
+    """Persist a billing lifecycle event to the ``billing_events`` table.
+
+    Args:
+        event_type: Short string identifier, e.g. ``"tenant.created"``.
+        tenant_id: UUID string of the affected tenant (may be None for
+            system-level events).
+        source: Origin of the event — ``"stripe"``, ``"admin"``, or
+            ``"system"``.
+        details: Arbitrary JSON-serialisable dict with event context.
+
+    This function never raises — it logs and returns on any error to ensure
+    billing operations are not blocked by audit failures.
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine  # noqa: PLC0415
+
+        if _build_engine() is None or _async_session_factory is None:
+            logger.debug(
+                "log_billing_event: no database configured — skipping audit for %s",
+                event_type,
+            )
+            return
+
+        from mannf.product.models import BillingEvent  # noqa: PLC0415
+
+        tid: uuid.UUID | None = None
+        if tenant_id:
+            try:
+                tid = uuid.UUID(tenant_id)
+            except ValueError:
+                logger.warning(
+                    "log_billing_event: invalid tenant_id %r for event %s",
+                    tenant_id,
+                    event_type,
+                )
+
+        async with _async_session_factory() as session:
+            event = BillingEvent(
+                id=uuid.uuid4(),
+                tenant_id=tid,
+                event_type=event_type,
+                source=source,
+                details=details or {},
+            )
+            session.add(event)
+            await session.commit()
+
+        logger.debug("Logged billing event: %s (tenant=%s)", event_type, tenant_id)
+
+    except ImportError:
+        logger.debug("log_billing_event: SQLAlchemy not available — skipping audit")
+    except Exception as exc:  # noqa: BLE001
+        logger.error(
+            "log_billing_event failed for event %s (tenant=%s): %s",
+            event_type,
+            tenant_id,
+            exc,
+        )
+
+
+async def get_billing_events(
+    tenant_id: str,
+    *,
+    limit: int = 100,
+    event_type: Optional[str] = None,
+) -> list[dict[str, Any]]:
+    """Return billing events for a tenant, newest first.
+
+    Args:
+        tenant_id: UUID string of the tenant.
+        limit: Maximum number of events to return (capped at 500).
+        event_type: Optional filter — return only events of this type.
+
+    Returns:
+        List of event dicts with keys: ``id``, ``tenant_id``, ``event_type``,
+        ``source``, ``details``, ``created_at``.
+
+    Raises:
+        RuntimeError: if the database is not configured.
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine  # noqa: PLC0415
+
+        if _build_engine() is None or _async_session_factory is None:
+            raise RuntimeError("Database not configured — set DATABASE_URL and restart.")
+
+        from sqlalchemy import select  # noqa: PLC0415
+        from mannf.product.models import BillingEvent  # noqa: PLC0415
+
+        tid = uuid.UUID(tenant_id)
+        effective_limit = min(limit, 500)
+
+        async with _async_session_factory() as session:
+            stmt = (
+                select(BillingEvent)
+                .where(BillingEvent.tenant_id == tid)
+                .order_by(BillingEvent.created_at.desc())
+                .limit(effective_limit)
+            )
+            if event_type:
+                stmt = stmt.where(BillingEvent.event_type == event_type)
+            result = await session.execute(stmt)
+            rows = result.scalars().all()
+
+        return [
+            {
+                "id": str(row.id),
+                "tenant_id": str(row.tenant_id) if row.tenant_id else None,
+                "event_type": row.event_type,
+                "source": row.source,
+                "details": row.details or {},
+                "created_at": row.created_at.isoformat() if row.created_at else None,
+            }
+            for row in rows
+        ]
+
+    except ImportError as exc:
+        raise RuntimeError("Database models not available") from exc
+    except ValueError as exc:
+        raise ValueError(f"Invalid tenant_id: {tenant_id!r}") from exc

mannf/product/billing/feature_gates.py

@@ -0,0 +1,180 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+
+"""Feature-gating configuration for NAT SaaS plan tiers.
+
+Some features are only available on higher-tier plans.  ``check_feature()``
+is the single public entry point — it returns ``True`` when the tenant's
+plan allows the requested feature.
+
+``api_specs`` is an integer: ``-1`` means unlimited, any positive integer
+is the maximum number of API specs the tenant may upload.
+
+``modules`` lists the scan modules included with the plan.  Use
+``check_module()`` to gate module access by plan.
+"""
+
+from __future__ import annotations
+
+from typing import Any, TypedDict
+
+
+class FeatureConfig(TypedDict):
+    security_scans: bool
+    webhooks: bool
+    export_csv: bool
+    export_html: bool
+    export_pdf: bool        # PDF report export; Pro+ only
+    pdf_whitelabel: bool    # Strip NAT branding from PDF; Enterprise only
+    api_specs: int          # -1 = unlimited
+    modules: list[str]      # included scan modules
+    scheduled_scans: dict   # {"enabled": bool, "max_schedules": int}
+    llm_test_plans: int     # plans/month; -1 = unlimited, 0 = disabled
+    llm_natural_language: int  # NL test authoring queries/month; -1 = unlimited, 0 = disabled
+    llm_root_cause: int     # root cause analyses/scan; -1 = unlimited, 0 = disabled, 5 = top-5
+
+
+_ALL_MODULES: list[str] = ["functional", "visual", "accessibility", "performance"]
+
+PLAN_FEATURES: dict[str, FeatureConfig] = {
+    "free": {
+        "security_scans": False,
+        "webhooks": False,
+        "export_csv": True,
+        "export_html": True,
+        "export_pdf": False,
+        "pdf_whitelabel": False,
+        "api_specs": 1,
+        "modules": ["functional"],
+        "scheduled_scans": {"enabled": False, "max_schedules": 0},
+        "llm_test_plans": 0,
+        "llm_natural_language": 0,
+        "llm_root_cause": 0,
+    },
+    "pro": {
+        "security_scans": True,
+        "webhooks": False,
+        "export_csv": True,
+        "export_html": True,
+        "export_pdf": True,
+        "pdf_whitelabel": False,
+        "api_specs": 5,
+        "modules": _ALL_MODULES,
+        "scheduled_scans": {"enabled": True, "max_schedules": 3},
+        "llm_test_plans": 5,
+        "llm_natural_language": 10,
+        "llm_root_cause": 5,
+    },
+    "team": {
+        "security_scans": True,
+        "webhooks": True,
+        "export_csv": True,
+        "export_html": True,
+        "export_pdf": True,
+        "pdf_whitelabel": False,
+        "api_specs": -1,  # unlimited
+        "modules": _ALL_MODULES,
+        "scheduled_scans": {"enabled": True, "max_schedules": 10},
+        "llm_test_plans": 20,
+        "llm_natural_language": 50,
+        "llm_root_cause": -1,  # unlimited
+    },
+    "enterprise": {
+        "security_scans": True,
+        "webhooks": True,
+        "export_csv": True,
+        "export_html": True,
+        "export_pdf": True,
+        "pdf_whitelabel": True,
+        "api_specs": -1,  # unlimited
+        "modules": _ALL_MODULES,
+        "scheduled_scans": {"enabled": True, "max_schedules": 999_999},
+        "llm_test_plans": -1,  # unlimited
+        "llm_natural_language": -1,  # unlimited
+        "llm_root_cause": -1,  # unlimited
+    },
+}
+
+_DEFAULT_PLAN = "free"
+
+
+def get_features(plan_tier: str) -> FeatureConfig:
+    """Return the feature-gate config for *plan_tier*, defaulting to free."""
+    return PLAN_FEATURES.get(plan_tier, PLAN_FEATURES[_DEFAULT_PLAN])
+
+
+def check_feature(plan_tier: str, feature: str) -> bool:
+    """Return ``True`` if *plan_tier* allows *feature*.
+
+    For boolean features, returns the boolean value directly.
+    For integer features (e.g. ``api_specs``), returns ``True`` when the
+    value is non-zero (i.e. the feature is available at all).
+    """
+    config: dict[str, Any] = get_features(plan_tier)
+    value = config.get(feature)
+    if value is None:
+        return False
+    if isinstance(value, bool):
+        return value
+    # Integer feature: -1 = unlimited (always allowed), 0 = blocked
+    return int(value) != 0
+
+
+def check_module(plan: str, module: str) -> bool:
+    """Return ``True`` if *plan* includes the scan *module*.
+
+    Args:
+        plan: Plan tier name, e.g. ``"free"``, ``"pro"``, ``"team"``,
+              ``"enterprise"``.
+        module: Module name, e.g. ``"functional"``, ``"visual"``,
+                ``"accessibility"``, ``"performance"``.
+
+    Returns:
+        ``True`` if the plan includes the module, ``False`` otherwise.
+    """
+    return module in get_features(plan).get("modules", [])
+
+
+def check_module_entitlement(
+    tenant_plan: str,
+    module_name: str,
+    tenant_addons: list[str] | None = None,
+) -> tuple[bool, str | None]:
+    """Check if a tenant is entitled to use a scan module.
+
+    First checks whether the module is included in the tenant's base plan via
+    ``included_modules`` on :class:`~mannf.product.billing.plans.PlanTier`.
+    If not included in the base plan, checks whether the tenant has purchased
+    the module as an add-on (present in *tenant_addons*).
+
+    Args:
+        tenant_plan: The tenant's plan tier name (e.g. ``"free"``, ``"pro"``).
+        module_name: The scan module to check (e.g. ``"visual"``,
+                     ``"performance"``).
+        tenant_addons: List of module names the tenant has purchased as
+                       add-ons.  ``None`` is treated the same as an empty list.
+
+    Returns:
+        A ``(allowed, denial_reason)`` tuple.  If allowed, ``denial_reason``
+        is ``None``.  If denied, ``denial_reason`` is a human-readable string
+        explaining why access was refused.
+    """
+    from mannf.product.billing.plans import get_plan  # noqa: PLC0415
+
+    plan = get_plan(tenant_plan)
+
+    # 1. Module is included in the base plan subscription.
+    if module_name in plan.included_modules:
+        return True, None
+
+    # 2. Tenant purchased the module as an add-on.
+    addons = tenant_addons or []
+    if module_name in addons:
+        return True, None
+
+    # 3. Denied — neither plan-included nor add-on purchased.
+    return False, (
+        f"Module '{module_name}' is not included in the '{tenant_plan}' plan "
+        "and has not been purchased as an add-on."
+    )