auth_middleware.py

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