"""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)