nat-engine 1__py3-none-any.whl

mannf/product/billing/tenant_auth.py

@@ -0,0 +1,233 @@
+# 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.
+
+"""Tenant authentication and quota checking helpers.
+
+These functions bridge the auth middleware with the Tenant/ApiKey/UsageRecord
+database models from PR 1. They are designed to gracefully degrade — if no
+DATABASE_URL is configured or no tenants exist in the DB, the system falls
+back to the legacy NAT_API_KEY env var behavior.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+async def resolve_tenant_from_key(api_key: str) -> Optional[Tuple[dict, dict]]:
+    """Look up an API key by its SHA-256 hash in the database.
+
+    Returns (tenant_dict, api_key_dict) if found and both are active,
+    or None if not found / inactive / no DB.
+
+    The tenant_dict includes: id, name, email, plan_tier,
+    monthly_scan_quota, monthly_security_scan_quota, is_active.
+
+    The api_key_dict includes: id, tenant_id, key_prefix, label, is_active.
+
+    Also updates the api_key's last_used_at timestamp (fire-and-forget,
+    don't block on it).
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine
+
+        if _build_engine() is None or _async_session_factory is None:
+            return None
+
+        from sqlalchemy import select
+        from mannf.product.models import ApiKey, Tenant
+
+        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
+
+        async with _async_session_factory() as session:
+            stmt = (
+                select(ApiKey, Tenant)
+                .join(Tenant, ApiKey.tenant_id == Tenant.id)
+                .where(ApiKey.key_hash == key_hash)
+                .where(ApiKey.is_active.is_(True))
+                .where(Tenant.is_active.is_(True))
+            )
+            result = await session.execute(stmt)
+            row = result.first()
+            if row is None:
+                return None
+
+            api_key_obj, tenant_obj = row
+
+            tenant_dict = {
+                "id": str(tenant_obj.id),
+                "name": tenant_obj.name,
+                "email": tenant_obj.email,
+                "plan_tier": tenant_obj.plan_tier,
+                "monthly_scan_quota": tenant_obj.monthly_scan_quota,
+                "monthly_security_scan_quota": tenant_obj.monthly_security_scan_quota,
+                "max_concurrent_scans": getattr(tenant_obj, "max_concurrent_scans", 1),
+                "is_active": tenant_obj.is_active,
+            }
+            api_key_dict = {
+                "id": str(api_key_obj.id),
+                "tenant_id": str(api_key_obj.tenant_id),
+                "key_prefix": api_key_obj.key_prefix,
+                "label": api_key_obj.label,
+                "is_active": api_key_obj.is_active,
+            }
+
+            # Fire-and-forget: update last_used_at
+            import asyncio
+
+            async def _update_last_used(ak_id: uuid.UUID) -> None:
+                try:
+                    async with _async_session_factory() as s:
+                        obj = await s.get(ApiKey, ak_id)
+                        if obj is not None:
+                            obj.last_used_at = datetime.now(timezone.utc)
+                            await s.commit()
+                except Exception as exc:  # noqa: BLE001
+                    logger.debug("last_used_at update failed: %s", exc)
+
+            asyncio.create_task(_update_last_used(api_key_obj.id))
+
+            return tenant_dict, api_key_dict
+
+    except ImportError:
+        return None
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("resolve_tenant_from_key failed: %s", exc)
+        return None
+
+
+async def get_monthly_usage(tenant_id: uuid.UUID, scan_type: str) -> int:
+    """Count the number of scans of the given type this tenant has used
+    in the current calendar month.
+
+    Uses the composite index ix_usage_records_tenant_period for fast lookup.
+    Returns 0 if no DB is configured.
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine
+
+        if _build_engine() is None or _async_session_factory is None:
+            return 0
+
+        from sqlalchemy import func, select
+        from mannf.product.models import UsageRecord
+
+        now = datetime.now(timezone.utc)
+        async with _async_session_factory() as session:
+            stmt = (
+                select(func.count())
+                .select_from(UsageRecord)
+                .where(UsageRecord.tenant_id == tenant_id)
+                .where(UsageRecord.scan_type == scan_type)
+                .where(UsageRecord.period_year == now.year)
+                .where(UsageRecord.period_month == now.month)
+            )
+            result = await session.execute(stmt)
+            return result.scalar() or 0
+
+    except ImportError:
+        return 0
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("get_monthly_usage failed: %s", exc)
+        return 0
+
+
+async def record_usage(
+    tenant_id: uuid.UUID,
+    scan_type: str,
+    scan_id: Optional[uuid.UUID] = None,
+) -> None:
+    """Insert a UsageRecord row for this tenant/scan.
+
+    Sets period_year and period_month from the current UTC date.
+    No-op if no DB is configured.
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine
+
+        if _build_engine() is None or _async_session_factory is None:
+            return
+
+        from mannf.product.models import UsageRecord
+
+        now = datetime.now(timezone.utc)
+        async with _async_session_factory() as session:
+            record = UsageRecord(
+                id=uuid.uuid4(),
+                tenant_id=tenant_id,
+                scan_type=scan_type,
+                scan_id=scan_id,
+                recorded_at=now,
+                period_year=now.year,
+                period_month=now.month,
+            )
+            session.add(record)
+            await session.commit()
+
+    except ImportError:
+        return
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("record_usage failed: %s", exc)
+
+
+async def get_daily_usage(tenant_id: uuid.UUID, scan_type: str) -> list:
+    """Return per-day usage counts for the current calendar month.
+
+    Each item is a dict with ``date`` (YYYY-MM-DD string) and ``count``
+    (number of scans of *scan_type* recorded on that date).
+    Returns an empty list when no DB is configured.
+    """
+    try:
+        from mannf.product.database import _async_session_factory, _build_engine
+
+        if _build_engine() is None or _async_session_factory is None:
+            return []
+
+        from sqlalchemy import cast, func, select
+        from sqlalchemy import Date
+        from mannf.product.models import UsageRecord
+
+        now = datetime.now(timezone.utc)
+        async with _async_session_factory() as session:
+            stmt = (
+                select(
+                    cast(UsageRecord.recorded_at, Date).label("day"),
+                    func.count().label("count"),
+                )
+                .where(UsageRecord.tenant_id == tenant_id)
+                .where(UsageRecord.scan_type == scan_type)
+                .where(UsageRecord.period_year == now.year)
+                .where(UsageRecord.period_month == now.month)
+                .group_by(cast(UsageRecord.recorded_at, Date))
+                .order_by(cast(UsageRecord.recorded_at, Date))
+            )
+            result = await session.execute(stmt)
+            rows = result.all()
+            return [{"date": str(row.day), "count": row.count} for row in rows]
+
+    except ImportError:
+        return []
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("get_daily_usage failed: %s", exc)
+        return []
+
+
+async def check_quota(
+    tenant_id: uuid.UUID,
+    scan_type: str,
+    tenant_quota: int,
+) -> Tuple[bool, int, int]:
+    """Check if the tenant is within their monthly quota.
+
+    Returns (allowed: bool, used: int, limit: int).
+    """
+    used = await get_monthly_usage(tenant_id, scan_type)
+    allowed = used < tenant_quota
+    return allowed, used, tenant_quota