123 lines
5.4 KiB
Python
123 lines
5.4 KiB
Python
"""
|
|
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 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)
|
|
if user_account:
|
|
request.account = user_account
|
|
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:
|
|
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)
|
|
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
|
|
|