drogue 0.1.0__py3-none-any.whl

drogue/__init__.py

@@ -0,0 +1,74 @@
+"""Drogue - Production-ready rate limiting and DDoS protection.
+
+Usage:
+    from drogue import DrogueLimiter, RateLimitRule
+
+    # FastAPI
+    app = FastAPI()
+    limiter = DrogueLimiter(app)
+
+    @app.get("/api/data")
+    @limiter.limit("100/minute")
+    async def get_data():
+        return {"data": "value"}
+"""
+
+from __future__ import annotations
+
+from drogue.core.abstracts import AcquireResult, Algorithm, IdentityExtractor, Storage
+from drogue.core.algorithms import (
+    FixedWindowAlgorithm,
+    SlidingWindowAlgorithm,
+    TokenBucketAlgorithm,
+)
+from drogue.core.config import DrogueConfig
+from drogue.core.errors import (
+    BackendFailure,
+    BanError,
+    ConfigurationError,
+    RateLimitExceeded,
+    StorageError,
+)
+from drogue.core.identity import (
+    CompositeExtractor,
+    HeaderExtractor,
+    RemoteAddressExtractor,
+    UserExtractor,
+)
+from drogue.core.rules.rule import AlgorithmType, RateLimitRule, parse_rule_string
+from drogue.core.storage.memory import MemoryStorage
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Core
+    "Algorithm",
+    "Storage",
+    "IdentityExtractor",
+    "AcquireResult",
+    # Config
+    "DrogueConfig",
+    # Errors
+    "RateLimitExceeded",
+    "BackendFailure",
+    "BanError",
+    "ConfigurationError",
+    "StorageError",
+    # Rules
+    "RateLimitRule",
+    "AlgorithmType",
+    "parse_rule_string",
+    # Storage
+    "MemoryStorage",
+    # Algorithms
+    "TokenBucketAlgorithm",
+    "SlidingWindowAlgorithm",
+    "FixedWindowAlgorithm",
+    # Identity
+    "RemoteAddressExtractor",
+    "UserExtractor",
+    "HeaderExtractor",
+    "CompositeExtractor",
+    # Version
+    "__version__",
+]

drogue/adapters/__init__.py

@@ -0,0 +1,3 @@
+from drogue.adapters.fastapi.limiter import DrogueLimiter
+
+__all__ = ["DrogueLimiter"]

drogue/adapters/django/__init__.py

@@ -0,0 +1,3 @@
+from drogue.adapters.django.limiter import DrogueMiddleware, DrogueRateLimiter
+
+__all__ = ["DrogueRateLimiter", "DrogueMiddleware"]

drogue/adapters/django/limiter.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+import functools
+import logging
+from typing import TYPE_CHECKING, Any
+
+from django.conf import settings
+from django.http import HttpRequest, JsonResponse
+from django.utils.deprecation import MiddlewareMixin
+
+from drogue.core.algorithms import (
+    FixedWindowAlgorithm,
+    SlidingWindowAlgorithm,
+    TokenBucketAlgorithm,
+)
+from drogue.core.config import DrogueConfig
+from drogue.core.headers import build_429_response
+from drogue.core.identity import RemoteAddressExtractor
+from drogue.core.rules.resolver import CostResolver
+from drogue.core.rules.rule import AlgorithmType, RateLimitRule, parse_rule_string
+from drogue.core.storage.memory import MemoryStorage
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from drogue.core.abstracts import AcquireResult, Algorithm, IdentityExtractor, Storage
+
+logger = logging.getLogger("drogue.django")
+
+_ALGORITHM_MAP: dict[AlgorithmType, type[Algorithm]] = {
+    AlgorithmType.TOKEN_BUCKET: TokenBucketAlgorithm,
+    AlgorithmType.SLIDING_WINDOW: SlidingWindowAlgorithm,
+    AlgorithmType.FIXED_WINDOW: FixedWindowAlgorithm,
+}
+
+# Thread-local storage for limiter instance
+_local: Any = None
+try:
+    import threading
+    _local = threading.local()
+except ImportError:
+    pass
+
+
+def _get_limiter() -> DrogueRateLimiter:
+    """Get the current limiter instance from thread-local or settings."""
+    if _local is not None and hasattr(_local, "limiter"):
+        return _local.limiter
+    return getattr(settings, "DROGUE_LIMITER", None)  # type: ignore[no-any-return]
+
+
+class DrogueRateLimiter:
+    """Main entry point for Django rate limiting.
+
+    Usage (in settings.py):
+        DROGUE_LIMITER = DrogueRateLimiter(default_limits=["100/minute"])
+
+    Usage (in views.py):
+        from drogue.django import ratelimit
+
+        @ratelimit("100/minute")
+        def my_view(request):
+            return JsonResponse({"data": "value"})
+    """
+
+    def __init__(
+        self,
+        *,
+        rules: list[RateLimitRule] | None = None,
+        storage: Storage | None = None,
+        key_func: IdentityExtractor | None = None,
+        config: DrogueConfig | None = None,
+        default_limits: list[str] | None = None,
+    ) -> None:
+        self.config = config or DrogueConfig()
+        self.storage = storage if storage is not None else MemoryStorage()
+        self.key_func = key_func or RemoteAddressExtractor(
+            trusted_proxies=self.config.trusted_proxies,
+            trust_x_real_ip=self.config.trust_x_real_ip,
+        )
+        self._algorithms: dict[str, Algorithm] = {}
+        self._route_rules: dict[str, list[RateLimitRule]] = {}
+        self._initialized = False
+
+        # Global rules
+        self._global_rules: list[RateLimitRule] = list(rules or [])
+        for limit_str in (default_limits or []):
+            rule = parse_rule_string(limit_str)
+            object.__setattr__(rule, "scope", "global")
+            self._global_rules.append(rule)
+
+    def initialize(self) -> None:
+        """Initialize storage backend. Call once at startup."""
+        _run_async(self.storage.initialize())
+        self._initialized = True
+
+    def _get_algorithm(self, rule: RateLimitRule, route_key: str = "") -> Algorithm:
+        """Get or create algorithm instance for a rule."""
+        algo_class = _ALGORITHM_MAP[rule.algorithm]
+        cache_key = f"{route_key}:{rule.algorithm.value}:{rule.limit}:{rule.window}"
+        if cache_key not in self._algorithms:
+            self._algorithms[cache_key] = algo_class(
+                storage=self.storage,
+                limit=rule.limit,
+                window=rule.window,
+            )
+        return self._algorithms[cache_key]
+
+    async def _check(
+        self,
+        key: str,
+        rule: RateLimitRule,
+        context: dict[str, Any] | None = None,
+        route_key: str = "",
+    ) -> AcquireResult:
+        """Check a single rate limit rule."""
+        algo = self._get_algorithm(rule, route_key)
+        cost = await CostResolver.resolve_cost(rule, context)
+        storage_key = f"{route_key}:{key}" if route_key else key
+        return await algo.acquire(storage_key, cost=cost, block=rule.block, timeout=rule.timeout)
+
+    def _check_sync(
+        self,
+        key: str,
+        rule: RateLimitRule,
+        context: dict[str, Any] | None = None,
+        route_key: str = "",
+    ) -> AcquireResult:
+        """Synchronous check (for Django views)."""
+        import asyncio
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = None
+
+        if loop and loop.is_running():
+            # We're in an async context, create a new event loop
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as pool:
+                future = pool.submit(
+                    asyncio.run,
+                    self._check(key, rule, context, route_key),
+                )
+                return future.result()
+        else:
+            return asyncio.run(self._check(key, rule, context, route_key))
+
+    def limit(
+        self,
+        rule_str: str,
+        *,
+        algorithm: AlgorithmType = AlgorithmType.TOKEN_BUCKET,
+        block: bool = False,
+        key_func: IdentityExtractor | None = None,
+    ) -> Callable:
+        """Decorator for view-level rate limiting.
+
+        Usage:
+            @ratelimit("100/minute")
+            def my_view(request):
+                return JsonResponse({"data": "value"})
+        """
+        rule = parse_rule_string(rule_str, algorithm=algorithm, block=block)
+
+        def decorator(view_func: Callable) -> Callable:
+            func_name = f"{view_func.__module__}.{view_func.__qualname__}"
+            self._route_rules.setdefault(func_name, []).append(rule)
+
+            @functools.wraps(view_func)
+            def wrapper(request: HttpRequest, *args: Any, **kwargs: Any) -> Any:
+                context = _request_to_context(request)
+                extractor = key_func or self.key_func
+                key = _run_async(extractor.extract(context))
+
+                for r in self._route_rules.get(func_name, [rule]):
+                    result = self._check_sync(key, r, context, route_key=func_name)
+                    if not result.allowed:
+                        return JsonResponse(
+                            build_429_response(result),
+                            status=429,
+                            headers=result.headers,
+                        )
+
+                response = view_func(request, *args, **kwargs)
+
+                # Inject headers from the first check result
+                if hasattr(response, "items") and rule.headers:
+                    for header, value in result.headers.items():
+                        response[header] = value
+
+                return response
+
+            return wrapper
+
+        return decorator
+
+
+class DrogueMiddleware(MiddlewareMixin):
+    """Django middleware for global rate limiting.
+
+    Add to MIDDLEWARE in settings.py:
+        MIDDLEWARE = [
+            ...
+            'drogue.django.middleware.DrogueMiddleware',
+        ]
+    """
+
+    def process_request(self, request: HttpRequest) -> JsonResponse | None:
+        from drogue.core.errors import BackendFailure
+
+        limiter = _get_limiter()
+        if limiter is None:
+            return None
+
+        try:
+            context = _request_to_context(request)
+            key = _run_async(limiter.key_func.extract(context))
+
+            for rule in limiter._global_rules:
+                result = limiter._check_sync(key, rule, context, route_key="__global__")
+                if not result.allowed:
+                    return JsonResponse(
+                        build_429_response(result),
+                        status=429,
+                        headers=result.headers,
+                    )
+
+            return None
+
+        except BackendFailure:
+            # Fail-closed: deny request when backend is unavailable
+            return JsonResponse(
+                {"error": "Rate limit service unavailable", "retry_after": 1},
+                status=429,
+                headers={"Retry-After": "1"},
+            )
+
+
+def _request_to_context(request: HttpRequest) -> dict[str, Any]:
+    """Convert a Django HttpRequest to a context dict."""
+    # Get real IP from proxy headers
+    x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
+    x_real_ip = request.META.get("HTTP_X_REAL_IP")
+
+    if x_real_ip:
+        client_host = x_real_ip.strip()
+    elif x_forwarded_for:
+        client_host = x_forwarded_for.split(",")[0].strip()
+    else:
+        client_host = request.META.get("REMOTE_ADDR", "127.0.0.1")
+
+    return {
+        "client": {"host": client_host},
+        "headers": {k.lower().replace("http_", ""): v for k, v in request.META.items() if k.startswith("HTTP_")},
+        "path": request.path,
+        "method": request.method,
+        "query_params": dict(request.GET),
+        "state": getattr(request, "_drogue_state", {}),
+        "request": request,
+    }
+
+
+def _run_async(coro: Any) -> Any:
+    """Run an async coroutine from sync code."""
+    from drogue.utils.async_utils import run_async
+    return run_async(coro)

drogue/adapters/django/throttle.py

@@ -0,0 +1,131 @@
+"""Django REST Framework throttle adapter using drogue.
+
+Provides DrogueThrottle that plugs into DRF's throttling system,
+using drogue's algorithms instead of DRF's built-in throttling.
+
+Usage in settings.py:
+    REST_FRAMEWORK = {
+        'DEFAULT_THROTTLE_CLASSES': ['drogue.adapters.django.throttle.DrogueThrottle'],
+        'DEFAULT_THROTTLE_RATES': {
+            'user': '100/hour',
+            'anon': '20/minute',
+        }
+    }
+
+Usage per-view:
+    from drogue.adapters.django.throttle import DrogueThrottle
+
+    class MyView(APIView):
+        throttle_classes = [DrogueThrottle]
+        throttle_scope = 'user'
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from django.conf import settings
+from rest_framework.throttling import BaseThrottle
+
+from drogue.core.algorithms import (
+    FixedWindowAlgorithm,
+    SlidingWindowAlgorithm,
+    TokenBucketAlgorithm,
+)
+from drogue.core.rules.rule import AlgorithmType, parse_rule_string
+from drogue.core.storage.memory import MemoryStorage
+
+if TYPE_CHECKING:
+    from drogue.core.abstracts import Algorithm
+
+_ALGORITHM_MAP: dict[AlgorithmType, type[Algorithm]] = {
+    AlgorithmType.TOKEN_BUCKET: TokenBucketAlgorithm,
+    AlgorithmType.SLIDING_WINDOW: SlidingWindowAlgorithm,
+    AlgorithmType.FIXED_WINDOW: FixedWindowAlgorithm,
+}
+
+# Shared storage and algorithms across throttle instances
+_storage: MemoryStorage | None = None
+_algorithms: dict[str, Algorithm] = {}
+_MAX_ALGORITHMS = 256
+
+
+def _get_storage() -> MemoryStorage:
+    global _storage
+    if _storage is None:
+        _storage = MemoryStorage()
+    return _storage
+
+
+def _get_algorithm(rule_str: str, route_key: str = "") -> Algorithm:
+    cache_key = f"{route_key}:{rule_str}"
+    if cache_key not in _algorithms:
+        if len(_algorithms) >= _MAX_ALGORITHMS:
+            _algorithms.clear()
+        rule = parse_rule_string(rule_str)
+        algo_class = _ALGORITHM_MAP[rule.algorithm]
+        _algorithms[cache_key] = algo_class(
+            storage=_get_storage(),
+            limit=rule.limit,
+            window=rule.window,
+        )
+    return _algorithms[cache_key]
+
+
+class DrogueThrottle(BaseThrottle):
+    """DRF throttle using drogue's rate limiting algorithms.
+
+    Configure rate via DRF's DEFAULT_THROTTLE_RATES or per-view:
+        throttle_scope = 'user'
+        throttle_rates = {'user': '100/hour'}
+
+    Supports:
+        - Per-user rate limiting (authenticated users)
+        - Per-IP rate limiting (anonymous users)
+        - All drogue algorithms (token_bucket, sliding_window, fixed_window)
+    """
+
+    def allow_request(self, request: Any, view: Any) -> bool:
+        """Check if the request should be allowed."""
+        self.key = self.get_cache_key(request, view)
+        if self.key is None:
+            return True
+
+        # Get rate from view or global settings
+        rate = self.get_rate(view)
+        if rate is None:
+            return True
+
+        algo = _get_algorithm(rate, route_key=f"drf:{getattr(view, 'throttle_scope', 'default')}")
+
+        # Run async acquire synchronously
+        from drogue.utils.async_utils import run_async
+        result = run_async(algo.acquire(self.key))
+
+        self._wait_after = result.retry_after
+        return result.allowed
+
+    def get_rate(self, view: Any) -> str | None:
+        """Get the throttle rate for the current view."""
+        # Check view-level throttle_rates first
+        throttle_rates = getattr(view, "throttle_rates", None)
+        scope = getattr(view, "throttle_scope", None)
+
+        if throttle_rates and scope and scope in throttle_rates:
+            return throttle_rates[scope]
+
+        # Fall back to DRF global settings
+        rates = getattr(settings, "REST_FRAMEWORK", {}).get("DEFAULT_THROTTLE_RATES", {})
+        if scope and scope in rates:
+            return rates[scope]
+
+        return None
+
+    def get_cache_key(self, request: Any, view: Any) -> str | None:
+        """Generate cache key for the request."""
+        if request.user and request.user.is_authenticated:
+            return f"throttle:user:{request.user.pk}"
+        return f"throttle:ip:{self.get_ident(request)}"
+
+    def wait(self) -> float | None:  # type: ignore[override]
+        """Seconds to wait before next request, or None if allowed."""
+        return getattr(self, "_wait_after", None)

drogue/adapters/fastapi/__init__.py

@@ -0,0 +1,3 @@
+from drogue.adapters.fastapi.limiter import DrogueLimiter
+
+__all__ = ["DrogueLimiter"]