""" auth_middleware.py - Authentication Middleware for HF Space Extracts user ID from headers sent by the main proxy and validates session tokens """ from functools import wraps from flask import request, jsonify, g, session import logging from datetime import datetime import hashlib import hmac import os logger = logging.getLogger(__name__) # =========================================== # CONFIGURATION # =========================================== # Secret key for validating admin tokens (should match the one in admin proxy) ADMIN_SECRET_KEY = os.environ.get('ADMIN_SECRET_KEY', 'admin-proxy-secret-key-change-me') # Session validation cache (simple in-memory, could be replaced with Redis) _session_cache = {} # {session_token: {'user_id': user_id, 'expires_at': timestamp}} _session_cache_ttl = 3600 # 1 hour # =========================================== # HELPER FUNCTIONS # =========================================== def is_admin_request(): """Check if request is from admin proxy (based on header)""" admin_header = request.headers.get('X-Admin-Authenticated', '').lower() admin_username = request.headers.get('X-Admin-Username', '') # For now, just check if admin header is present # In production, you should validate a signature return admin_header == 'true' and admin_username != '' def validate_admin_request(): """Validate that the request is from the admin proxy""" admin_header = request.headers.get('X-Admin-Authenticated', '').lower() if admin_header == 'true': # This is a request from admin proxy # You could also check a signature or token here for additional security return True return False def get_user_from_session_token(session_token): """Get user ID from session token (cached)""" # Check cache first if session_token in _session_cache: cache_entry = _session_cache[session_token] if datetime.now().timestamp() < cache_entry['expires_at']: return cache_entry['user_id'] else: # Expired, remove from cache del _session_cache[session_token] # Not in cache or expired - would need to validate with database # For now, we'll assume the token is valid if it's passed # In production, you should validate against a database table of active sessions # Since we don't have a database yet, we'll just trust the token # and let the routes handle validation via X-User-ID header return None def cache_session_token(session_token, user_id, ttl=3600): """Cache a session token for faster validation""" _session_cache[session_token] = { 'user_id': user_id, 'expires_at': datetime.now().timestamp() + ttl } # =========================================== # MAIN AUTHENTICATION MIDDLEWARE # =========================================== def require_auth(f): """ Decorator to require authentication for routes. Extracts user_id from headers and stores in Flask's g object. """ @wraps(f) def decorated_function(*args, **kwargs): # Check if this is an admin request if validate_admin_request(): # Admin requests bypass normal auth g.is_admin = True g.user_id = None g.user_email = None g.user_username = 'admin' logger.info(f"Admin request to {request.path}") return f(*args, **kwargs) # Regular user authentication user_id = request.headers.get('X-User-ID') session_token = request.headers.get('X-Session-Token') user_email = request.headers.get('X-User-Email') user_username = request.headers.get('X-User-Username') authenticated = request.headers.get('X-Authenticated', '').lower() # Validate authentication if not user_id or authenticated != 'true': logger.warning(f"Unauthenticated request to {request.path} from {request.remote_addr}") return jsonify({'error': 'Authentication required'}), 401 # Optional: Validate session token against database # For now, we trust the proxy's authentication # In production, you should validate the token # Store user info in Flask's g object for use in routes g.user_id = user_id g.user_email = user_email g.user_username = user_username g.is_admin = False g.session_token = session_token logger.debug(f"Authenticated request: user_id={user_id}, path={request.path}") return f(*args, **kwargs) return decorated_function def optional_auth(f): """ Decorator for routes that can work with or without authentication. Used for public API endpoints like health checks. """ @wraps(f) def decorated_function(*args, **kwargs): user_id = request.headers.get('X-User-ID') authenticated = request.headers.get('X-Authenticated', '').lower() if user_id and authenticated == 'true': g.user_id = user_id g.user_email = request.headers.get('X-User-Email') g.user_username = request.headers.get('X-User-Username') g.is_authenticated = True else: g.user_id = None g.user_email = None g.user_username = None g.is_authenticated = False g.is_admin = validate_admin_request() return f(*args, **kwargs) return decorated_function def admin_required(f): """ Decorator to require admin access. Must be used after require_auth or optional_auth. """ @wraps(f) def decorated_function(*args, **kwargs): if not hasattr(g, 'is_admin') or not g.is_admin: logger.warning(f"Non-admin access attempt to {request.path}") return jsonify({'error': 'Admin access required'}), 403 return f(*args, **kwargs) return decorated_function def get_current_user_id(): """Get the current authenticated user's ID""" return getattr(g, 'user_id', None) def get_current_user_email(): """Get the current authenticated user's email""" return getattr(g, 'user_email', None) def get_current_username(): """Get the current authenticated user's username""" return getattr(g, 'user_username', None) def is_admin(): """Check if current request is from admin""" return getattr(g, 'is_admin', False) # =========================================== # SESSION MANAGEMENT (for user API calls) # =========================================== def create_user_session(user_id, session_token): """Create a session cache entry for a user""" cache_session_token(session_token, user_id) def invalidate_user_session(session_token): """Invalidate a user's session token""" if session_token in _session_cache: del _session_cache[session_token] return True return False def validate_session(session_token): """Validate a session token and return user_id if valid""" # Check cache if session_token in _session_cache: cache_entry = _session_cache[session_token] if datetime.now().timestamp() < cache_entry['expires_at']: return cache_entry['user_id'] return None # =========================================== # CLEANUP THREAD (optional) # =========================================== import threading import time def cleanup_expired_sessions(): """Background thread to clean up expired session cache entries""" while True: try: now = datetime.now().timestamp() expired = [] for token, entry in _session_cache.items(): if now >= entry['expires_at']: expired.append(token) for token in expired: del _session_cache[token] if expired: logger.debug(f"Cleaned up {len(expired)} expired session cache entries") except Exception as e: logger.error(f"Session cleanup error: {e}") time.sleep(300) # Run every 5 minutes # Start cleanup thread _cleanup_thread = threading.Thread(target=cleanup_expired_sessions, daemon=True) _cleanup_thread.start()