igny8/backend/igny8_core/auth/middleware.py

"""
Multi-Account Middleware
Extracts account from JWT token and injects into request context
"""
from django.utils.deprecation import MiddlewareMixin
from django.http import JsonResponse
from django.contrib.auth import logout
from rest_framework import status

try:
    import jwt
    JWT_AVAILABLE = True
except ImportError:
    JWT_AVAILABLE = False

from django.conf import settings


class AccountContextMiddleware(MiddlewareMixin):
    """
    Middleware that extracts account information from JWT token
    and adds it to request context for account isolation.
    """
    
    def process_request(self, request):
        """Extract account from JWT token in Authorization header or session."""
        # Skip for admin and auth endpoints
        if request.path.startswith('/admin/') or request.path.startswith('/api/v1/auth/'):
            return None
        
        # First, try to get user from Django session (cookie-based auth)
        # This handles cases where frontend uses credentials: 'include' with session cookies
        if hasattr(request, 'user') and request.user and request.user.is_authenticated:
            # User is authenticated via session - refresh from DB to get latest account/plan data
            # This ensures changes to account/plan are reflected immediately without re-login
            try:
                from .models import User as UserModel
                # Refresh user from DB with account and plan relationships to get latest data
                # This is important so account/plan changes are reflected immediately
                user = UserModel.objects.select_related('account', 'account__plan').get(id=request.user.id)
                # Update request.user with fresh data
                request.user = user
                # Get account from refreshed user
                user_account = getattr(user, 'account', None)
                validation_error = self._validate_account_and_plan(request, user)
                if validation_error:
                    return validation_error
                request.account = getattr(user, 'account', None)
                return None
            except (AttributeError, UserModel.DoesNotExist, Exception):
                # If refresh fails, fallback to cached account
                try:
                    user_account = getattr(request.user, 'account', None)
                    if user_account:
                        validation_error = self._validate_account_and_plan(request, request.user)
                        if validation_error:
                            return validation_error
                        request.account = user_account
                        return None
                except (AttributeError, Exception):
                    pass
                # If account access fails (e.g., column mismatch), set to None
                request.account = None
                return None
        
        # Get token from Authorization header (JWT auth - for future implementation)
        auth_header = request.META.get('HTTP_AUTHORIZATION', '')
        if not auth_header.startswith('Bearer '):
            # No JWT token - if session auth didn't work, set account to None
            # But don't set request.user to None - it might be set by Django's auth middleware
            if not hasattr(request, 'account'):
                request.account = None
            return None
        
        token = auth_header.split(' ')[1] if len(auth_header.split(' ')) > 1 else None
        if not token:
            if not hasattr(request, 'account'):
                request.account = None
            return None
        
        try:
            if not JWT_AVAILABLE:
                # JWT library not installed yet - skip for now
                request.account = None
                return None
            
            # Decode JWT token with signature verification
            # Use JWT_SECRET_KEY from settings (falls back to SECRET_KEY if not set)
            jwt_secret = getattr(settings, 'JWT_SECRET_KEY', getattr(settings, 'SECRET_KEY', None))
            if not jwt_secret:
                raise ValueError("JWT_SECRET_KEY or SECRET_KEY must be set in settings")
            
            decoded = jwt.decode(token, jwt_secret, algorithms=[getattr(settings, 'JWT_ALGORITHM', 'HS256')])
            
            # Extract user and account info from token
            user_id = decoded.get('user_id')
            account_id = decoded.get('account_id')
            
            if user_id:
                from .models import User, Account
                try:
                    # Get user from DB (but don't set request.user - let DRF authentication handle that)
                    # Only set request.account for account context
                    user = User.objects.select_related('account', 'account__plan').get(id=user_id)
                    validation_error = self._validate_account_and_plan(request, user)
                    if validation_error:
                        return validation_error
                    if account_id:
                        # Verify account still exists
                        try:
                            account = Account.objects.get(id=account_id)
                            request.account = account
                        except Account.DoesNotExist:
                            # Account from token doesn't exist - don't fallback, set to None
                            request.account = None
                    else:
                        # No account_id in token - set to None (don't fallback to user.account)
                        request.account = None
                except (User.DoesNotExist, Account.DoesNotExist):
                    request.account = None
            else:
                request.account = None
                
        except jwt.InvalidTokenError:
            request.account = None
        except Exception:
            # Fail silently for now - allow unauthenticated access
            request.account = None
        
        return None
    
    def _validate_account_and_plan(self, request, user):
        """
        Ensure the authenticated user has an account and an active plan.
        If not, logout the user (for session auth) and block the request.
        """
        try:
            account = getattr(user, 'account', None)
        except Exception:
            account = None
        
        if not account:
            return self._deny_request(
                request,
                error='Account not configured for this user. Please contact support.',
                status_code=status.HTTP_403_FORBIDDEN,
            )
        
        plan = getattr(account, 'plan', None)
        if plan is None or getattr(plan, 'is_active', False) is False:
            return self._deny_request(
                request,
                error='Active subscription required. Visit igny8.com/pricing to subscribe.',
                status_code=status.HTTP_402_PAYMENT_REQUIRED,
            )
        
        return None
    
    def _deny_request(self, request, error, status_code):
        """Logout session users (if any) and return a consistent JSON error."""
        try:
            if hasattr(request, 'user') and request.user and request.user.is_authenticated:
                logout(request)
        except Exception:
            pass
        
        return JsonResponse(
            {
                'success': False,
                'error': error,
            },
            status=status_code,
        )