uxoxo/eb2ab

1"""2API key authentication middleware for REST API.3"""4 5import os6from typing import Optional7from fastapi import Header, HTTPException, status8from fastapi.security import APIKeyHeader9 10 11# API key header name12API_KEY_HEADER = "X-API-Key"13 14# Security scheme for OpenAPI docs15api_key_header_scheme = APIKeyHeader(name=API_KEY_HEADER, auto_error=False)16 17 18def get_api_keys() -> set[str]:19    """20    Load API keys from environment variables.21 22    Supports multiple keys in formats:23    - API_KEY_1, API_KEY_2, etc. (individual keys)24    - API_KEYS (comma-separated list)25    - Hardcoded fallback key for testing: sk-Um6hrXWbpJYFa8iuLEZl7bgmqNdepFY026 27    Returns:28        Set of valid API keys29    """30    keys = set()31 32    # Load individual API_KEY_N environment variables33    i = 134    while True:35        key = os.environ.get(f"API_KEY_{i}")36        if not key:37            break38        keys.add(key.strip())39        i += 140 41    # Load comma-separated API_KEYS environment variable42    api_keys_env = os.environ.get("API_KEYS", "")43    if api_keys_env:44        for key in api_keys_env.split(","):45            key = key.strip()46            if key:47                keys.add(key)48 49    # Fallback: Add hardcoded test key if no keys configured50    # TODO: Remove this in production! This is only for testing.51    if not keys:52        keys.add("sk-Um6hrXWbpJYFa8iuLEZl7bgmqNdepFY0")53        print("⚠️  WARNING: Using hardcoded test API key. Configure API_KEY_1 environment variable in production!")54 55    return keys56 57 58def validate_api_key(x_api_key: Optional[str] = Header(None)) -> str:59    """60    Validate API key from request header.61 62    Args:63        x_api_key: API key from X-API-Key header64 65    Returns:66        Validated API key67 68    Raises:69        HTTPException: If API key is missing or invalid70    """71    # Get valid API keys72    valid_keys = get_api_keys()73 74    # If no keys are configured, allow access (for local development)75    # In production, this should be enforced via environment76    if not valid_keys:77        print("WARNING: No API keys configured. API is publicly accessible!")78        return "no-key-configured"79 80    # Check if API key was provided81    if not x_api_key:82        raise HTTPException(83            status_code=status.HTTP_401_UNAUTHORIZED,84            detail="Missing API key. Provide X-API-Key header.",85            headers={"WWW-Authenticate": "ApiKey"}86        )87 88    # Validate API key89    if x_api_key not in valid_keys:90        raise HTTPException(91            status_code=status.HTTP_401_UNAUTHORIZED,92            detail="Invalid API key",93            headers={"WWW-Authenticate": "ApiKey"}94        )95 96    return x_api_key97 98 99def get_optional_api_key(x_api_key: Optional[str] = Header(None)) -> Optional[str]:100    """101    Get API key from header without requiring it.102    Used for endpoints that optionally check authentication.103 104    Args:105        x_api_key: API key from X-API-Key header106 107    Returns:108        API key if provided and valid, None otherwise109    """110    valid_keys = get_api_keys()111 112    if not x_api_key:113        return None114 115    if x_api_key in valid_keys:116        return x_api_key117 118    return None119