webscout 8.3__py3-none-any.whl → 8.3.2__py3-none-any.whl

webscout/auth/middleware.py

@@ -0,0 +1,248 @@
+# webscout/auth/middleware.py
+
+import secrets
+from datetime import datetime
+from typing import Optional, Tuple
+import logging
+
+from fastapi import Request, HTTPException, status
+from fastapi.security import APIKeyHeader
+from fastapi.responses import JSONResponse
+
+from .api_key_manager import APIKeyManager
+from .rate_limiter import RateLimiter
+from .models import APIKey
+from .schemas import ErrorResponse, RateLimitStatus
+
+logger = logging.getLogger(__name__)
+
+
+class AuthMiddleware:
+    """Authentication and rate limiting middleware."""
+    
+    def __init__(self, api_key_manager: APIKeyManager, rate_limiter: RateLimiter, auth_required: bool = True, rate_limit_enabled: bool = True):
+        self.api_key_manager = api_key_manager
+        self.rate_limiter = rate_limiter
+        self.api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
+        self.auth_required = auth_required
+        self.rate_limit_enabled = rate_limit_enabled
+
+        # Paths that don't require authentication
+        self.public_paths = {
+            "/",
+            "/docs",
+            "/redoc",
+            "/openapi.json",
+            "/health",
+            "/v1/auth/generate-key"  # API key generation endpoint
+        }
+
+        # Paths that require authentication (when auth is enabled)
+        self.protected_path_prefixes = [
+            "/v1/chat/completions",
+            "/v1/images/generations",
+            "/v1/models",
+            "/v1/TTI/models"
+        ]
+    
+    def is_protected_path(self, path: str) -> bool:
+        """Check if a path requires authentication."""
+        if not self.auth_required:
+            return False  # No authentication required in no-auth mode
+
+        if path in self.public_paths:
+            return False
+
+        return any(path.startswith(prefix) for prefix in self.protected_path_prefixes)
+    
+    def extract_api_key(self, authorization_header: Optional[str]) -> Optional[str]:
+        """Extract API key from Authorization header."""
+        if not authorization_header:
+            return None
+        
+        # Support both "Bearer <key>" and direct key formats
+        if authorization_header.startswith("Bearer "):
+            return authorization_header[7:]
+        elif authorization_header.startswith("ws_"):
+            return authorization_header
+        else:
+            return None
+    
+    async def authenticate_request(self, request: Request) -> Tuple[bool, Optional[APIKey], Optional[dict]]:
+        """
+        Authenticate a request.
+        
+        Returns:
+            Tuple of (is_authenticated, api_key_object, error_response)
+        """
+        path = request.url.path
+        
+        # Check if path requires authentication
+        if not self.is_protected_path(path):
+            return True, None, None
+        
+        # Get authorization header
+        auth_header = request.headers.get("authorization")
+        api_key = self.extract_api_key(auth_header)
+        
+        if not api_key:
+            error_response = {
+                "error": "API key required",
+                "code": "missing_api_key",
+                "details": {
+                    "message": "Please provide a valid API key in the Authorization header",
+                    "format": "Authorization: Bearer <your-api-key>"
+                }
+            }
+            return False, None, error_response
+        
+        # Validate API key
+        is_valid, api_key_obj, error_msg = await self.api_key_manager.validate_api_key(api_key)
+        
+        if not is_valid:
+            error_response = {
+                "error": error_msg or "Invalid API key",
+                "code": "invalid_api_key",
+                "details": {
+                    "message": "The provided API key is invalid, expired, or inactive"
+                }
+            }
+            return False, None, error_response
+        
+        return True, api_key_obj, None
+    
+    async def check_rate_limit(self, api_key: Optional[APIKey], client_ip: str = "unknown") -> Tuple[bool, dict]:
+        """
+        Check rate limit for an API key or IP address.
+
+        Returns:
+            Tuple of (is_allowed, rate_limit_info)
+        """
+        if not self.rate_limit_enabled:
+            # Rate limiting disabled
+            return True, {
+                "allowed": True,
+                "limit": 999999,
+                "remaining": 999999,
+                "reset_at": None,
+                "retry_after": None
+            }
+
+        if api_key:
+            return await self.rate_limiter.check_rate_limit(api_key)
+        else:
+            # No-auth mode: use IP-based rate limiting
+            return await self.rate_limiter.check_ip_rate_limit(client_ip)
+    
+    async def process_request(self, request: Request) -> Tuple[bool, Optional[dict], Optional[dict]]:
+        """
+        Process a request through authentication and rate limiting.
+        
+        Returns:
+            Tuple of (is_allowed, error_response, rate_limit_headers)
+        """
+        # Authenticate request
+        is_authenticated, api_key_obj, auth_error = await self.authenticate_request(request)
+        
+        if not is_authenticated:
+            return False, auth_error, None
+        
+        # If no API key (public endpoint or no-auth mode), handle rate limiting
+        if not api_key_obj:
+            if not self.auth_required:
+                # No-auth mode: still check rate limiting by IP
+                client_ip = request.client.host if request.client else "unknown"
+                is_allowed, rate_limit_info = await self.check_rate_limit(None, client_ip)
+
+                if not is_allowed:
+                    error_response = {
+                        "error": "Rate limit exceeded",
+                        "code": "rate_limit_exceeded",
+                        "details": {
+                            "message": f"Rate limit of {rate_limit_info['limit']} requests per minute exceeded",
+                            "retry_after": rate_limit_info["retry_after"],
+                            "reset_at": rate_limit_info["reset_at"].isoformat() if rate_limit_info["reset_at"] else None
+                        }
+                    }
+                    rate_limit_headers = {
+                        "X-RateLimit-Limit": str(rate_limit_info["limit"]),
+                        "X-RateLimit-Remaining": str(rate_limit_info["remaining"]),
+                        "Retry-After": str(rate_limit_info["retry_after"])
+                    }
+                    if rate_limit_info["reset_at"]:
+                        rate_limit_headers["X-RateLimit-Reset"] = rate_limit_info["reset_at"].isoformat()
+                    return False, error_response, rate_limit_headers
+
+                # Store rate limit info for no-auth mode
+                request.state.rate_limit_info = rate_limit_info
+                rate_limit_headers = {
+                    "X-RateLimit-Limit": str(rate_limit_info["limit"]),
+                    "X-RateLimit-Remaining": str(rate_limit_info["remaining"])
+                }
+                if rate_limit_info["reset_at"]:
+                    rate_limit_headers["X-RateLimit-Reset"] = rate_limit_info["reset_at"].isoformat()
+                return True, None, rate_limit_headers
+            else:
+                return True, None, None
+
+        # Check rate limit for authenticated requests
+        client_ip = request.client.host if request.client else "unknown"
+        is_allowed, rate_limit_info = await self.check_rate_limit(api_key_obj, client_ip)
+        
+        # Prepare rate limit headers
+        rate_limit_headers = {
+            "X-RateLimit-Limit": str(rate_limit_info["limit"]),
+            "X-RateLimit-Remaining": str(rate_limit_info["remaining"]),
+            "X-RateLimit-Reset": rate_limit_info["reset_at"].isoformat()
+        }
+        
+        if not is_allowed:
+            error_response = {
+                "error": "Rate limit exceeded",
+                "code": "rate_limit_exceeded",
+                "details": {
+                    "message": f"Rate limit of {rate_limit_info['limit']} requests per minute exceeded",
+                    "retry_after": rate_limit_info["retry_after"],
+                    "reset_at": rate_limit_info["reset_at"].isoformat()
+                }
+            }
+            rate_limit_headers["Retry-After"] = str(rate_limit_info["retry_after"])
+            return False, error_response, rate_limit_headers
+        
+        # Store API key info in request state for use in endpoints
+        request.state.api_key = api_key_obj
+        request.state.rate_limit_info = rate_limit_info
+        
+        return True, None, rate_limit_headers
+    
+    def create_error_response(self, error_data: dict, status_code: int = 401, headers: Optional[dict] = None) -> JSONResponse:
+        """Create a standardized error response."""
+        response = JSONResponse(
+            status_code=status_code,
+            content=error_data
+        )
+        
+        if headers:
+            for key, value in headers.items():
+                response.headers[key] = value
+        
+        return response
+    
+    async def __call__(self, request: Request, call_next):
+        """Middleware callable for FastAPI."""
+        # Process request
+        is_allowed, error_response, rate_limit_headers = await self.process_request(request)
+        
+        if not is_allowed:
+            status_code = 429 if error_response.get("code") == "rate_limit_exceeded" else 401
+            return self.create_error_response(error_response, status_code, rate_limit_headers)
+        
+        # Continue with request
+        response = await call_next(request)
+        
+        # Add rate limit headers to response
+        if rate_limit_headers:
+            for key, value in rate_limit_headers.items():
+                response.headers[key] = value
+        
+        return response

webscout/auth/models.py

@@ -0,0 +1,130 @@
+# webscout/auth/models.py
+
+from datetime import datetime, timezone
+from typing import Optional, Dict, Any, List
+from dataclasses import dataclass, field
+import uuid
+import json
+
+
+@dataclass
+class User:
+    """User model for authentication system."""
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    username: str = ""
+    telegram_id: int = field(default_factory=lambda: 0)  # Required Telegram ID as number only
+
+    def validate_telegram_id(self) -> None:
+        """Ensure telegram_id is an integer."""
+        if not isinstance(self.telegram_id, int):
+            raise ValueError("telegram_id must be an integer.")
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    is_active: bool = True
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert user to dictionary for storage."""
+        return {
+            "id": self.id,
+            "username": self.username,
+            "telegram_id": self.telegram_id,
+            "created_at": self.created_at.isoformat(),
+            "updated_at": self.updated_at.isoformat(),
+            "is_active": self.is_active,
+            "metadata": self.metadata
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "User":
+        """Create user from dictionary."""
+        return cls(
+            id=data.get("id", str(uuid.uuid4())),
+            username=data.get("username", ""),
+            telegram_id=int(data.get("telegram_id", 0)),
+            created_at=datetime.fromisoformat(data.get("created_at", datetime.now(timezone.utc).isoformat())),
+            updated_at=datetime.fromisoformat(data.get("updated_at", datetime.now(timezone.utc).isoformat())),
+            is_active=data.get("is_active", True),
+            metadata=data.get("metadata", {})
+        )
+
+
+@dataclass
+class APIKey:
+    """API Key model for authentication system."""
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    key: str = ""
+    user_id: str = ""
+    name: Optional[str] = None
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    last_used_at: Optional[datetime] = None
+    expires_at: Optional[datetime] = None
+    is_active: bool = True
+    rate_limit: int = 10  # requests per minute
+    usage_count: int = 0
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert API key to dictionary for storage."""
+        return {
+            "id": self.id,
+            "key": self.key,
+            "user_id": self.user_id,
+            "name": self.name,
+            "created_at": self.created_at.isoformat(),
+            "last_used_at": self.last_used_at.isoformat() if self.last_used_at else None,
+            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
+            "is_active": self.is_active,
+            "rate_limit": self.rate_limit,
+            "usage_count": self.usage_count,
+            "metadata": self.metadata
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "APIKey":
+        """Create API key from dictionary."""
+        return cls(
+            id=data.get("id", str(uuid.uuid4())),
+            key=data.get("key", ""),
+            user_id=data.get("user_id", ""),
+            name=data.get("name"),
+            created_at=datetime.fromisoformat(data.get("created_at", datetime.now(timezone.utc).isoformat())),
+            last_used_at=datetime.fromisoformat(data["last_used_at"]) if data.get("last_used_at") else None,
+            expires_at=datetime.fromisoformat(data["expires_at"]) if data.get("expires_at") else None,
+            is_active=data.get("is_active", True),
+            rate_limit=data.get("rate_limit", 10),
+            usage_count=data.get("usage_count", 0),
+            metadata=data.get("metadata", {})
+        )
+    
+    def is_expired(self) -> bool:
+        """Check if API key is expired."""
+        if not self.expires_at:
+            return False
+        return datetime.now(timezone.utc) > self.expires_at
+    
+    def is_valid(self) -> bool:
+        """Check if API key is valid (active and not expired)."""
+        return self.is_active and not self.is_expired()
+
+
+@dataclass
+class RateLimitEntry:
+    """Rate limit tracking entry."""
+    api_key_id: str
+    requests: List[datetime] = field(default_factory=list)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for storage."""
+        return {
+            "api_key_id": self.api_key_id,
+            "requests": [req.isoformat() for req in self.requests]
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "RateLimitEntry":
+        """Create from dictionary."""
+        return cls(
+            api_key_id=data.get("api_key_id", ""),
+            requests=[datetime.fromisoformat(req) for req in data.get("requests", [])]
+        )

webscout/auth/providers.py

@@ -0,0 +1,257 @@
+"""
+Provider management and initialization for the Webscout API.
+"""
+
+import sys
+import inspect
+from typing import Any, Dict, Tuple
+from starlette.status import HTTP_404_NOT_FOUND, HTTP_500_INTERNAL_SERVER_ERROR
+
+from webscout.Litlogger import Logger, LogLevel, LogFormat, ConsoleHandler
+from .config import AppConfig
+from .exceptions import APIError
+
+# Setup logger
+logger = Logger(
+    name="webscout.api",
+    level=LogLevel.INFO,
+    handlers=[ConsoleHandler(stream=sys.stdout)],
+    fmt=LogFormat.DEFAULT
+)
+
+# Cache for provider instances to avoid reinitialization on every request
+provider_instances: Dict[str, Any] = {}
+tti_provider_instances: Dict[str, Any] = {}
+
+
+def initialize_provider_map() -> None:
+    """Initialize the provider map by discovering available providers."""
+    logger.info("Initializing provider map...")
+
+    try:
+        from webscout.Provider.OPENAI.base import OpenAICompatibleProvider
+        module = sys.modules["webscout.Provider.OPENAI"]
+
+        provider_count = 0
+        model_count = 0
+
+        for name, obj in inspect.getmembers(module):
+            if (
+                inspect.isclass(obj)
+                and issubclass(obj, OpenAICompatibleProvider)
+                and obj.__name__ != "OpenAICompatibleProvider"
+            ):
+                provider_name = obj.__name__
+                AppConfig.provider_map[provider_name] = obj
+                provider_count += 1
+
+                # Register available models for this provider
+                if hasattr(obj, "AVAILABLE_MODELS") and isinstance(
+                    obj.AVAILABLE_MODELS, (list, tuple, set)
+                ):
+                    for model in obj.AVAILABLE_MODELS:
+                        if model and isinstance(model, str):
+                            model_key = f"{provider_name}/{model}"
+                            AppConfig.provider_map[model_key] = obj
+                            model_count += 1
+
+        # Fallback to ChatGPT if no providers found
+        if not AppConfig.provider_map:
+            logger.warning("No providers found, using ChatGPT fallback")
+            try:
+                from webscout.Provider.OPENAI.chatgpt import ChatGPT
+                fallback_models = ["gpt-4", "gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"]
+
+                AppConfig.provider_map["ChatGPT"] = ChatGPT
+
+                for model in fallback_models:
+                    model_key = f"ChatGPT/{model}"
+                    AppConfig.provider_map[model_key] = ChatGPT
+
+                AppConfig.default_provider = "ChatGPT"
+                provider_count = 1
+                model_count = len(fallback_models)
+            except ImportError as e:
+                logger.error(f"Failed to import ChatGPT fallback: {e}")
+                raise APIError("No providers available", HTTP_500_INTERNAL_SERVER_ERROR)
+
+        logger.info(f"Initialized {provider_count} providers with {model_count} models")
+
+    except Exception as e:
+        logger.error(f"Failed to initialize provider map: {e}")
+        raise APIError(f"Provider initialization failed: {e}", HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+def initialize_tti_provider_map() -> None:
+    """Initialize the TTI provider map by discovering available TTI providers."""
+    logger.info("Initializing TTI provider map...")
+
+    try:
+        import webscout.Provider.TTI as tti_module
+        from webscout.Provider.TTI.base import TTICompatibleProvider
+        
+        provider_count = 0
+        model_count = 0
+
+        for name, obj in inspect.getmembers(tti_module):
+            if (
+                inspect.isclass(obj)
+                and issubclass(obj, TTICompatibleProvider)
+                and obj.__name__ != "TTICompatibleProvider"
+                and obj.__name__ != "BaseImages"
+            ):
+                provider_name = obj.__name__
+                AppConfig.tti_provider_map[provider_name] = obj
+                provider_count += 1
+
+                # Register available models for this TTI provider
+                if hasattr(obj, "AVAILABLE_MODELS") and isinstance(
+                    obj.AVAILABLE_MODELS, (list, tuple, set)
+                ):
+                    for model in obj.AVAILABLE_MODELS:
+                        if model and isinstance(model, str):
+                            model_key = f"{provider_name}/{model}"
+                            AppConfig.tti_provider_map[model_key] = obj
+                            model_count += 1
+
+        # Fallback to PollinationsAI if no TTI providers found
+        if not AppConfig.tti_provider_map:
+            logger.warning("No TTI providers found, using PollinationsAI fallback")
+            try:
+                from webscout.Provider.TTI.pollinations import PollinationsAI
+                fallback_models = ["flux", "turbo", "gptimage"]
+
+                AppConfig.tti_provider_map["PollinationsAI"] = PollinationsAI
+
+                for model in fallback_models:
+                    model_key = f"PollinationsAI/{model}"
+                    AppConfig.tti_provider_map[model_key] = PollinationsAI
+
+                AppConfig.default_tti_provider = "PollinationsAI"
+                provider_count = 1
+                model_count = len(fallback_models)
+            except ImportError as e:
+                logger.error(f"Failed to import PollinationsAI fallback: {e}")
+                raise APIError("No TTI providers available", HTTP_500_INTERNAL_SERVER_ERROR)
+
+        logger.info(f"Initialized {provider_count} TTI providers with {model_count} models")
+
+    except Exception as e:
+        logger.error(f"Failed to initialize TTI provider map: {e}")
+        raise APIError(f"TTI Provider initialization failed: {e}", HTTP_500_INTERNAL_SERVER_ERROR)
+
+
+def resolve_provider_and_model(model_identifier: str) -> Tuple[Any, str]:
+    """Resolve provider class and model name from model identifier."""
+    provider_class = None
+    model_name = None
+
+    # Check for explicit provider/model syntax
+    if model_identifier in AppConfig.provider_map and "/" in model_identifier:
+        provider_class = AppConfig.provider_map[model_identifier]
+        _, model_name = model_identifier.split("/", 1)
+    elif "/" in model_identifier:
+        provider_name, model_name = model_identifier.split("/", 1)
+        provider_class = AppConfig.provider_map.get(provider_name)
+    else:
+        provider_class = AppConfig.provider_map.get(AppConfig.default_provider)
+        model_name = model_identifier
+
+    if not provider_class:
+        available_providers = list(set(v.__name__ for v in AppConfig.provider_map.values()))
+        raise APIError(
+            f"Provider for model '{model_identifier}' not found. Available providers: {available_providers}",
+            HTTP_404_NOT_FOUND,
+            "model_not_found",
+            param="model"
+        )
+
+    # Validate model availability
+    if hasattr(provider_class, "AVAILABLE_MODELS") and model_name is not None:
+        available = getattr(provider_class, "AVAILABLE_MODELS", None)
+        # If it's a property, get from instance
+        if isinstance(available, property):
+            try:
+                available = getattr(provider_class(), "AVAILABLE_MODELS", [])
+            except Exception:
+                available = []
+        # If still not iterable, fallback to empty list
+        if not isinstance(available, (list, tuple, set)):
+            available = list(available) if hasattr(available, "__iter__") and not isinstance(available, str) else []
+        if available and model_name not in available:
+            raise APIError(
+                f"Model '{model_name}' not supported by provider '{provider_class.__name__}'. Available models: {available}",
+                HTTP_404_NOT_FOUND,
+                "model_not_found",
+                param="model"
+            )
+
+    return provider_class, model_name
+
+
+def resolve_tti_provider_and_model(model_identifier: str) -> Tuple[Any, str]:
+    """Resolve TTI provider class and model name from model identifier."""
+    provider_class = None
+    model_name = None
+
+    # Check for explicit provider/model syntax
+    if model_identifier in AppConfig.tti_provider_map and "/" in model_identifier:
+        provider_class = AppConfig.tti_provider_map[model_identifier]
+        _, model_name = model_identifier.split("/", 1)
+    elif "/" in model_identifier:
+        provider_name, model_name = model_identifier.split("/", 1)
+        provider_class = AppConfig.tti_provider_map.get(provider_name)
+    else:
+        provider_class = AppConfig.tti_provider_map.get(AppConfig.default_tti_provider)
+        model_name = model_identifier
+
+    if not provider_class:
+        available_providers = list(set(v.__name__ for v in AppConfig.tti_provider_map.values()))
+        raise APIError(
+            f"TTI Provider for model '{model_identifier}' not found. Available TTI providers: {available_providers}",
+            HTTP_404_NOT_FOUND,
+            "model_not_found",
+            param="model"
+        )
+
+    # Validate model availability
+    if hasattr(provider_class, "AVAILABLE_MODELS") and model_name is not None:
+        available = getattr(provider_class, "AVAILABLE_MODELS", None)
+        # If it's a property, get from instance
+        if isinstance(available, property):
+            try:
+                available = getattr(provider_class(), "AVAILABLE_MODELS", [])
+            except Exception:
+                available = []
+        # If still not iterable, fallback to empty list
+        if not isinstance(available, (list, tuple, set)):
+            available = list(available) if hasattr(available, "__iter__") and not isinstance(available, str) else []
+        if available and model_name not in available:
+            raise APIError(
+                f"Model '{model_name}' not supported by TTI provider '{provider_class.__name__}'. Available models: {available}",
+                HTTP_404_NOT_FOUND,
+                "model_not_found",
+                param="model"
+            )
+
+    return provider_class, model_name
+
+
+def get_provider_instance(provider_class: Any):
+    """Return a cached instance of the provider, creating it if necessary."""
+    key = provider_class.__name__
+    instance = provider_instances.get(key)
+    if instance is None:
+        instance = provider_class()
+        provider_instances[key] = instance
+    return instance
+
+
+def get_tti_provider_instance(provider_class: Any):
+    """Return a cached instance of the TTI provider, creating it if needed."""
+    key = provider_class.__name__
+    instance = tti_provider_instances.get(key)
+    if instance is None:
+        instance = provider_class()
+        tti_provider_instances[key] = instance
+    return instance