""" Quart-compatible JWT Authentication Replacement for Flask-JWT-Extended to work with Quart ASGI applications. Provides jwt_required decorator and token management functions. """ import os import jwt import functools import json from datetime import datetime, timedelta, timezone from typing import Optional, Dict, Any from quart import request, g, current_app, jsonify, Response # JWT Configuration — reads SECRET_KEY from env, crashes if missing/weak _raw_secret = os.environ.get('SECRET_KEY', '') _weak_defaults = {'dev-secret-key', 'your-secret-key-for-sessions-and-tokens', '', 'change-me'} if not _raw_secret or _raw_secret in _weak_defaults: raise RuntimeError( "SECRET_KEY environment variable is not set or uses a weak default. " "Set a strong random value in backend/.env before starting the server." ) JWT_SECRET_KEY = _raw_secret JWT_ACCESS_TOKEN_EXPIRES = timedelta(hours=24) JWT_ALGORITHM = 'HS256' class QuartJWTError(Exception): """Base exception for JWT errors in Quart.""" pass def create_access_token(identity: str, expires_delta: Optional[timedelta] = None) -> str: """ Create a JWT access token. Args: identity: User identifier (usually user ID) expires_delta: Optional expiration time override Returns: JWT token string """ if expires_delta: expire = datetime.now(timezone.utc) + expires_delta else: expire = datetime.now(timezone.utc) + JWT_ACCESS_TOKEN_EXPIRES payload = { 'sub': identity, # Subject (user ID) 'exp': expire, 'iat': datetime.now(timezone.utc), 'type': 'access' } return jwt.encode(payload, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM) def decode_token(token: str) -> Dict[str, Any]: """ Decode and validate a JWT token. Args: token: JWT token string Returns: Decoded token payload Raises: QuartJWTError: If token is invalid """ try: payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM]) return payload except jwt.ExpiredSignatureError: raise QuartJWTError("Token has expired") except jwt.InvalidTokenError as e: raise QuartJWTError(f"Invalid token: {str(e)}") def get_jwt_identity() -> Optional[str]: """ Get the identity (user ID) from the current JWT token. Returns: User ID from token, or None if no valid token """ try: return getattr(g, 'current_user_id', None) except Exception: return None def jwt_required(optional: bool = False): """ Decorator to require valid JWT token for route access. Args: optional: If True, allow access without token (but still decode if present) """ def decorator(func): @functools.wraps(func) async def wrapper(*args, **kwargs): try: # Get token from Authorization header auth_header = request.headers.get('Authorization') token = None if auth_header: # Expected format: "Bearer " parts = auth_header.split() if len(parts) == 2 and parts[0].lower() == 'bearer': token = parts[1] if not token: if optional: # No token provided but optional - allow access g.current_user_id = None result = await func(*args, **kwargs) # Handle tuple returns if isinstance(result, tuple) and len(result) == 2: response, status_code = result if hasattr(response, 'status_code'): response.status_code = status_code return response else: from quart import make_response return make_response(response, status_code) else: return result else: # Token required but not provided return Response( json.dumps({'error': 'Missing authorization token'}), status=401, mimetype="application/json" ) # Validate token try: payload = decode_token(token) user_id = payload.get('sub') if not user_id: raise QuartJWTError("No user ID in token") # Store user ID in request context g.current_user_id = user_id # Propagate user_id into the LLM usage ContextVar for this request. # Each Quart request runs in its own asyncio Task, so setting the ContextVar # here is request-scoped. Child tasks (create_task) and thread submissions # (run_coroutine_threadsafe) inherit this context automatically. try: from app.services.llm_usage_context import _ctx, current_context from dataclasses import replace as _dc_replace _ctx.set(_dc_replace(current_context(), user_id=user_id)) except Exception: pass # Non-fatal — telemetry only # Call the actual route function and handle tuple returns result = await func(*args, **kwargs) # Handle tuple returns (response, status_code) if isinstance(result, tuple) and len(result) == 2: response, status_code = result if hasattr(response, 'status_code'): response.status_code = status_code return response else: # Use make_response for non-response objects from quart import make_response return make_response(response, status_code) else: return result except QuartJWTError as e: if optional: # Invalid token but optional - allow access without user ID g.current_user_id = None result = await func(*args, **kwargs) # Handle tuple returns here too if isinstance(result, tuple) and len(result) == 2: response, status_code = result if hasattr(response, 'status_code'): response.status_code = status_code return response else: from quart import make_response return make_response(response, status_code) else: return result else: # Invalid token and required return Response( json.dumps({'error': f'Invalid token: {str(e)}'}), status=401, mimetype="application/json" ) except Exception as e: current_app.logger.error(f"JWT validation error: {e}") if optional: g.current_user_id = None result = await func(*args, **kwargs) # Handle tuple returns here too if isinstance(result, tuple) and len(result) == 2: response, status_code = result if hasattr(response, 'status_code'): response.status_code = status_code return response else: from quart import make_response return make_response(response, status_code) else: return result else: return Response( json.dumps({'error': 'Authentication error'}), status=500, mimetype="application/json" ) return wrapper return decorator # For backward compatibility, provide the same function names as Flask-JWT-Extended def get_current_user(): """Get current user ID (alias for get_jwt_identity).""" return get_jwt_identity()