pycurb 0.1.0__py3-none-any.whl

pycurb/adapters/django/__init__.py

@@ -0,0 +1,11 @@
+from .extractors import ip_extractor, user_id_extractor, api_key_extractor
+from .decorators import rate_limit
+from .middleware import create_rate_limit_middleware
+
+__all__ = [
+    "ip_extractor",
+    "user_id_extractor",
+    "api_key_extractor",
+    "rate_limit",
+    "create_rate_limit_middleware",
+]

pycurb/adapters/django/decorators.py

@@ -0,0 +1,77 @@
+import functools
+import inspect
+from typing import Union, List, Callable, Optional, cast
+from django.http import JsonResponse, HttpRequest
+from pycurb.core import RateLimiter, AsyncRateLimiter
+from pycurb.core.models import RateLimitHeaders
+from .extractors import ip_extractor
+
+
+def rate_limit(
+    limiter: Union[RateLimiter, AsyncRateLimiter],
+    rule_name: Union[str, List[str]],
+    key_extractor: Callable[[HttpRequest], str] = ip_extractor,
+    on_limit: Optional[Callable] = None,
+):
+    """
+    Django view decorator for rate limiting.
+    Works for sync and async views.
+    Returns 429 JSON response when limit exceeded.
+    """
+
+    def decorator(view_func):
+        is_async = inspect.iscoroutinefunction(view_func)
+        limiter_is_async = isinstance(limiter, AsyncRateLimiter)  # RateLimiter is async
+
+        if is_async and not limiter_is_async:
+            raise TypeError(
+                "Async view requires an async RateLimiter (use AsyncRateLimiter, not RateLimiter)"
+            )
+        if not is_async and limiter_is_async:
+            raise TypeError(
+                "Sync view requires a sync RateLimiter (use RateLimiter, not AsyncRateLimiter)"
+            )
+
+        @functools.wraps(view_func)
+        def sync_wrapper(request, *args, **kwargs):
+            key = key_extractor(request)
+            limiter_sync = cast(RateLimiter, limiter)
+            result = limiter_sync.check(key, rule_name)
+            if not result.allowed:
+                if on_limit:
+                    return on_limit(request, result)
+                headers = RateLimitHeaders.from_result(result)
+                response = JsonResponse({"detail": "Rate limit exceeded"}, status=429)
+                for name, value in headers.to_dict().items():
+                    response[name] = value
+                return response
+
+            response = view_func(request, *args, **kwargs)
+            headers = RateLimitHeaders.from_result(result)
+            for name, value in headers.to_dict().items():
+                response[name] = value
+            return response
+
+        @functools.wraps(view_func)
+        async def async_wrapper(request, *args, **kwargs):
+            key = key_extractor(request)
+            limiter_async = cast(AsyncRateLimiter, limiter)
+            result = await limiter_async.check(key, rule_name)
+            if not result.allowed:
+                if on_limit:
+                    return on_limit(request, result)
+                headers = RateLimitHeaders.from_result(result)
+                response = JsonResponse({"detail": "Rate limit exceeded"}, status=429)
+                for name, value in headers.to_dict().items():
+                    response[name] = value
+                return response
+
+            response = await view_func(request, *args, **kwargs)
+            headers = RateLimitHeaders.from_result(result)
+            for name, value in headers.to_dict().items():
+                response[name] = value
+            return response
+
+        return async_wrapper if is_async else sync_wrapper
+
+    return decorator

pycurb/adapters/django/extractors.py

@@ -0,0 +1,25 @@
+from django.http import HttpRequest
+
+
+def ip_extractor(request: HttpRequest) -> str:
+    """Extract client IP from request, respecting X-Forwarded-For."""
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0].strip()
+    return request.META.get("REMOTE_ADDR", "unknown")
+
+
+def user_id_extractor(request: HttpRequest) -> str:
+    """Extract authenticated user ID (or 'anon')."""
+    if hasattr(request, "user") and request.user.is_authenticated:
+        return str(request.user.pk)
+    return "anon"
+
+
+def api_key_extractor(request: HttpRequest, header: str = "X-API-Key") -> str:
+    """Extract API key from a custom header."""
+    return request.headers.get(header, "")
+
+def custom_extractor(extractor_func):
+    """Decorator to register custom extractors (optional)."""
+    return extractor_func

pycurb/adapters/django/middleware.py

@@ -0,0 +1,46 @@
+from django.http import JsonResponse
+from typing import Callable, Union
+from pycurb.core import RateLimiter, AsyncRateLimiter
+from pycurb.core.models import RateLimitHeaders
+from .extractors import ip_extractor
+
+
+def create_rate_limit_middleware(
+    limiter: Union[RateLimiter, AsyncRateLimiter],
+    rule_name: str,
+    key_extractor: Callable = ip_extractor,
+):
+    """
+    Factory that returns a Django middleware class for global rate limiting.
+    Works with sync limiter for now.
+    If you use async limiter, you must ensure it is compatible with sync calls
+    (e.g., by using asyncio.run). For simplicity, use RateLimiter.
+    """
+    if isinstance(limiter, AsyncRateLimiter):
+        raise TypeError(
+            "Out-of-box Middleware only supports a sync limiter for now. Use RateLimiter."
+        )
+
+    class RateLimitMiddleware:
+        def __init__(self, get_response):
+            self.get_response = get_response
+            self.limiter = limiter
+            self.rule_name = rule_name
+            self.key_extractor = key_extractor
+
+        def __call__(self, request):
+            key = self.key_extractor(request)
+            result = self.limiter.check(key, self.rule_name)
+            if not result.allowed:
+                headers = RateLimitHeaders.from_result(result)
+                response = JsonResponse({"detail": "Rate limit exceeded"}, status=429)
+                for name, value in headers.to_dict().items():
+                    response[name] = value
+                return response
+            response = self.get_response(request)
+            headers = RateLimitHeaders.from_result(result)
+            for name, value in headers.to_dict().items():
+                response[name] = value
+            return response
+
+    return RateLimitMiddleware

pycurb/adapters/fastapi/__init__.py

@@ -0,0 +1,17 @@
+from .extractors import (
+    ip_extractor,
+    api_key_extractor,
+    user_id_extractor,
+    custom_extractor,
+)
+from .dependencies import rate_limiter
+from .middleware import RateLimitMiddleware
+
+__all__ = [
+    "ip_extractor",
+    "api_key_extractor",
+    "user_id_extractor",
+    "custom_extractor",
+    "rate_limiter",
+    "RateLimitMiddleware",
+]

pycurb/adapters/fastapi/dependencies.py

@@ -0,0 +1,30 @@
+from typing import Union, List, Callable
+from fastapi import Request, HTTPException
+from pycurb.core import AsyncRateLimiter, RateLimitResult
+from .extractors import ip_extractor
+import time
+
+
+def rate_limiter(
+    limiter: AsyncRateLimiter,
+    rule_name: Union[str, List[str]],
+    key_extractor: Callable[[Request], str] = ip_extractor,
+):
+    """
+    Factory that creates a FastAPI dependency for rate limiting.
+    """
+
+    async def dependency(request: Request) -> RateLimitResult:
+        key = key_extractor(request)
+        result = await limiter.check(key, rule_name)
+        if not result.allowed:
+            retry_after = max(1, int(result.reset_at - time.time()))
+            raise HTTPException(
+                status_code=429,
+                detail="Rate limit exceeded",
+                headers={"Retry-After": str(retry_after)},
+            )
+        request.state.rate_limit_result = result
+        return result
+
+    return dependency

pycurb/adapters/fastapi/extractors.py

@@ -0,0 +1,24 @@
+from fastapi import Request
+
+
+def ip_extractor(request: Request) -> str:
+    """Extract client IP address, respecting X-Forwarded-For."""
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0].strip()
+    return request.client.host if request.client else "unknown"
+
+
+def api_key_extractor(request: Request, header_name: str = "X-API-Key") -> str:
+    """Extract API key from a header."""
+    return request.headers.get(header_name, "")
+
+
+def user_id_extractor(request: Request) -> str:
+    """Extract user ID from request state (set by auth middleware)."""
+    return getattr(request.state, "user_id", "anon")
+
+
+def custom_extractor(extractor_func):
+    """Decorator to register custom extractors (optional)."""
+    return extractor_func

pycurb/adapters/fastapi/middleware.py

@@ -0,0 +1,53 @@
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+from typing import Callable, Optional
+from pycurb.core import AsyncRateLimiter
+from pycurb.core.models import RateLimitHeaders
+from .extractors import ip_extractor
+
+
+class RateLimitMiddleware(BaseHTTPMiddleware):
+    """
+    ASGI middleware to apply a global rate limit rule to all requests.
+    Optionally exclude certain paths.
+    """
+
+    def __init__(
+        self,
+        app: ASGIApp,
+        limiter: AsyncRateLimiter,
+        rule_name: str,
+        key_extractor: Callable[[Request], str] = ip_extractor,
+        exclude_paths: Optional[list[str]] = None,
+    ):
+        super().__init__(app)
+        self.limiter = limiter
+        self.rule_name = rule_name
+        self.key_extractor = key_extractor
+        self.exclude_paths = exclude_paths or []
+
+    async def dispatch(self, request: Request, call_next):
+        # Skip excluded paths
+        for path in self.exclude_paths:
+            if request.url.path.startswith(path):
+                return await call_next(request)
+
+        key = self.key_extractor(request)
+        result = await self.limiter.check(key, self.rule_name)
+
+        if not result.allowed:
+            headers = RateLimitHeaders.from_result(result)
+            return JSONResponse(
+                status_code=429,
+                content={"detail": "Rate limit exceeded"},
+                headers=headers.to_dict(),
+            )
+
+        # Add rate limit headers to response
+        response = await call_next(request)
+        headers = RateLimitHeaders.from_result(result)
+        for name, value in headers.to_dict().items():
+            response.headers[name] = value
+        return response

pycurb/adapters/flask/__init__.py

@@ -0,0 +1,17 @@
+from .extractors import (
+    ip_extractor,
+    flask_login_user_extractor,
+    session_user_extractor,
+    api_key_extractor,
+)
+from .decorators import rate_limit
+from .middleware import RateLimit
+
+__all__ = [
+    "ip_extractor",
+    "flask_login_user_extractor",
+    "session_user_extractor",
+    "api_key_extractor",
+    "rate_limit",
+    "RateLimit",
+]

pycurb/adapters/flask/decorators.py

@@ -0,0 +1,44 @@
+import functools
+from flask import make_response, jsonify
+from typing import Union, Callable, List, Optional
+from pycurb.core import RateLimiter
+from pycurb.core.models import RateLimitHeaders
+from .extractors import ip_extractor
+
+
+def rate_limit(
+    limiter: RateLimiter,
+    rule_name: Union[str, List[str]],
+    key_extractor: Callable[..., str] = ip_extractor,
+    on_limit: Optional[Callable] = None,
+):
+    """
+    Flask view decorator for rate limiting.
+    Works with sync views.
+    """
+
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            key = key_extractor()
+            result = limiter.check(key, rule_name)
+            if not result.allowed:
+                if on_limit:
+                    response = on_limit(result)
+                    if response is not None:
+                        return response
+                headers = RateLimitHeaders.from_result(result)
+                resp = make_response(jsonify({"detail": "Rate limit exceeded"}), 429)
+                for name, value in headers.to_dict().items():
+                    resp.headers[name] = value
+                return resp
+            rv = f(*args, **kwargs)
+            resp = make_response(rv)
+            headers = RateLimitHeaders.from_result(result)
+            for name, value in headers.to_dict().items():
+                resp.headers[name] = value
+            return resp
+
+        return wrapper
+
+    return decorator

pycurb/adapters/flask/extractors.py

@@ -0,0 +1,35 @@
+from flask import request, session
+
+
+def ip_extractor():
+    """Extract client IP from request, respecting X-Forwarded-For."""
+    forwarded = request.headers.get("X-Forwarded-For")
+    if forwarded:
+        return forwarded.split(",")[0].strip()
+    return request.remote_addr or "unknown"
+
+
+def flask_login_user_extractor():
+    """Extract authenticated user ID (using flask_login)."""
+    try:
+        from flask_login import current_user
+
+        if current_user.is_authenticated:
+            return str(current_user.get_id())
+    except (ImportError, AttributeError):
+        pass
+    return "anon"
+
+
+def session_user_extractor():
+    """Extract authenticated user ID directly from session."""
+    return session.get("user_id", "anon")
+
+
+def api_key_extractor(header="X-API-Key"):
+    """Extract API key from a custom header."""
+    return request.headers.get(header, "")
+
+def custom_extractor(extractor_func):
+    """Decorator to register custom extractors (optional)."""
+    return extractor_func

pycurb/adapters/flask/middleware.py

@@ -0,0 +1,54 @@
+from flask import make_response, jsonify, request
+from typing import Optional, Union, List, Callable
+from pycurb.core import RateLimiter
+from pycurb.core.models import RateLimitHeaders
+from .extractors import ip_extractor
+
+
+class RateLimit:
+    """
+    Flask extension for global rate limiting.
+    Usage:
+        app = Flask(__name__)
+        limiter = RateLimiterSync(...)
+        RateLimit(app, limiter, rule_name="global")
+    """
+
+    def __init__(
+        self,
+        app,
+        limiter: RateLimiter,
+        rule_name: Union[str, List[str]],
+        key_extractor: Callable[..., str],
+    ):
+        self.app = app
+        self.limiter = limiter
+        self.rule_name = rule_name
+        self.key_extractor = key_extractor or ip_extractor
+        if app is not None:
+            self.init_app(app)
+
+    def init_app(self, app):
+        app.before_request(self.before_request)
+        app.after_request(self.after_request)
+        app.extensions["ratelimit"] = self
+
+    def before_request(self):
+        key = self.key_extractor()
+        result = self.limiter.check(key, self.rule_name)
+        if not result.allowed:
+            headers = RateLimitHeaders.from_result(result)
+            resp = make_response(jsonify({"detail": "Rate limit exceeded"}), 429)
+            for name, value in headers.to_dict().items():
+                resp.headers[name] = value
+            return resp
+        # Store result in request context for after_request
+        request.environ["_rate_limit_result"] = result
+
+    def after_request(self, response):
+        result = request.environ.get("_rate_limit_result", None)
+        if result:
+            headers = RateLimitHeaders.from_result(result)
+            for name, value in headers.to_dict().items():
+                response.headers[name] = value
+        return response

pycurb/core/__init__.py

@@ -0,0 +1,24 @@
+from .models import LimitRule, RateLimitResult, RateLimitHeaders, RateLimitExceeded
+from .limiter_async import AsyncRateLimiter
+from .limiter import RateLimiter
+from .resolver import RuleResolver, AsyncRuleResolver
+from .decorators import rate_limit, arg_extractor
+from .storage import MemoryStorage, AsyncMemoryStorage, RedisStorage, AsyncRedisStorage
+
+
+__all__ = [
+    "LimitRule",
+    "RateLimitResult",
+    "RateLimitHeaders",
+    "RateLimiter",
+    "AsyncRateLimiter",
+    "RuleResolver",
+    "AsyncRuleResolver",
+    "MemoryStorage",
+    "AsyncMemoryStorage",
+    "RedisStorage",
+    "AsyncRedisStorage",
+    "rate_limit",
+    "RateLimitExceeded",
+    "arg_extractor",
+]

pycurb/core/algorithms/__init__.py

@@ -0,0 +1,12 @@
+from .fixed_window_async import AsyncFixedWindowAlgorithm
+from .fixed_window import FixedWindowAlgorithm
+from .sliding_window_async import AsyncSlidingWindowAlgorithm
+from .sliding_window import SlidingWindowAlgorithm
+from .token_bucket_async import AsyncTokenBucketAlgorithm
+from .token_bucket import TokenBucketAlgorithm
+from .leaky_bucket_async import AsyncLeakyBucketAlgorithm
+from .leaky_bucket import LeakyBucketAlgorithm
+from .gcra import GcraAlgorithm
+from .gcra_async import AsyncGcraAlgorithm
+from .base_async import AsyncRateLimiterAlgorithm
+from .base import RateLimiterAlgorithm

pycurb/core/algorithms/base.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+from ..models import LimitRule, RateLimitResult
+from ..storage import Storage
+
+
+class RateLimiterAlgorithm(ABC):
+    """Abstract base for all rate limiting algorithms."""
+
+    @abstractmethod
+    def check(self, key: str, rule: LimitRule, storage: Storage) -> RateLimitResult:
+        """
+        Evaluate rate limit for a given key and rule.
+        Args:
+            key: Unique client identifier (extracted by adapter) prefixed with rule name
+            rule: The limit rule applied
+            storage: Storage backend for atomic counters
+
+        Returns:
+            RateLimitResult: Decision and metadata.
+        """
+        pass

pycurb/core/algorithms/base_async.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+from ..models import LimitRule, RateLimitResult
+from ..storage import AsyncStorage
+
+
+class AsyncRateLimiterAlgorithm(ABC):
+    """Async abstract base for all rate limiting algorithms."""
+
+    @abstractmethod
+    async def check(
+        self, key: str, rule: LimitRule, storage: AsyncStorage
+    ) -> RateLimitResult:
+        """
+        Evaluate rate limit for a given key and rule.
+        Args:
+            key: Unique client identifier (extracted by adapter) prefixed with rule name
+            rule: The limit rule applied
+            storage: Storage backend for atomic counters
+
+        Returns:
+            RateLimitResult: Decision and metadata.
+        """
+        pass

pycurb/core/algorithms/fixed_window.py

@@ -0,0 +1,26 @@
+import time
+import math
+from .base import RateLimiterAlgorithm
+from ..models import LimitRule, RateLimitResult
+from ..storage import Storage
+
+
+class FixedWindowAlgorithm(RateLimiterAlgorithm):
+    def check(self, key: str, rule: LimitRule, storage: Storage) -> RateLimitResult:
+        if rule.limit is None or rule.window is None:
+            raise ValueError("Fixed window algorithm requires 'limit' and 'window'.")
+
+        now = time.time()
+        storage_key = f"{rule.name}:{key}"
+        allowed, remaining, reset_at = storage.fixed_window(
+            key=storage_key, window=rule.window, limit=rule.limit, now=now
+        )
+        retry_after = max(0, math.ceil(reset_at - now))
+        return RateLimitResult(
+            allowed=allowed,
+            remaining=remaining,
+            reset_at=reset_at,
+            limit=rule.limit,
+            retry_after=retry_after if not allowed else None,
+            rule_name=rule.name,
+        )

pycurb/core/algorithms/fixed_window_async.py

@@ -0,0 +1,28 @@
+import time
+import math
+from .base_async import AsyncRateLimiterAlgorithm
+from ..models import LimitRule, RateLimitResult
+from ..storage import AsyncStorage
+
+
+class AsyncFixedWindowAlgorithm(AsyncRateLimiterAlgorithm):
+    async def check(
+        self, key: str, rule: LimitRule, storage: AsyncStorage
+    ) -> RateLimitResult:
+        if rule.limit is None or rule.window is None:
+            raise ValueError("Fixed window algorithm requires 'limit' and 'window'.")
+
+        now = time.time()
+        storage_key = f"{rule.name}:{key}"
+        allowed, remaining, reset_at = await storage.fixed_window(
+            key=storage_key, window=rule.window, limit=rule.limit, now=now
+        )
+        retry_after = max(0, math.ceil(reset_at - now))
+        return RateLimitResult(
+            allowed=allowed,
+            remaining=remaining,
+            reset_at=reset_at,
+            limit=rule.limit,
+            retry_after=retry_after if not allowed else None,
+            rule_name=rule.name,
+        )

pycurb/core/algorithms/gcra.py

@@ -0,0 +1,38 @@
+import time
+import math
+from .base import RateLimiterAlgorithm
+from ..models import LimitRule, RateLimitResult
+from ..storage import Storage
+
+
+class GcraAlgorithm(RateLimiterAlgorithm):
+    def check(self, key: str, rule: LimitRule, storage: Storage) -> RateLimitResult:
+        capacity = rule.capacity if rule.capacity is not None else rule.limit
+        if capacity is None:
+            raise ValueError("Gcra algorithm requires 'capacity' or 'limit'.")
+
+        if rule.refill_rate is not None:
+            rate = rule.refill_rate
+        else:
+            if rule.window is None:
+                raise ValueError("Gcra algorithm requires rate equal to 'refill_rate'")
+            rate = capacity / rule.window
+
+        if rate <= 0:
+            raise ValueError(f"rate must be positive, got {rate}")
+
+        now = time.time()
+        storage_key = f"{rule.name}:{key}"
+        allowed, remaining, reset_at = storage.gcra(
+            key=storage_key, capacity=capacity, rate=rate, now=now
+        )
+        retry_after = max(0, math.ceil(reset_at - now))
+
+        return RateLimitResult(
+            allowed=allowed,
+            remaining=remaining,
+            reset_at=reset_at,
+            limit=capacity,
+            retry_after=retry_after if not allowed else None,
+            rule_name=rule.name,
+        )

pycurb/core/algorithms/gcra_async.py

@@ -0,0 +1,40 @@
+import time
+import math
+from .base_async import AsyncRateLimiterAlgorithm
+from ..models import LimitRule, RateLimitResult
+from ..storage import AsyncStorage
+
+
+class AsyncGcraAlgorithm(AsyncRateLimiterAlgorithm):
+    async def check(
+        self, key: str, rule: LimitRule, storage: AsyncStorage
+    ) -> RateLimitResult:
+        capacity = rule.capacity if rule.capacity is not None else rule.limit
+        if capacity is None:
+            raise ValueError("Gcra algorithm requires 'capacity' or 'limit'.")
+
+        if rule.refill_rate is not None:
+            rate = rule.refill_rate
+        else:
+            if rule.window is None:
+                raise ValueError("Gcra algorithm requires rate equal to 'refill_rate'")
+            rate = capacity / rule.window
+
+        if rate <= 0:
+            raise ValueError(f"rate must be positive, got {rate}")
+
+        now = time.time()
+        storage_key = f"{rule.name}:{key}"
+        allowed, remaining, reset_at = await storage.gcra(
+            key=storage_key, capacity=capacity, rate=rate, now=now
+        )
+        retry_after = max(0, math.ceil(reset_at - now))
+
+        return RateLimitResult(
+            allowed=allowed,
+            remaining=remaining,
+            reset_at=reset_at,
+            limit=capacity,
+            retry_after=retry_after if not allowed else None,
+            rule_name=rule.name,
+        )