mdb-engine 0.1.6__py3-none-any.whl → 0.2.0__py3-none-any.whl

mdb_engine/auth/rate_limiter.py

@@ -0,0 +1,504 @@
+"""
+Rate Limiting for Authentication Endpoints
+
+Provides rate limiting middleware to protect auth endpoints from brute-force attacks.
+Supports both in-memory storage (single instance) and MongoDB-backed storage (distributed).
+
+This module is part of MDB_ENGINE - MongoDB Engine.
+
+Features:
+    - Sliding window rate limiting algorithm
+    - Per-endpoint configurable limits via manifest
+    - IP + optional email-based tracking
+    - In-memory (default) or MongoDB storage
+    - 429 Too Many Requests with Retry-After header
+
+Usage:
+    # Via middleware (recommended for shared auth)
+    app.add_middleware(
+        AuthRateLimitMiddleware,
+        limits={
+            "/login": RateLimit(max_attempts=5, window_seconds=300),
+            "/register": RateLimit(max_attempts=3, window_seconds=3600),
+        }
+    )
+
+    # Via decorator (for specific endpoints)
+    @app.post("/login")
+    @rate_limit(max_attempts=5, window_seconds=300)
+    async def login(request: Request):
+        ...
+"""
+
+import logging
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+from pymongo.errors import OperationFailure
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RateLimit:
+    """Rate limit configuration for an endpoint."""
+
+    max_attempts: int = 5
+    window_seconds: int = 300  # 5 minutes
+
+    def to_dict(self) -> Dict[str, int]:
+        return {
+            "max_attempts": self.max_attempts,
+            "window_seconds": self.window_seconds,
+        }
+
+
+# Default rate limits for auth endpoints
+DEFAULT_AUTH_RATE_LIMITS: Dict[str, RateLimit] = {
+    "/login": RateLimit(max_attempts=5, window_seconds=300),
+    "/register": RateLimit(max_attempts=3, window_seconds=3600),
+    "/logout": RateLimit(max_attempts=10, window_seconds=60),
+}
+
+
+class InMemoryRateLimitStore:
+    """
+    In-memory rate limit storage using sliding window algorithm.
+
+    Suitable for single-instance deployments. For distributed systems,
+    use MongoDBRateLimitStore instead.
+    """
+
+    def __init__(self):
+        # Structure: {identifier: [(timestamp, count), ...]}
+        self._storage: Dict[str, List[Tuple[float, int]]] = defaultdict(list)
+
+    async def record_attempt(
+        self,
+        identifier: str,
+        window_seconds: int,
+    ) -> int:
+        """
+        Record an attempt and return current count in window.
+
+        Args:
+            identifier: Unique identifier (e.g., "login:192.168.1.1:user@example.com")
+            window_seconds: Time window in seconds
+
+        Returns:
+            Number of attempts in the current window (including this one)
+        """
+        now = time.time()
+        cutoff = now - window_seconds
+
+        # Clean old entries and count current
+        entries = self._storage[identifier]
+        entries[:] = [(ts, c) for ts, c in entries if ts > cutoff]
+
+        # Add new attempt
+        entries.append((now, 1))
+
+        # Return total count
+        return sum(c for _, c in entries)
+
+    async def get_count(
+        self,
+        identifier: str,
+        window_seconds: int,
+    ) -> int:
+        """Get current attempt count without recording."""
+        now = time.time()
+        cutoff = now - window_seconds
+
+        entries = self._storage.get(identifier, [])
+        return sum(c for ts, c in entries if ts > cutoff)
+
+    async def reset(self, identifier: str) -> None:
+        """Reset rate limit for an identifier (e.g., after successful login)."""
+        self._storage.pop(identifier, None)
+
+    def cleanup(self, max_age_seconds: int = 7200) -> int:
+        """
+        Clean up old entries to prevent memory growth.
+
+        Args:
+            max_age_seconds: Remove entries older than this (default: 2 hours)
+
+        Returns:
+            Number of identifiers cleaned up
+        """
+        now = time.time()
+        cutoff = now - max_age_seconds
+        cleaned = 0
+
+        identifiers_to_remove = []
+        for identifier, entries in self._storage.items():
+            entries[:] = [(ts, c) for ts, c in entries if ts > cutoff]
+            if not entries:
+                identifiers_to_remove.append(identifier)
+
+        for identifier in identifiers_to_remove:
+            del self._storage[identifier]
+            cleaned += 1
+
+        return cleaned
+
+
+class MongoDBRateLimitStore:
+    """
+    MongoDB-backed rate limit storage for distributed deployments.
+
+    Uses TTL indexes for automatic cleanup.
+    """
+
+    COLLECTION = "_mdb_engine_rate_limits"
+
+    def __init__(self, db):
+        """
+        Initialize MongoDB rate limit store.
+
+        Args:
+            db: MongoDB database instance (Motor AsyncIOMotorDatabase)
+        """
+        self._db = db
+        self._collection = db[self.COLLECTION]
+        self._indexes_created = False
+
+    async def ensure_indexes(self) -> None:
+        """Create necessary indexes."""
+        if self._indexes_created:
+            return
+
+        try:
+            # Compound index for lookups
+            await self._collection.create_index(
+                [("identifier", 1), ("timestamp", -1)], name="identifier_timestamp_idx"
+            )
+            # TTL index for cleanup
+            await self._collection.create_index(
+                "expires_at", expireAfterSeconds=0, name="expires_at_ttl_idx"
+            )
+            self._indexes_created = True
+            logger.info("Rate limit indexes ensured")
+        except OperationFailure as e:
+            logger.warning(f"Failed to create rate limit indexes: {e}")
+
+    async def record_attempt(
+        self,
+        identifier: str,
+        window_seconds: int,
+    ) -> int:
+        """Record an attempt and return current count in window."""
+        await self.ensure_indexes()
+
+        now = datetime.utcnow()
+        expires_at = now + timedelta(seconds=window_seconds)
+        cutoff = now - timedelta(seconds=window_seconds)
+
+        # Insert attempt
+        await self._collection.insert_one(
+            {
+                "identifier": identifier,
+                "timestamp": now,
+                "expires_at": expires_at,
+            }
+        )
+
+        # Count attempts in window
+        count = await self._collection.count_documents(
+            {
+                "identifier": identifier,
+                "timestamp": {"$gte": cutoff},
+            }
+        )
+
+        return count
+
+    async def get_count(
+        self,
+        identifier: str,
+        window_seconds: int,
+    ) -> int:
+        """Get current attempt count without recording."""
+        await self.ensure_indexes()
+
+        cutoff = datetime.utcnow() - timedelta(seconds=window_seconds)
+
+        count = await self._collection.count_documents(
+            {
+                "identifier": identifier,
+                "timestamp": {"$gte": cutoff},
+            }
+        )
+
+        return count
+
+    async def reset(self, identifier: str) -> None:
+        """Reset rate limit for an identifier."""
+        await self._collection.delete_many({"identifier": identifier})
+
+
+# Global in-memory store (shared across middleware instances in same process)
+_default_store = InMemoryRateLimitStore()
+
+
+class AuthRateLimitMiddleware(BaseHTTPMiddleware):
+    """
+    ASGI middleware for rate limiting authentication endpoints.
+
+    Automatically protects /login, /register, and other auth endpoints
+    from brute-force attacks.
+
+    Features:
+    - Configurable per-endpoint limits
+    - IP + email tracking for login attempts
+    - 429 responses with Retry-After header
+    - Skips rate limiting for non-auth endpoints
+
+    Usage:
+        # Basic usage with defaults
+        app.add_middleware(AuthRateLimitMiddleware)
+
+        # Custom limits
+        app.add_middleware(
+            AuthRateLimitMiddleware,
+            limits={"/login": RateLimit(max_attempts=3, window_seconds=60)}
+        )
+
+        # With MongoDB storage for distributed deployments
+        app.add_middleware(
+            AuthRateLimitMiddleware,
+            store=MongoDBRateLimitStore(db)
+        )
+    """
+
+    def __init__(
+        self,
+        app: Callable,
+        limits: Optional[Dict[str, RateLimit]] = None,
+        store: Optional[InMemoryRateLimitStore] = None,
+        include_email_in_key: bool = True,
+    ):
+        """
+        Initialize rate limit middleware.
+
+        Args:
+            app: ASGI application
+            limits: Dict of path -> RateLimit config. Defaults to DEFAULT_AUTH_RATE_LIMITS.
+            store: Rate limit storage backend. Defaults to in-memory store.
+            include_email_in_key: Include email in rate limit key for more granular limits.
+        """
+        super().__init__(app)
+        self._limits = limits or DEFAULT_AUTH_RATE_LIMITS
+        self._store = store or _default_store
+        self._include_email_in_key = include_email_in_key
+
+        logger.info(
+            f"AuthRateLimitMiddleware initialized with limits for: {list(self._limits.keys())}"
+        )
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable[[Request], Response],
+    ) -> Response:
+        """Process request through rate limiter."""
+        path = request.url.path
+        method = request.method
+
+        # Only rate limit POST requests to configured endpoints
+        if method != "POST" or path not in self._limits:
+            return await call_next(request)
+
+        limit = self._limits[path]
+        identifier = await self._build_identifier(request, path)
+
+        # Check current count (before recording this attempt)
+        current_count = await self._store.get_count(identifier, limit.window_seconds)
+
+        if current_count >= limit.max_attempts:
+            logger.warning(
+                f"Rate limit exceeded: {identifier} "
+                f"({current_count}/{limit.max_attempts} in {limit.window_seconds}s)"
+            )
+            return self._rate_limit_response(limit.window_seconds)
+
+        # Record this attempt
+        await self._store.record_attempt(identifier, limit.window_seconds)
+
+        # Process request
+        response = await call_next(request)
+
+        # Reset on successful login (2xx response)
+        if response.status_code < 300 and path == "/login":
+            await self._store.reset(identifier)
+
+        return response
+
+    async def _build_identifier(self, request: Request, path: str) -> str:
+        """Build rate limit identifier from request."""
+        parts = [path]
+
+        # Add client IP
+        client_ip = self._get_client_ip(request)
+        parts.append(client_ip)
+
+        # Optionally add email for more granular rate limiting
+        if self._include_email_in_key:
+            email = await self._extract_email(request)
+            if email:
+                parts.append(email)
+
+        return ":".join(parts)
+
+    def _get_client_ip(self, request: Request) -> str:
+        """Get client IP, respecting proxy headers."""
+        # Check X-Forwarded-For header (set by proxies/load balancers)
+        forwarded_for = request.headers.get("X-Forwarded-For")
+        if forwarded_for:
+            # Take the first IP (original client)
+            return forwarded_for.split(",")[0].strip()
+
+        # Check X-Real-IP header
+        real_ip = request.headers.get("X-Real-IP")
+        if real_ip:
+            return real_ip.strip()
+
+        # Fall back to direct client IP
+        if request.client:
+            return request.client.host
+
+        return "unknown"
+
+    async def _extract_email(self, request: Request) -> Optional[str]:
+        """Try to extract email from request body."""
+        try:
+            # Only try to read body for JSON requests
+            content_type = request.headers.get("content-type", "")
+            if "application/json" in content_type:
+                # Note: This consumes the body, so we need to be careful
+                # In practice, this is called before the body is read by the route
+                body = await request.body()
+                if body:
+                    import json
+
+                    data = json.loads(body)
+                    return data.get("email")
+        except (ValueError, UnicodeDecodeError, KeyError):
+            pass
+        return None
+
+    @staticmethod
+    def _rate_limit_response(retry_after: int) -> JSONResponse:
+        """Return 429 Too Many Requests response."""
+        return JSONResponse(
+            status_code=429,
+            content={
+                "detail": f"Too many attempts. Please try again in {retry_after} seconds.",
+                "error": "rate_limit_exceeded",
+                "retry_after": retry_after,
+            },
+            headers={"Retry-After": str(retry_after)},
+        )
+
+
+def create_rate_limit_middleware(
+    manifest_auth: Dict[str, Any],
+    store: Optional[InMemoryRateLimitStore] = None,
+) -> type:
+    """
+    Factory function to create rate limit middleware from manifest config.
+
+    Args:
+        manifest_auth: Auth section from manifest
+        store: Optional storage backend
+
+    Returns:
+        Configured middleware class
+
+    Manifest format:
+        {
+            "auth": {
+                "rate_limits": {
+                    "/login": {"max_attempts": 5, "window_seconds": 300},
+                    "/register": {"max_attempts": 3, "window_seconds": 3600}
+                }
+            }
+        }
+    """
+    rate_limits_config = manifest_auth.get("rate_limits", {})
+
+    limits: Dict[str, RateLimit] = {}
+    for path, config in rate_limits_config.items():
+        limits[path] = RateLimit(
+            max_attempts=config.get("max_attempts", 5),
+            window_seconds=config.get("window_seconds", 300),
+        )
+
+    # Use defaults if no config provided
+    if not limits:
+        limits = DEFAULT_AUTH_RATE_LIMITS.copy()
+
+    class ConfiguredRateLimitMiddleware(AuthRateLimitMiddleware):
+        def __init__(self, app: Callable):
+            super().__init__(app, limits=limits, store=store)
+
+    return ConfiguredRateLimitMiddleware
+
+
+def rate_limit(
+    max_attempts: int = 5,
+    window_seconds: int = 300,
+    key_func: Optional[Callable[[Request], str]] = None,
+):
+    """
+    Decorator for rate limiting individual endpoints.
+
+    Usage:
+        @app.post("/login")
+        @rate_limit(max_attempts=5, window_seconds=300)
+        async def login(request: Request):
+            ...
+
+    Args:
+        max_attempts: Maximum attempts in window
+        window_seconds: Time window in seconds
+        key_func: Optional function to generate rate limit key from request
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        async def wrapper(request: Request, *args, **kwargs):
+            # Build identifier
+            if key_func:
+                identifier = key_func(request)
+            else:
+                client_ip = request.client.host if request.client else "unknown"
+                identifier = f"{func.__name__}:{client_ip}"
+
+            # Check rate limit
+            count = await _default_store.record_attempt(identifier, window_seconds)
+
+            if count > max_attempts:
+                logger.warning(f"Rate limit exceeded: {identifier} ({count}/{max_attempts})")
+                return JSONResponse(
+                    status_code=429,
+                    content={
+                        "detail": f"Too many attempts. Try again in {window_seconds} seconds.",
+                        "error": "rate_limit_exceeded",
+                    },
+                    headers={"Retry-After": str(window_seconds)},
+                )
+
+            return await func(request, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

mdb_engine/auth/restrictions.py

@@ -27,9 +27,7 @@ from .users import get_app_user
 logger = logging.getLogger(__name__)
 
 
-def is_demo_user(
-    user: Optional[Dict[str, Any]] = None, email: Optional[str] = None
-) -> bool:
+def is_demo_user(user: Optional[Dict[str, Any]] = None, email: Optional[str] = None) -> bool:
     """
     Check if a user is a demo user.
 
@@ -85,9 +83,7 @@ async def _get_sub_auth_user(
         if not (config and users_config.get("enabled", False)):
             return None
 
-        app_user = await get_app_user(
-            request, slug_id, db, config, allow_demo_fallback=False
-        )
+        app_user = await get_app_user(request, slug_id, db, config, allow_demo_fallback=False)
         if not app_user:
             return None
 
@@ -119,9 +115,7 @@ async def _get_authenticated_user(
 
     # Try sub-auth if platform auth didn't work
     if get_app_db_func and get_app_config_func:
-        return await _get_sub_auth_user(
-            request, slug_id, get_app_config_func, get_app_db_func
-        )
+        return await _get_sub_auth_user(request, slug_id, get_app_config_func, get_app_db_func)
 
     return None
 
@@ -158,9 +152,7 @@ def _validate_dependencies(
 
 async def require_non_demo_user(
     request: Request,
-    get_app_config_func: Optional[
-        Callable[[Request, str, Dict], Awaitable[Dict]]
-    ] = None,
+    get_app_config_func: Optional[Callable[[Request, str, Dict], Awaitable[Dict]]] = None,
     get_app_db_func: Optional[Callable[[Request], Awaitable[Any]]] = None,
 ) -> Dict[str, Any]:
     """
@@ -189,9 +181,7 @@ async def require_non_demo_user(
     if get_app_db_func and get_app_config_func:
         _validate_dependencies(get_app_config_func, get_app_db_func)
 
-    user = await _get_authenticated_user(
-        request, slug_id, get_app_config_func, get_app_db_func
-    )
+    user = await _get_authenticated_user(request, slug_id, get_app_config_func, get_app_db_func)
 
     # Check if user is demo
     if user and is_demo_user(user):
@@ -215,9 +205,7 @@ async def require_non_demo_user(
 
 async def block_demo_users(
     request: Request,
-    get_app_config_func: Optional[
-        Callable[[Request, str, Dict], Awaitable[Dict]]
-    ] = None,
+    get_app_config_func: Optional[Callable[[Request, str, Dict], Awaitable[Dict]]] = None,
     get_app_db_func: Optional[Callable[[Request], Awaitable[Any]]] = None,
 ):
     """
@@ -253,15 +241,11 @@ async def block_demo_users(
 
     # Check sub-auth if platform auth didn't work
     if not user and get_app_db_func and get_app_config_func and slug_id:
-        user = await _get_sub_auth_user(
-            request, slug_id, get_app_config_func, get_app_db_func
-        )
+        user = await _get_sub_auth_user(request, slug_id, get_app_config_func, get_app_db_func)
 
     # Block if demo user (only if user exists)
     if user and is_demo_user(user):
-        logger.info(
-            f"Demo user '{user.get('email')}' blocked from accessing: {request.url.path}"
-        )
+        logger.info(f"Demo user '{user.get('email')}' blocked from accessing: {request.url.path}")
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail="Demo users cannot access this endpoint. Demo mode is read-only.",

mdb_engine/auth/session_manager.py

@@ -13,8 +13,11 @@ from typing import Any, Dict, List, Optional
 from bson.objectid import ObjectId
 
 try:
-    from pymongo.errors import (ConnectionFailure, OperationFailure,
-                                ServerSelectionTimeoutError)
+    from pymongo.errors import (
+        ConnectionFailure,
+        OperationFailure,
+        ServerSelectionTimeoutError,
+    )
 except ImportError:
     ConnectionFailure = Exception
     OperationFailure = Exception
@@ -120,9 +123,7 @@ class SessionManager:
                 # Remove oldest inactive session
                 await self.cleanup_inactive_sessions(user_id)
                 # Check again
-                active_sessions = await self.get_user_sessions(
-                    user_id, active_only=True
-                )
+                active_sessions = await self.get_user_sessions(user_id, active_only=True)
                 if len(active_sessions) >= self.max_sessions:
                     # Force remove oldest session
                     if active_sessions:
@@ -168,9 +169,7 @@ class SessionManager:
             ValueError,
             TypeError,
         ) as e:
-            logger.error(
-                f"Error creating session for user {user_id}: {e}", exc_info=True
-            )
+            logger.error(f"Error creating session for user {user_id}: {e}", exc_info=True)
             return None
 
     async def update_session_activity(
@@ -205,14 +204,10 @@ class SessionManager:
             ValueError,
             TypeError,
         ) as e:
-            logger.error(
-                f"Error updating session activity for {refresh_jti}: {e}", exc_info=True
-            )
+            logger.error(f"Error updating session activity for {refresh_jti}: {e}", exc_info=True)
             return False
 
-    async def get_session_by_refresh_token(
-        self, refresh_jti: str
-    ) -> Optional[Dict[str, Any]]:
+    async def get_session_by_refresh_token(self, refresh_jti: str) -> Optional[Dict[str, Any]]:
         """
         Get session by refresh token JWT ID.
 
@@ -223,9 +218,7 @@ class SessionManager:
             Session document or None if not found
         """
         try:
-            session = await self.collection.find_one(
-                {"refresh_jti": refresh_jti, "active": True}
-            )
+            session = await self.collection.find_one({"refresh_jti": refresh_jti, "active": True})
             return session
         except (
             OperationFailure,
@@ -267,9 +260,7 @@ class SessionManager:
             stored_fingerprint = session.get("session_fingerprint")
 
             if not stored_fingerprint:
-                strict_mode = (
-                    strict if strict is not None else self.fingerprinting_strict
-                )
+                strict_mode = strict if strict is not None else self.fingerprinting_strict
                 return not strict_mode
 
             return stored_fingerprint == current_fingerprint
@@ -313,9 +304,7 @@ class SessionManager:
             if active_only:
                 query["active"] = True
 
-            sessions = (
-                await self.collection.find(query).sort("last_seen", -1).to_list(None)
-            )
+            sessions = await self.collection.find(query).sort("last_seen", -1).to_list(None)
             return sessions
         except (
             OperationFailure,
@@ -324,9 +313,7 @@ class SessionManager:
             ValueError,
             TypeError,
         ) as e:
-            logger.error(
-                f"Error getting sessions for user {user_id}: {e}", exc_info=True
-            )
+            logger.error(f"Error getting sessions for user {user_id}: {e}", exc_info=True)
             return []
 
     async def revoke_session(self, session_id: Any) -> bool:
@@ -400,9 +387,7 @@ class SessionManager:
             ValueError,
             TypeError,
         ) as e:
-            logger.error(
-                f"Error revoking sessions for user {user_id}: {e}", exc_info=True
-            )
+            logger.error(f"Error revoking sessions for user {user_id}: {e}", exc_info=True)
             return 0
 
     async def cleanup_inactive_sessions(self, user_id: Optional[str] = None) -> int: