app/core/access/subscription.py

"""Author RAG Chatbot SaaS — Subscription Validation.

This module is called by access_middleware on every chat request.
Validation order (fail-fast):
  1. HMAC signature check (tamper detection)
  2. Token expiry check
  3. Redis revocation blacklist check (instant revocation)
  4. DB record validation (is_revoked flag, author_id match)
  5. Token budget check (exhausted = soft block)
"""

import structlog
from redis.asyncio import Redis
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.access.token_crypto import validate_subscription_token_signature, hash_subscription_token
from app.exceptions.access import (
    AccessRevokedError,
    BudgetExhaustedError,
    InvalidSubscriptionTokenError,
    SubscriptionExpiredError,
    SubscriptionNotFoundError,
)
from app.models.user import User
from app.repositories.access_repo import AccessRepository
from app.repositories.user_repo import UserRepository

logger = structlog.get_logger(__name__)

REVOCATION_REDIS_PREFIX = "revoked_token:"
BUDGET_EXHAUSTED_REDIS_PREFIX = "budget_exhausted:"


async def validate_subscription_token(
    token: str,
    db: AsyncSession,
    redis: Redis,
) -> User:
    """Validate a subscription token through all security layers.

    Args:
        token: Raw subscription token from request header.
        db: Active database session.
        redis: Redis connection.

    Returns:
        Authenticated author User object.

    Raises:
        InvalidSubscriptionTokenError: HMAC invalid or malformed.
        SubscriptionExpiredError: Token past expiry.
        AccessRevokedError: Token in revocation blacklist or DB flagged.
        BudgetExhaustedError: Monthly token budget is at 100%.
        SubscriptionNotFoundError: No matching DB record found.
    """
    # Step 1 + 2: HMAC validation + expiry (both in signature check)
    try:
        payload = validate_subscription_token_signature(token)
    except Exception as e:
        logger.warning("Subscription token validation failed", error=str(e))
        raise InvalidSubscriptionTokenError()

    author_id: str = payload["author_id"]
    grant_id: str = payload["grant_id"]
    token_hash = hash_subscription_token(token)

    # Step 3: Redis revocation blacklist (fastest check — O(1))
    try:
        is_revoked = await redis.exists(f"{REVOCATION_REDIS_PREFIX}{token_hash}")
    except Exception as e:
        logger.warning("Redis revocation check skipped", error=str(e))
        is_revoked = 0
    if is_revoked:
        logger.warning("Revoked token used", author_id=author_id, grant_id=grant_id)
        raise AccessRevokedError()

    # Step 4: DB record validation
    access_repo = AccessRepository(db)
    access_record = await access_repo.get_by_grant_id(grant_id)

    if access_record is None:
        raise SubscriptionNotFoundError()

    if access_record.author_id != author_id:
        logger.error("Token author_id mismatch — possible token theft attempt", grant_id=grant_id)
        raise InvalidSubscriptionTokenError()

    if access_record.is_revoked:
        # Sync Redis blacklist (should already be there but safety net)
        await _add_to_revocation_blacklist(redis, token_hash, access_record.expires_at)
        raise AccessRevokedError(access_record.revoke_reason or "Access revoked")

    # Step 5: Token budget check
    try:
        budget_exhausted = await redis.exists(f"{BUDGET_EXHAUSTED_REDIS_PREFIX}{author_id}")
    except Exception as e:
        logger.warning("Redis budget check skipped", error=str(e))
        budget_exhausted = 0
    if budget_exhausted:
        raise BudgetExhaustedError()

    # All checks passed — return author
    user_repo = UserRepository(db)
    author = await user_repo.get_by_id(author_id)
    if not author or not author.is_active:
        raise SubscriptionNotFoundError()

    if not author.chatbot_is_active:
        raise AccessRevokedError("Account suspended")

    return author


async def revoke_token_in_redis(
    redis: Redis,
    token_hash: str,
    expires_at,
) -> None:
    """Add a token hash to the Redis revocation blacklist.

    TTL is set to the token's original expiry — after that, the token
    is naturally expired anyway and doesn't need to be in the blacklist.

    Args:
        redis: Redis connection.
        token_hash: SHA-256 hash of the token.
        expires_at: Token expiry datetime (sets Redis TTL).
    """
    await _add_to_revocation_blacklist(redis, token_hash, expires_at)


async def _add_to_revocation_blacklist(redis: Redis, token_hash: str, expires_at) -> None:
    """Internal: add token to Redis blacklist with appropriate TTL."""
    from datetime import datetime, timezone
    now = datetime.now(timezone.utc)
    exp = expires_at
    if exp.tzinfo is None:
        exp = exp.replace(tzinfo=timezone.utc)
    ttl_seconds = max(int((exp - now).total_seconds()), 1)
    await redis.setex(
        name=f"{REVOCATION_REDIS_PREFIX}{token_hash}",
        time=ttl_seconds,
        value="revoked",
    )
    logger.info("Token added to revocation blacklist", ttl_seconds=ttl_seconds)