This commit is contained in:
CarterPerez-dev 2025-12-28 11:31:36 -05:00
parent 43b2b5ca45
commit b34cc4a0db
35 changed files with 8182 additions and 0 deletions

View File

@ -0,0 +1,46 @@
[style]
based_on_style = pep8
column_limit = 75
indent_width = 4
continuation_indent_width = 4
indent_closing_brackets = false
dedent_closing_brackets = true
indent_blank_lines = false
spaces_before_comment = 2
spaces_around_power_operator = false
spaces_around_default_or_named_assign = true
space_between_ending_comma_and_closing_bracket = false
space_inside_brackets = false
spaces_around_subscript_colon = true
blank_line_before_nested_class_or_def = false
blank_line_before_class_docstring = false
blank_lines_around_top_level_definition = 2
blank_lines_between_top_level_imports_and_variables = 2
blank_line_before_module_docstring = false
split_before_logical_operator = true
split_before_first_argument = true
split_before_named_assigns = true
split_complex_comprehension = true
split_before_expression_after_opening_paren = false
split_before_closing_bracket = true
split_all_comma_separated_values = true
split_all_top_level_comma_separated_values = false
coalesce_brackets = false
each_dict_entry_on_separate_line = true
allow_multiline_lambdas = false
allow_multiline_dictionary_keys = false
split_penalty_import_names = 0
join_multiple_lines = false
align_closing_bracket_with_visual_indent = true
arithmetic_precedence_indication = false
split_penalty_for_added_line_split = 275
use_tabs = false
split_before_dot = false
split_arguments_when_comma_terminated = true
i18n_function_call = ['_', 'N_', 'gettext', 'ngettext']
i18n_comment = ['# Translators:', '# i18n:']
split_penalty_comprehension = 80
split_penalty_after_opening_bracket = 280
split_penalty_before_if_expr = 0
split_penalty_bitwise_operator = 290
split_penalty_logical_operator = 0

View File

View File

@ -0,0 +1,242 @@
[project]
name = "fastapi-420"
version = "0.1.0"
description = "Enhance Your Calm - Advanced Rate Limiting & DDoS Protection for FastAPI"
readme = "README.md"
license = {text = "MIT"}
requires-python = ">=3.12"
keywords = ["fastapi", "rate-limiting", "ddos", "security", "api", "420"]
classifiers = [
"Development Status :: 4 - Beta",
"Environment :: Web Environment",
"Framework :: FastAPI",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Topic :: Internet :: WWW/HTTP",
"Topic :: Security",
"Typing :: Typed",
]
dependencies = [
"fastapi[standard]>=0.123.0,<1.0.0",
"pydantic>=2.12.5,<3.0.0",
"pydantic-settings>=2.12.0,<3.0.0",
"redis>=7.1.0",
"pyjwt>=2.10.0",
]
[project.optional-dependencies]
dev = [
"pytest>=9.0.2",
"pytest-asyncio>=1.3.0",
"pytest-cov>=6.0.0",
"httpx>=0.28.1",
"fakeredis>=2.26.0",
"time-machine>=2.16.0",
"asgi-lifespan>=2.1.0",
"mypy>=1.19.0",
"types-redis>=4.6.0",
"ruff>=0.14.8",
"pylint>=4.0.4",
"pylint-pydantic>=0.4.1",
"pylint-per-file-ignores>=3.2.0",
"pre-commit>=4.2.0",
]
[project.urls]
Homepage = "https://github.com/CarterPerez-dev/fastapi-420"
Documentation = "https://github.com/CarterPerez-dev/fastapi-420#readme"
Repository = "https://github.com/CarterPerez-dev/fastapi-420"
Issues = "https://github.com/CarterPerez-dev/fastapi-420/issues"
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["src/fastapi_420"]
[tool.ruff]
target-version = "py312"
line-length = 88
src = ["src"]
exclude = ["alembic"]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"ARG", # flake8-unused-arguments
"SIM", # flake8-simplify
"PTH", # flake8-use-pathlib
"RUF", # ruff-specific
"ASYNC", # flake8-async
"S", # flake8-bandit (security)
"N", # pep8-naming
]
ignore = [
"E501", # line too long (formatter handles this)
"B008", # function call in default argument (FastAPI Depends)
"S101", # assert usage (needed for tests)
"S104", # 0.0.0.0 binding (intentional for Docker)
"S105", # "bearer" token_type is not a password
"ARG001", # unused function argument (common in FastAPI deps)
"E712", # == False is REQUIRED for SQLAlchemy WHERE clauses
"N999", # PascalCase module names (intentional: Base.py, User.py)
"N818", # exception naming convention (style preference)
"UP046", # Generic[T] syntax (keep for compatibility)
"RUF005", # list concatenation style (preference)
]
[tool.ruff.lint.per-file-ignores]
[tool.mypy]
python_version = "3.12"
strict = true
warn_return_any = true
warn_unused_ignores = true
disallow_untyped_defs = true
disallow_incomplete_defs = true
plugins = ["pydantic.mypy"]
exclude = ["alembic", ".venv", "venv"]
[[tool.mypy.overrides]]
module = ["tests.*", "conftest"]
ignore_errors = true
[[tool.mypy.overrides]]
module = ["core.logging"]
disable_error_code = ["no-any-return"]
[[tool.mypy.overrides]]
module = [
"uuid6",
"structlog",
"structlog.*",
"pwdlib",
"slowapi",
"slowapi.*",
]
ignore_missing_imports = true
[tool.pydantic-mypy]
init_forbid_extra = true
init_typed = true
warn_required_dynamic_aliases = true
[tool.pylint.main]
py-version = "3.11"
jobs = 4
load-plugins = [
"pylint_pydantic",
"pylint_per_file_ignores",
]
persistent = true
ignore = [
"alembic",
"venv",
".venv",
"__pycache__",
"build",
"dist",
".git",
".pytest_cache",
".mypy_cache",
".ruff_cache",
]
ignore-paths = [
"^alembic/.*",
"^venv/.*",
"^.venv/.*",
"^build/.*",
"^dist/.*",
]
[tool.pylint.messages_control]
disable = [
"C0103", # invalid-name
"C0116", # missing-function-docstring (we use minimal docs)
"C0121", # singleton-comparison (== False required for SQLAlchemy)
"C0301", # line-too-long
"C0302", # too-many-lines
"C0303", # trailing-whitespace
"C0304", # final-newline-missing
"C0305", # trailing-newlines
"C0411", # wrong-import-order
"C0412", # ungrouped-imports (style preference)
"E0401", # import-error (uuid6/structlog/pwdlib not found by pylint)
"E0611", # no-name-in-module (false positive for config re-exports)
"E1102", # not-callable (false positive for SQLAlchemy func.now/count)
"E1136", # unsubscriptable-object (false positive for generics)
"R0801", # similar-lines
"R0901", # too-many-ancestors (SQLAlchemy inheritance)
"R0903", # too-few-public-methods
"R0917", # too-many-positional-arguments (FastAPI patterns)
"W0611", # unused-import (handled by ruff, config.py re-exports)
"W0612", # unused-variable (handled by ruff)
"W0613", # unused-argument (handled by ruff)
"W0621", # redefined-outer-name (FastAPI route/param naming)
"W0622", # redefined-builtin
"W0718", # broad-exception-caught (intentional in health/rate-limit)
]
[tool.pylint.format]
max-line-length = 95
[tool.pylint.design]
max-args = 12
max-attributes = 10
max-branches = 15
max-locals = 20
max-statements = 55
[tool.pytest.ini_options]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
testpaths = ["tests"]
addopts = "-ra -q"
filterwarnings = [
"ignore::DeprecationWarning",
]
[tool.coverage.run]
branch = true
source = ["src"]
[tool.coverage.report]
exclude_lines = [
"pragma: no cover",
"if TYPE_CHECKING:",
"raise NotImplementedError",
]
[tool.ty.src]
include = ["src", "tests"]
exclude = ["alembic/versions/**", ".venv/**"]
respect-ignore-files = true
[tool.ty.environment]
python-version = "3.12"
root = ["./src"]
python = "./.venv"
[tool.ty.rules]
possibly-missing-attribute = "error"
possibly-missing-import = "error"
unused-ignore-comment = "warn"
redundant-cast = "warn"
undefined-reveal = "warn"

View File

@ -0,0 +1,67 @@
"""
AngelaMos | 2025
__init__.py
"""
from fastapi_420.config import RateLimiterSettings, get_settings
from fastapi_420.defense import CircuitBreaker, LayeredDefense
from fastapi_420.dependencies import (
LimiterDep,
RateLimitDep,
ScopedRateLimiter,
create_rate_limit_dep,
get_limiter,
require_rate_limit,
set_global_limiter,
)
from fastapi_420.exceptions import (
EnhanceYourCalm,
HTTP_420_ENHANCE_YOUR_CALM,
RateLimitError,
RateLimitExceeded,
StorageError,
)
from fastapi_420.limiter import RateLimiter
from fastapi_420.middleware import RateLimitMiddleware, SlowDownMiddleware
from fastapi_420.types import (
Algorithm,
DefenseMode,
FingerprintData,
FingerprintLevel,
Layer,
RateLimitResult,
RateLimitRule,
)
__version__ = "0.1.0"
__all__ = [
"HTTP_420_ENHANCE_YOUR_CALM",
"Algorithm",
"CircuitBreaker",
"DefenseMode",
"EnhanceYourCalm",
"FingerprintData",
"FingerprintLevel",
"Layer",
"LayeredDefense",
"LimiterDep",
"RateLimitDep",
"RateLimitError",
"RateLimitExceeded",
"RateLimitMiddleware",
"RateLimitResult",
"RateLimitRule",
"RateLimiter",
"RateLimiterSettings",
"ScopedRateLimiter",
"SlowDownMiddleware",
"StorageError",
"__version__",
"create_rate_limit_dep",
"get_limiter",
"get_settings",
"require_rate_limit",
"set_global_limiter",
]

View File

@ -0,0 +1,39 @@
"""
AngelaMos | 2025
__init__.py
"""
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
from fastapi_420.types import Algorithm
def create_algorithm(algorithm_type: Algorithm) -> BaseAlgorithm:
"""
Factory function to create appropriate algorithm instance
"""
algorithm_map: dict[Algorithm,
type[BaseAlgorithm]] = {
Algorithm.SLIDING_WINDOW:
SlidingWindowAlgorithm,
Algorithm.TOKEN_BUCKET: TokenBucketAlgorithm,
Algorithm.FIXED_WINDOW: FixedWindowAlgorithm,
Algorithm.LEAKY_BUCKET: SlidingWindowAlgorithm,
}
algorithm_class = algorithm_map.get(
algorithm_type,
SlidingWindowAlgorithm
)
return algorithm_class()
__all__ = [
"BaseAlgorithm",
"FixedWindowAlgorithm",
"SlidingWindowAlgorithm",
"TokenBucketAlgorithm",
"create_algorithm",
]

View File

@ -0,0 +1,52 @@
"""
AngelaMos | 2025
base.py
"""
# pylint: disable=unnecessary-ellipsis
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from fastapi_420.storage import Storage
from fastapi_420.types import RateLimitResult, RateLimitRule
class BaseAlgorithm(ABC):
"""
Abstract base class for rate limiting algorithms
"""
@property
@abstractmethod
def name(self) -> str:
"""
Algorithm name for logging and debugging
"""
...
@abstractmethod
async def check(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Check if request is allowed under rate limit
"""
...
@abstractmethod
async def get_current_usage(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
) -> int:
"""
Get current usage count without incrementing
"""
...

View File

@ -0,0 +1,76 @@
"""
AngelaMos | 2025
fixed_window.py
"""
from __future__ import annotations
import time
from typing import TYPE_CHECKING
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.storage.redis_backend import RedisStorage
from fastapi_420.types import Algorithm
if TYPE_CHECKING:
from fastapi_420.storage import Storage
from fastapi_420.types import RateLimitResult, RateLimitRule
class FixedWindowAlgorithm(BaseAlgorithm):
"""
Fixed window counter algorithm
Simple implementation but suffers from boundary burst problem:
clients can make 2x the limit by timing requests at window edges.
Use sliding_window for production unless simplicity is required.
"""
@property
def name(self) -> str:
return Algorithm.FIXED_WINDOW.value
async def check(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Check and increment counter using fixed window algorithm
"""
if isinstance(storage, RedisStorage):
return await storage.increment_fixed_window(
key = key,
window_seconds = rule.window_seconds,
limit = rule.requests,
timestamp = timestamp,
)
return await storage.increment(
key = key,
window_seconds = rule.window_seconds,
limit = rule.requests,
timestamp = timestamp,
)
async def get_current_usage(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
) -> int:
"""
Get current window count without incrementing
"""
now = time.time()
current_window = int(now // rule.window_seconds)
window_key = f"{key}:{current_window}"
state = await storage.get_window_state(
key = window_key,
window_seconds = rule.window_seconds,
)
return state.current_count

View File

@ -0,0 +1,66 @@
"""
AngelaMos | 2025
sliding_window.py
"""
from __future__ import annotations
import time
from typing import TYPE_CHECKING
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.types import Algorithm
if TYPE_CHECKING:
from fastapi_420.storage import Storage
from fastapi_420.types import RateLimitResult, RateLimitRule
class SlidingWindowAlgorithm(BaseAlgorithm):
"""
Sliding window counter algorithm
The recommended default for production rate limiting.
Achieves ~99.997% accuracy with O(1) memory per client.
Uses weighted interpolation between two fixed windows to
approximate a true sliding window.
"""
@property
def name(self) -> str:
return Algorithm.SLIDING_WINDOW.value
async def check(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Check and increment counter using sliding window algorithm
"""
return await storage.increment(
key = key,
window_seconds = rule.window_seconds,
limit = rule.requests,
timestamp = timestamp,
)
async def get_current_usage(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
) -> int:
"""
Get current weighted usage count without incrementing
"""
now = time.time()
elapsed_ratio = (now % rule.window_seconds) / rule.window_seconds
state = await storage.get_window_state(
key = key,
window_seconds = rule.window_seconds,
)
return int(state.weighted_count(elapsed_ratio))

View File

@ -0,0 +1,66 @@
"""
AngelaMos | 2025
token_bucket.py
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.types import Algorithm
if TYPE_CHECKING:
from fastapi_420.storage import Storage
from fastapi_420.types import RateLimitResult, RateLimitRule
class TokenBucketAlgorithm(BaseAlgorithm):
"""
Token bucket algorithm
Allows controlled bursting up to bucket capacity while
enforcing average rate limits. Tokens refill at a constant
rate and are consumed per request.
Best for APIs that need burst tolerance with eventual rate enforcement.
"""
@property
def name(self) -> str:
return Algorithm.TOKEN_BUCKET.value
async def check(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
timestamp: float | None = None, # noqa: ARG002
) -> RateLimitResult:
"""
Check and consume token using token bucket algorithm
"""
capacity = rule.requests
refill_rate = rule.requests / rule.window_seconds
return await storage.consume_token(
key = key,
capacity = capacity,
refill_rate = refill_rate,
tokens_to_consume = 1,
)
async def get_current_usage(
self,
storage: Storage,
key: str,
rule: RateLimitRule,
) -> int:
"""
Get current token count (inverted as usage)
"""
state = await storage.get_token_bucket_state(key = key)
if state is None:
return 0
return rule.requests - int(state.tokens)

View File

@ -0,0 +1,188 @@
"""
AngelaMos | 2025
config.py
"""
from __future__ import annotations
from functools import lru_cache
from typing import Annotated, Literal
from pydantic import (
Field,
RedisDsn,
model_validator,
)
from pydantic_settings import (
BaseSettings,
SettingsConfigDict,
)
from fastapi_420.types import (
Algorithm,
DefenseMode,
FingerprintLevel,
RateLimitRule,
)
class StorageSettings(BaseSettings):
"""
Storage backend configuration
"""
model_config = SettingsConfigDict(
env_prefix = "RATELIMIT_",
env_file = ".env",
env_file_encoding = "utf-8",
extra = "ignore",
)
REDIS_URL: RedisDsn | None = None
REDIS_MAX_CONNECTIONS: Annotated[int, Field(ge = 1, le = 1000)] = 100
REDIS_SOCKET_TIMEOUT: Annotated[float,
Field(ge = 0.1,
le = 60.0)] = 5.0
REDIS_RETRY_ON_TIMEOUT: bool = True
REDIS_DECODE_RESPONSES: bool = True
MEMORY_MAX_KEYS: Annotated[int,
Field(ge = 100,
le = 10_000_000)] = 100_000
MEMORY_CLEANUP_INTERVAL: Annotated[int, Field(ge = 1, le = 3600)] = 60
FALLBACK_TO_MEMORY: bool = True
class FingerprintSettings(BaseSettings):
"""
Fingerprinting configuration
"""
model_config = SettingsConfigDict(
env_prefix = "RATELIMIT_FP_",
env_file = ".env",
env_file_encoding = "utf-8",
extra = "ignore",
)
LEVEL: FingerprintLevel = FingerprintLevel.NORMAL
USE_IP: bool = True
USE_USER_AGENT: bool = True
USE_ACCEPT_HEADERS: bool = False
USE_HEADER_ORDER: bool = False
USE_AUTH: bool = True
USE_TLS: bool = False
USE_GEO: bool = False
IPV6_PREFIX_LENGTH: Annotated[int, Field(ge = 32, le = 128)] = 64
TRUSTED_PROXIES: list[str] = []
TRUST_X_FORWARDED_FOR: bool = False
class DefenseSettings(BaseSettings):
"""
DDoS defense layer configuration
"""
model_config = SettingsConfigDict(
env_prefix = "RATELIMIT_DEFENSE_",
env_file = ".env",
env_file_encoding = "utf-8",
extra = "ignore",
)
MODE: DefenseMode = DefenseMode.ADAPTIVE
GLOBAL_LIMIT: str = "50000/minute"
CIRCUIT_THRESHOLD: Annotated[int, Field(ge = 1)] = 10000
CIRCUIT_WINDOW: Annotated[int, Field(ge = 1, le = 3600)] = 60
CIRCUIT_RECOVERY_TIME: Annotated[int, Field(ge = 1, le = 3600)] = 30
ADAPTIVE_REDUCTION_FACTOR: Annotated[float,
Field(ge = 0.1,
le = 1.0)] = 0.5
ENDPOINT_LIMIT_MULTIPLIER: Annotated[int, Field(ge = 1, le = 100)] = 10
LOCKDOWN_ALLOW_AUTHENTICATED: bool = True
LOCKDOWN_ALLOW_KNOWN_GOOD: bool = True
@model_validator(mode = "after")
def validate_global_limit(self) -> DefenseSettings:
"""
Validate global limit can be parsed.
"""
RateLimitRule.parse(self.GLOBAL_LIMIT)
return self
class RateLimiterSettings(BaseSettings):
"""
Main rate limiter settings with environment variable support
"""
model_config = SettingsConfigDict(
env_prefix = "RATELIMIT_",
env_file = ".env",
env_file_encoding = "utf-8",
extra = "ignore",
)
ENABLED: bool = True
ALGORITHM: Algorithm = Algorithm.SLIDING_WINDOW
DEFAULT_LIMIT: str = "100/minute"
DEFAULT_LIMITS: list[str] = ["100/minute", "1000/hour"]
FAIL_OPEN: bool = True
KEY_PREFIX: str = "ratelimit"
KEY_VERSION: str = "v1"
INCLUDE_HEADERS: bool = True
LOG_VIOLATIONS: bool = True
ENVIRONMENT: Literal["development",
"staging",
"production"] = "development"
HTTP_420_MESSAGE: str = "Enhance your calm"
HTTP_420_DETAIL: str = "Rate limit exceeded. Take a breather."
endpoint_limits: dict[str, list[RateLimitRule]] = {}
storage: StorageSettings = StorageSettings()
fingerprint: FingerprintSettings = FingerprintSettings()
defense: DefenseSettings = DefenseSettings()
@model_validator(mode = "after")
def validate_limits(self) -> RateLimiterSettings:
"""
Validate all limit strings can be parsed
"""
RateLimitRule.parse(self.DEFAULT_LIMIT)
for limit in self.DEFAULT_LIMITS:
RateLimitRule.parse(limit)
return self
@model_validator(mode = "after")
def validate_production_settings(self) -> RateLimiterSettings:
"""
Enforce stricter settings in production
"""
if self.ENVIRONMENT == "production": # noqa: SIM102
if self.storage.REDIS_URL is None and not self.storage.FALLBACK_TO_MEMORY:
raise ValueError(
"Production requires Redis URL or FALLBACK_TO_MEMORY=True"
)
return self
def get_default_rules(self) -> list[RateLimitRule]:
"""
Parse and return default rate limit rules
"""
return [
RateLimitRule.parse(limit) for limit in self.DEFAULT_LIMITS
]
def get_global_limit_rule(self) -> RateLimitRule:
"""
Parse and return global defense limit rule
"""
return RateLimitRule.parse(self.defense.GLOBAL_LIMIT)
@lru_cache
def get_settings() -> RateLimiterSettings:
"""
Cached settings instance
"""
return RateLimiterSettings()

View File

@ -0,0 +1,14 @@
"""
AngelaMos | 2025
__init__.py
"""
from fastapi_420.defense.circuit_breaker import CircuitBreaker
from fastapi_420.defense.layers import LayeredDefense, LayerResult
__all__ = [
"CircuitBreaker",
"LayerResult",
"LayeredDefense",
]

View File

@ -0,0 +1,142 @@
"""
AngelaMos | 2025
circuit_breaker.py
"""
from __future__ import annotations
import time
import asyncio
import logging
from dataclasses import (
field,
dataclass,
)
from typing import TYPE_CHECKING
from fastapi_420.types import CircuitState, DefenseMode
if TYPE_CHECKING:
from fastapi_420.storage import Storage
logger = logging.getLogger("fastapi_420")
@dataclass
class CircuitBreaker:
"""
Global circuit breaker for API-wide protection
When request volume exceeds threshold, the circuit opens
and applies configured defense strategy.
"""
threshold: int = 10000
window_seconds: int = 60
recovery_time: int = 30
defense_mode: DefenseMode = DefenseMode.ADAPTIVE
_state: CircuitState = field(default_factory = CircuitState)
_lock: asyncio.Lock = field(default_factory = asyncio.Lock)
_counter_key: str = "circuit:global:requests"
async def check(self, storage: Storage) -> bool:
"""
Check if circuit is allowing requests
Returns True if requests should be allowed
"""
async with self._lock:
now = time.time()
if self._state.is_open:
if now - self._state.last_failure_time >= self.recovery_time:
await self._enter_half_open()
return True
return False
request_count = await self._get_request_count(storage)
self._state.total_requests_in_window = request_count
if request_count >= self.threshold:
await self._trip(now)
return False
return True
async def record_request(self, storage: Storage) -> None:
"""
Record a request in the circuit breaker counter
"""
now = time.time()
window = int(now // self.window_seconds)
key = f"{self._counter_key}:{window}"
await storage.increment(
key = key,
window_seconds = self.window_seconds,
limit = self.threshold * 10,
timestamp = now,
)
async def _get_request_count(self, storage: Storage) -> int:
"""
Get current request count in window
"""
now = time.time()
window = int(now // self.window_seconds)
key = f"{self._counter_key}:{window}"
state = await storage.get_window_state(
key = key,
window_seconds = self.window_seconds,
)
return state.current_count
async def _trip(self, now: float) -> None:
"""
Trip the circuit breaker
"""
self._state.is_open = True
self._state.last_failure_time = now
self._state.failure_count += 1
self._state.half_open_requests = 0
logger.warning(
"Circuit breaker tripped",
extra = {
"threshold": self.threshold,
"total_requests": self._state.total_requests_in_window,
"defense_mode": self.defense_mode.value,
},
)
async def _enter_half_open(self) -> None:
"""
Enter half-open state for recovery testing
"""
self._state.half_open_requests = 0
logger.info("Circuit breaker entering half-open state")
async def reset(self) -> None:
"""
Reset circuit breaker to closed state
"""
async with self._lock:
self._state = CircuitState()
logger.info("Circuit breaker reset to closed state")
@property
def is_open(self) -> bool:
"""
Check if circuit is open
"""
return self._state.is_open
@property
def current_state(self) -> CircuitState:
"""
Get current circuit state
"""
return self._state

View File

@ -0,0 +1,314 @@
"""
AngelaMos | 2025
layers.py
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
from fastapi_420.algorithms import create_algorithm
from fastapi_420.exceptions import EnhanceYourCalm
from fastapi_420.types import (
DefenseContext,
DefenseMode,
FingerprintData,
Layer,
RateLimitKey,
RateLimitResult,
RateLimitRule,
)
if TYPE_CHECKING:
from starlette.requests import Request
from fastapi_420.config import RateLimiterSettings
from fastapi_420.defense.circuit_breaker import CircuitBreaker
from fastapi_420.storage import Storage
logger = logging.getLogger("fastapi_420")
@dataclass
class LayerResult:
"""
Result from a defense layer check
"""
layer: Layer
allowed: bool
result: RateLimitResult
should_continue: bool = True
class LayeredDefense:
"""
Three-layer defense system for comprehensive rate limiting
Layer 1: Per-User Per-Endpoint (granular)
Layer 2: Per-Endpoint Global (endpoint protection)
Layer 3: Global Circuit Breaker (DDoS protection)
"""
def __init__(
self,
storage: Storage,
settings: RateLimiterSettings,
circuit_breaker: CircuitBreaker | None = None,
) -> None:
self._storage = storage
self._settings = settings
self._circuit_breaker = circuit_breaker
self._algorithm = create_algorithm(settings.ALGORITHM)
async def check_all_layers(
self,
request: Request,
fingerprint: FingerprintData,
endpoint: str,
rules: list[RateLimitRule],
) -> RateLimitResult:
"""
Check all defense layers in order
"""
context = DefenseContext(
fingerprint = fingerprint,
endpoint = endpoint,
method = request.method,
is_authenticated = fingerprint.auth_identifier is not None,
)
layer3_result = await self._check_layer3_global(context)
if not layer3_result.allowed:
self._log_violation(layer3_result, context)
raise EnhanceYourCalm(
result = layer3_result.result,
message = self._settings.HTTP_420_MESSAGE,
detail = "API is under heavy load. Please try again later.",
)
layer2_result = await self._check_layer2_endpoint(context, rules)
if not layer2_result.allowed:
self._log_violation(layer2_result, context)
raise EnhanceYourCalm(
result = layer2_result.result,
message = self._settings.HTTP_420_MESSAGE,
detail = self._settings.HTTP_420_DETAIL,
)
layer1_result = await self._check_layer1_user(context, rules)
if not layer1_result.allowed:
self._log_violation(layer1_result, context)
raise EnhanceYourCalm(
result = layer1_result.result,
message = self._settings.HTTP_420_MESSAGE,
detail = self._settings.HTTP_420_DETAIL,
)
return layer1_result.result
async def _check_layer1_user(
self,
context: DefenseContext,
rules: list[RateLimitRule],
) -> LayerResult:
"""
Layer 1: Per-user per-endpoint rate limiting
"""
identifier = context.fingerprint.to_composite_key(
self._settings.fingerprint.LEVEL
)
worst_result: RateLimitResult | None = None
for rule in rules:
key = RateLimitKey(
prefix = self._settings.KEY_PREFIX,
version = self._settings.KEY_VERSION,
layer = Layer.USER,
endpoint = context.endpoint,
identifier = identifier,
window = rule.window_seconds,
).build()
result = await self._algorithm.check(
storage = self._storage,
key = key,
rule = rule,
)
if not result.allowed: # noqa: SIM102
if worst_result is None or (result.retry_after or 0) > (
worst_result.retry_after or 0):
worst_result = result
if worst_result:
return LayerResult(
layer = Layer.USER,
allowed = False,
result = worst_result,
)
return LayerResult(
layer = Layer.USER,
allowed = True,
result = result,
)
async def _check_layer2_endpoint(
self,
context: DefenseContext,
rules: list[RateLimitRule], # noqa: ARG002
) -> LayerResult:
"""
Layer 2: Per-endpoint global rate limiting
"""
endpoint_rules = self._settings.endpoint_limits.get(
context.endpoint,
self._settings.get_default_rules()
)
for rule in endpoint_rules:
key = RateLimitKey(
prefix = self._settings.KEY_PREFIX,
version = self._settings.KEY_VERSION,
layer = Layer.ENDPOINT,
endpoint = context.endpoint,
identifier = "global",
window = rule.window_seconds,
).build()
endpoint_rule = RateLimitRule(
requests = rule.requests * self._settings.defense.ENDPOINT_LIMIT_MULTIPLIER,
window_seconds = rule.window_seconds,
)
result = await self._algorithm.check(
storage = self._storage,
key = key,
rule = endpoint_rule,
)
if not result.allowed:
return LayerResult(
layer = Layer.ENDPOINT,
allowed = False,
result = result,
)
return LayerResult(
layer = Layer.ENDPOINT,
allowed = True,
result = result,
)
async def _check_layer3_global(
self,
context: DefenseContext,
) -> LayerResult:
"""
Layer 3: Global circuit breaker
"""
if self._circuit_breaker is None:
return LayerResult(
layer = Layer.GLOBAL,
allowed = True,
result = RateLimitResult(
allowed = True,
limit = 0,
remaining = 0,
reset_after = 0,
),
)
is_allowed = await self._circuit_breaker.check(self._storage)
if not is_allowed:
if self._should_bypass_circuit(context):
return LayerResult(
layer = Layer.GLOBAL,
allowed = True,
result = RateLimitResult(
allowed = True,
limit = 0,
remaining = 0,
reset_after = 0,
),
)
return LayerResult(
layer = Layer.GLOBAL,
allowed = False,
result = RateLimitResult(
allowed = False,
limit = self._circuit_breaker.threshold,
remaining = 0,
reset_after = float(
self._circuit_breaker.recovery_time
),
retry_after = float(
self._circuit_breaker.recovery_time
),
),
)
await self._circuit_breaker.record_request(self._storage)
return LayerResult(
layer = Layer.GLOBAL,
allowed = True,
result = RateLimitResult(
allowed = True,
limit = self._circuit_breaker.threshold,
remaining = max(
0,
self._circuit_breaker.threshold - self._circuit_breaker
.current_state.total_requests_in_window,
),
reset_after = float(self._circuit_breaker.window_seconds),
),
)
def _should_bypass_circuit(self, context: DefenseContext) -> bool:
"""
Determine if request should bypass open circuit
"""
mode = self._settings.defense.MODE
if mode == DefenseMode.DISABLED:
return True
if mode == DefenseMode.LOCKDOWN:
if self._settings.defense.LOCKDOWN_ALLOW_AUTHENTICATED and context.is_authenticated:
return True
return self._settings.defense.LOCKDOWN_ALLOW_KNOWN_GOOD and context.reputation_score >= 0.9
if mode == DefenseMode.ADAPTIVE:
return bool(context.is_authenticated)
return False
def _log_violation(
self,
layer_result: LayerResult,
context: DefenseContext,
) -> None:
"""
Log rate limit violation
"""
if not self._settings.LOG_VIOLATIONS:
return
logger.warning(
"Rate limit exceeded",
extra = {
"layer": layer_result.layer.value,
"endpoint": context.endpoint,
"method": context.method,
"is_authenticated": context.is_authenticated,
"remaining": layer_result.result.remaining,
"reset_after": layer_result.result.reset_after,
},
)

View File

@ -0,0 +1,177 @@
"""
AngelaMos | 2025
dependencies.py
"""
from __future__ import annotations
from typing import Annotated
from collections.abc import Callable
from fastapi import Depends, Request
from fastapi_420.limiter import RateLimiter
from fastapi_420.types import RateLimitResult, RateLimitRule
_global_limiter: RateLimiter | None = None
def set_global_limiter(limiter: RateLimiter) -> None:
"""
Set the global rate limiter instance for dependency injection
"""
global _global_limiter # pylint: disable=global-statement
_global_limiter = limiter
def get_limiter() -> RateLimiter:
"""
Get the global rate limiter instance
"""
if _global_limiter is None:
raise RuntimeError(
"Rate limiter not initialized. "
"Call set_global_limiter() or use RateLimiterDep with explicit limiter."
)
return _global_limiter
class RateLimitDep:
"""
FastAPI dependency for rate limiting
Usage:
@app.get("/api/data", dependencies=[Depends(RateLimitDep("100/minute"))])
async def get_data():
return {"data": "value"}
# Or with result access:
@app.get("/api/data")
async def get_data(limit_result: Annotated[RateLimitResult, Depends(RateLimitDep("100/minute"))]):
return {"remaining": limit_result.remaining}
"""
def __init__(
self,
*rules: str,
limiter: RateLimiter | None = None,
key_func: Callable[[Request],
str] | None = None,
) -> None:
self.rules = [RateLimitRule.parse(rule) for rule in rules]
self._limiter = limiter
self.key_func = key_func
@property
def limiter(self) -> RateLimiter:
"""
Get limiter instance
"""
if self._limiter is not None:
return self._limiter
return get_limiter()
async def __call__(self, request: Request) -> RateLimitResult:
"""
Check rate limit and return result
"""
rule_strings = [str(rule) for rule in self.rules]
return await self.limiter.check(
request,
*rule_strings,
key_func = self.key_func,
raise_on_limit = True,
)
def create_rate_limit_dep(
*rules: str,
limiter: RateLimiter | None = None,
key_func: Callable[[Request],
str] | None = None,
) -> RateLimitDep:
"""
Factory function to create rate limit dependency
Usage:
rate_limit = create_rate_limit_dep("100/minute", "1000/hour")
@app.get("/api/data", dependencies=[Depends(rate_limit)])
async def get_data():
return {"data": "value"}
"""
return RateLimitDep(*rules, limiter = limiter, key_func = key_func)
LimiterDep = Annotated[RateLimiter, Depends(get_limiter)]
async def require_rate_limit(
request: Request,
limiter: LimiterDep,
) -> RateLimitResult:
"""
Dependency that applies default rate limits
Usage:
@app.get("/api/data")
async def get_data(
limit_result: Annotated[RateLimitResult, Depends(require_rate_limit)]
):
return {"remaining": limit_result.remaining}
"""
return await limiter.check(request, raise_on_limit = True)
class ScopedRateLimiter:
"""
Rate limiter scoped to specific endpoints or route groups
Usage:
api_limiter = ScopedRateLimiter(
prefix="/api/v1",
default_rules=["100/minute"],
endpoint_rules={
"POST:/api/v1/upload": ["10/minute"],
"POST:/api/v1/login": ["5/minute"],
}
)
@app.post("/api/v1/upload", dependencies=[Depends(api_limiter)])
async def upload():
return {"status": "ok"}
"""
def __init__(
self,
prefix: str = "",
default_rules: list[str] | None = None,
endpoint_rules: dict[str,
list[str]] | None = None,
limiter: RateLimiter | None = None,
) -> None:
self.prefix = prefix
self.default_rules = default_rules or ["100/minute"]
self.endpoint_rules = endpoint_rules or {}
self._limiter = limiter
@property
def limiter(self) -> RateLimiter:
"""
Get limiter instance
"""
if self._limiter is not None:
return self._limiter
return get_limiter()
async def __call__(self, request: Request) -> RateLimitResult:
"""
Apply appropriate rate limit based on endpoint
"""
endpoint = f"{request.method}:{request.url.path}"
rules = self.endpoint_rules.get(endpoint, self.default_rules)
return await self.limiter.check(
request,
*rules,
raise_on_limit = True,
)

View File

@ -0,0 +1,211 @@
"""
AngelaMos | 2025
exceptions.py
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from fastapi import HTTPException
from fastapi_420.types import Layer, StorageType
if TYPE_CHECKING:
from fastapi_420.types import RateLimitResult
HTTP_420_ENHANCE_YOUR_CALM = 420
class RateLimitError(Exception):
"""
Base exception for rate limiting errors
"""
def __init__(
self,
message: str,
details: dict[str,
Any] | None = None
) -> None:
self.message = message
self.details = details or {}
super().__init__(self.message)
class EnhanceYourCalm(HTTPException):
"""
HTTP 420 response - the signature rate limit response
"""
def __init__(
self,
result: RateLimitResult | None = None,
message: str = "Enhance your calm",
detail: str = "Rate limit exceeded. Take a breather.",
headers: dict[str,
str] | None = None,
) -> None:
final_headers = {}
if result:
final_headers.update(result.headers)
if headers:
final_headers.update(headers)
super().__init__(
status_code = HTTP_420_ENHANCE_YOUR_CALM,
detail = {
"message": message,
"detail": detail,
"limit_info": result.headers if result else {}
},
headers = final_headers if final_headers else None,
)
self.result = result
class RateLimitExceeded(RateLimitError):
"""
Raised when a rate limit is exceeded at any layer
"""
def __init__(
self,
result: RateLimitResult,
layer: Layer = Layer.USER,
endpoint: str = "",
identifier: str = "",
) -> None:
self.result = result
self.layer = layer
self.endpoint = endpoint
self.identifier = identifier
super().__init__(
message = f"Rate limit exceeded on {layer.value} layer",
details = {
"layer": layer.value,
"endpoint": endpoint,
"remaining": result.remaining,
"reset_after": result.reset_after,
},
)
class StorageError(RateLimitError):
"""
Raised when storage backend operations fail
"""
def __init__(
self,
operation: str,
backend: StorageType | None = None,
original_error: Exception | None = None,
) -> None:
self.operation = operation
self.backend = backend
self.original_error = original_error
backend_name = backend.value if backend else "unknown"
super().__init__(
message =
f"Storage operation '{operation}' failed on {backend_name}",
details = {
"operation": operation,
"backend": backend_name,
"error": str(original_error) if original_error else None,
},
)
class StorageConnectionError(StorageError):
"""
Raised when unable to connect to storage backend
"""
def __init__(
self,
backend: StorageType,
original_error: Exception | None = None,
) -> None:
super().__init__(
operation = "connect",
backend = backend,
original_error = original_error,
)
class StorageUnavailable(StorageError):
"""
Raised when storage backend is temporarily unavailable
"""
def __init__(
self,
backend: StorageType,
original_error: Exception | None = None,
) -> None:
super().__init__(
operation = "health_check",
backend = backend,
original_error = original_error,
)
class FingerprintError(RateLimitError):
"""
Raised when fingerprint extraction fails
"""
def __init__(
self,
reason: str,
original_error: Exception | None = None,
) -> None:
self.reason = reason
self.original_error = original_error
super().__init__(
message = f"Fingerprint extraction failed: {reason}",
details = {
"reason": reason,
"error": str(original_error) if original_error else None,
},
)
class CircuitBreakerOpen(RateLimitError):
"""
Raised when global circuit breaker is open
"""
def __init__(
self,
recovery_time: float,
total_requests: int,
threshold: int,
) -> None:
self.recovery_time = recovery_time
self.total_requests = total_requests
self.threshold = threshold
super().__init__(
message = "Circuit breaker is open - API is in defense mode",
details = {
"recovery_time": recovery_time,
"total_requests": total_requests,
"threshold": threshold,
},
)
class ConfigurationError(RateLimitError):
"""
Raised when rate limiter configuration is invalid
"""
def __init__(self, reason: str) -> None:
super().__init__(
message = f"Invalid configuration: {reason}",
details = {"reason": reason},
)
class InvalidRuleError(ConfigurationError):
"""
Raised when a rate limit rule string is invalid
"""
def __init__(self, rule_string: str, reason: str) -> None:
self.rule_string = rule_string
super().__init__(
reason = f"Invalid rule '{rule_string}': {reason}"
)

View File

@ -0,0 +1,17 @@
"""
AngelaMos | 2025
__init__.py
"""
from fastapi_420.fingerprinting.auth import AuthExtractor
from fastapi_420.fingerprinting.composite import CompositeFingerprinter
from fastapi_420.fingerprinting.headers import HeadersExtractor
from fastapi_420.fingerprinting.ip import IPExtractor
__all__ = [
"AuthExtractor",
"CompositeFingerprinter",
"HeadersExtractor",
"IPExtractor",
]

View File

@ -0,0 +1,138 @@
"""
AngelaMos | 2025
auth.py
"""
from __future__ import annotations
import jwt
import json
import base64
import hashlib
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from starlette.requests import Request
class AuthExtractor:
"""
Extract authentication identifiers from requests
Supports JWT tokens, API keys in headers/query params,
and session-based authentication.
"""
def __init__(
self,
jwt_secret: str | None = None,
jwt_algorithms: list[str] | None = None,
api_key_header: str = "X-API-Key",
api_key_query_param: str = "api_key",
session_cookie: str = "session_id",
hash_identifiers: bool = True,
hash_length: int = 16,
) -> None:
self.jwt_secret = jwt_secret
self.jwt_algorithms = jwt_algorithms or ["HS256"]
self.api_key_header = api_key_header
self.api_key_query_param = api_key_query_param
self.session_cookie = session_cookie
self.hash_identifiers = hash_identifiers
self.hash_length = hash_length
def extract(self, request: Request) -> str | None:
"""
Extract authentication identifier using fallback chain
Order: JWT -> API Key (header) -> API Key (query) -> Session -> None
"""
identifier = (
self._extract_jwt_subject(request)
or self._extract_api_key_header(request)
or self._extract_api_key_query(request)
or self._extract_session(request)
)
if identifier and self.hash_identifiers:
return self._hash_identifier(identifier)
return identifier
def _extract_jwt_subject(self, request: Request) -> str | None:
"""
Extract subject claim from JWT Bearer token
"""
auth_header = request.headers.get("Authorization", "")
if not auth_header.startswith("Bearer "):
return None
token = auth_header[7 :]
if not self.jwt_secret:
return self._extract_jwt_subject_unsafe(token)
try:
payload = jwt.decode(
token,
self.jwt_secret,
algorithms = self.jwt_algorithms,
)
return str(payload.get("sub", ""))
except Exception:
return None
def _extract_jwt_subject_unsafe(self, token: str) -> str | None:
"""
Extract subject from JWT without verification
Used when jwt_secret is not configured - only extracts
the claim for rate limiting purposes, does NOT validate.
"""
try:
parts = token.split(".")
if len(parts) != 3:
return None
payload_b64 = parts[1]
padding = 4 - len(payload_b64) % 4
if padding != 4:
payload_b64 += "=" * padding
payload_bytes = base64.urlsafe_b64decode(payload_b64)
payload = json.loads(payload_bytes)
return str(payload.get("sub", ""))
except Exception:
return None
def _extract_api_key_header(self, request: Request) -> str | None:
"""
Extract API key from header
"""
return request.headers.get(self.api_key_header)
def _extract_api_key_query(self, request: Request) -> str | None:
"""
Extract API key from query parameter
"""
return request.query_params.get(self.api_key_query_param)
def _extract_session(self, request: Request) -> str | None:
"""
Extract session ID from cookie
"""
return request.cookies.get(self.session_cookie)
def _hash_identifier(self, identifier: str) -> str:
"""
Hash identifier for privacy
"""
hash_bytes = hashlib.sha256(identifier.encode()).hexdigest()
return hash_bytes[: self.hash_length]
def is_authenticated(self, request: Request) -> bool:
"""
Check if request has any authentication
"""
return self.extract(request) is not None

View File

@ -0,0 +1,198 @@
"""
AngelaMos | 2025
composite.py
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from fastapi_420.fingerprinting.auth import AuthExtractor
from fastapi_420.fingerprinting.headers import HeadersExtractor
from fastapi_420.fingerprinting.ip import IPExtractor
from fastapi_420.types import FingerprintData, FingerprintLevel
if TYPE_CHECKING:
from starlette.requests import Request
from fastapi_420.config import FingerprintSettings
class CompositeFingerprinter: # pylint: disable=too-many-instance-attributes
"""
Combines multiple fingerprinting methods based on configuration
Preset levels determine which extractors are used:
- strict: All methods (IP, headers, auth, TLS, geo)
- normal: IP + User-Agent + Auth (default)
- relaxed: IP + Auth only
- custom: Configured via settings
"""
def __init__(
self,
level: FingerprintLevel = FingerprintLevel.NORMAL,
ip_extractor: IPExtractor | None = None,
headers_extractor: HeadersExtractor | None = None,
auth_extractor: AuthExtractor | None = None,
use_ip: bool = True,
use_user_agent: bool = True,
use_accept_headers: bool = False,
use_header_order: bool = False,
use_auth: bool = True,
use_tls: bool = False,
use_geo: bool = False,
) -> None:
self.level = level
self._ip_extractor = ip_extractor or IPExtractor()
self._headers_extractor = headers_extractor or HeadersExtractor(
use_header_order = use_header_order,
)
self._auth_extractor = auth_extractor or AuthExtractor()
if level == FingerprintLevel.STRICT:
self.use_ip = True
self.use_user_agent = True
self.use_accept_headers = True
self.use_header_order = True
self.use_auth = True
self.use_tls = True
self.use_geo = True
elif level == FingerprintLevel.RELAXED:
self.use_ip = True
self.use_user_agent = False
self.use_accept_headers = False
self.use_header_order = False
self.use_auth = True
self.use_tls = False
self.use_geo = False
elif level == FingerprintLevel.CUSTOM:
self.use_ip = use_ip
self.use_user_agent = use_user_agent
self.use_accept_headers = use_accept_headers
self.use_header_order = use_header_order
self.use_auth = use_auth
self.use_tls = use_tls
self.use_geo = use_geo
else:
self.use_ip = True
self.use_user_agent = True
self.use_accept_headers = False
self.use_header_order = False
self.use_auth = True
self.use_tls = False
self.use_geo = False
@classmethod
def from_settings(
cls,
settings: FingerprintSettings
) -> CompositeFingerprinter:
"""
Create fingerprinter from settings
"""
ip_extractor = IPExtractor(
ipv6_prefix_length = settings.IPV6_PREFIX_LENGTH,
trusted_proxies = settings.TRUSTED_PROXIES,
trust_x_forwarded_for = settings.TRUST_X_FORWARDED_FOR,
)
headers_extractor = HeadersExtractor(
use_header_order = settings.USE_HEADER_ORDER,
)
return cls(
level = settings.LEVEL,
ip_extractor = ip_extractor,
headers_extractor = headers_extractor,
use_ip = settings.USE_IP,
use_user_agent = settings.USE_USER_AGENT,
use_accept_headers = settings.USE_ACCEPT_HEADERS,
use_header_order = settings.USE_HEADER_ORDER,
use_auth = settings.USE_AUTH,
use_tls = settings.USE_TLS,
use_geo = settings.USE_GEO,
)
async def extract(self, request: Request) -> FingerprintData:
"""
Extract fingerprint data from request
"""
raw_ip = ""
normalized_ip = ""
user_agent = None
accept_language = None
accept_encoding = None
headers_hash = None
auth_identifier = None
tls_fingerprint = None
geo_asn = None
if self.use_ip:
raw_ip, normalized_ip = self._ip_extractor.extract(request)
if self.use_user_agent:
user_agent = self._headers_extractor.extract_user_agent(
request
)
if self.use_accept_headers:
accept_language = self._headers_extractor.extract_accept_language(
request
)
accept_encoding = self._headers_extractor.extract_accept_encoding(
request
)
if self.use_header_order:
headers_hash = self._headers_extractor.compute_headers_hash(
request
)
if self.use_auth:
auth_identifier = self._auth_extractor.extract(request)
if self.use_tls:
tls_fingerprint = self._extract_tls_fingerprint(request)
if self.use_geo:
geo_asn = self._extract_geo_asn(request)
return FingerprintData(
ip = raw_ip,
ip_normalized = normalized_ip,
user_agent = user_agent,
accept_language = accept_language,
accept_encoding = accept_encoding,
headers_hash = headers_hash,
auth_identifier = auth_identifier,
tls_fingerprint = tls_fingerprint,
geo_asn = geo_asn,
)
def _extract_tls_fingerprint(self, request: Request) -> str | None:
"""
Extract TLS/JA3 fingerprint if available
Requires proxy to pass fingerprint in header
"""
return (
request.headers.get("X-JA3-Fingerprint")
or request.headers.get("X-TLS-Fingerprint")
)
def _extract_geo_asn(self, request: Request) -> str | None:
"""
Extract geographic ASN if available
Requires geo lookup service to populate header
"""
return (
request.headers.get("X-Client-ASN")
or request.headers.get("CF-IPCountry")
)
def is_authenticated(self, request: Request) -> bool:
"""
Check if request has authentication
"""
return self._auth_extractor.is_authenticated(request)

View File

@ -0,0 +1,95 @@
"""
AngelaMos | 2025
headers.py
"""
from __future__ import annotations
import hashlib
from typing import TYPE_CHECKING, ClassVar
if TYPE_CHECKING:
from starlette.requests import Request
class HeadersExtractor:
"""
Extract fingerprint data from HTTP headers
Header ordering is browser-specific and not user-configurable,
making it useful for fingerprinting even when other headers are spoofed.
"""
FINGERPRINT_HEADERS: ClassVar[list[str]] = [
"user-agent",
"accept",
"accept-language",
"accept-encoding",
"connection",
"upgrade-insecure-requests",
"sec-fetch-site",
"sec-fetch-mode",
"sec-fetch-user",
"sec-fetch-dest",
"sec-ch-ua",
"sec-ch-ua-mobile",
"sec-ch-ua-platform",
]
def __init__(
self,
use_header_order: bool = False,
hash_length: int = 16,
) -> None:
self.use_header_order = use_header_order
self.hash_length = hash_length
def extract_user_agent(self, request: Request) -> str | None:
"""
Extract User-Agent header
"""
return request.headers.get("user-agent")
def extract_accept_language(self, request: Request) -> str | None:
"""
Extract Accept-Language header
"""
return request.headers.get("accept-language")
def extract_accept_encoding(self, request: Request) -> str | None:
"""
Extract Accept-Encoding header
"""
return request.headers.get("accept-encoding")
def compute_headers_hash(self, request: Request) -> str:
"""
Compute hash of fingerprint-relevant headers
Includes header ordering if configured, which is
browser-specific and harder to spoof.
"""
components: list[str] = []
if self.use_header_order:
header_keys = [k.lower() for k in request.headers]
components.append("|".join(header_keys))
for header_name in self.FINGERPRINT_HEADERS:
value = request.headers.get(header_name, "")
components.append(f"{header_name}={value}")
fingerprint_string = "\n".join(components)
hash_bytes = hashlib.sha256(fingerprint_string.encode()
).hexdigest()
return hash_bytes[: self.hash_length]
def extract_all(self, request: Request) -> dict[str, str | None]:
"""
Extract all header-based fingerprint data
"""
return {
"user_agent": self.extract_user_agent(request),
"accept_language": self.extract_accept_language(request),
"accept_encoding": self.extract_accept_encoding(request),
"headers_hash": self.compute_headers_hash(request),
}

View File

@ -0,0 +1,134 @@
"""
AngelaMos | 2025
ip.py
"""
from __future__ import annotations
from ipaddress import (
IPv6Address,
ip_address,
ip_network,
)
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from starlette.requests import Request
class IPExtractor:
"""
Extract and normalize client IP addresses
Handles IPv6 /64 prefix normalization since users control
entire /64 prefixes (~18 quintillion addresses), making
per-IP limits trivially bypassable without normalization.
"""
def __init__(
self,
ipv6_prefix_length: int = 64,
trusted_proxies: list[str] | None = None,
trust_x_forwarded_for: bool = False,
) -> None:
self.ipv6_prefix_length = ipv6_prefix_length
self.trusted_proxies = set(trusted_proxies or [])
self.trust_x_forwarded_for = trust_x_forwarded_for
def extract(self, request: Request) -> tuple[str, str]:
"""
Extract raw IP and normalized IP from request
Returns tuple of (raw_ip, normalized_ip)
"""
raw_ip = self._get_client_ip(request)
normalized_ip = self._normalize_ip(raw_ip)
return raw_ip, normalized_ip
def _get_client_ip(self, request: Request) -> str:
"""
Get client IP address, handling proxy headers if configured
"""
if self.trust_x_forwarded_for:
forwarded_for = request.headers.get("X-Forwarded-For")
if forwarded_for:
return self._parse_x_forwarded_for(forwarded_for, request)
x_real_ip = request.headers.get("X-Real-IP")
if x_real_ip and self._is_trusted_proxy(
request.client.host if request.client else ""):
return x_real_ip.strip()
if request.client:
return request.client.host
return "0.0.0.0"
def _parse_x_forwarded_for(self, header: str, request: Request) -> str:
"""
Parse X-Forwarded-For header safely
Uses rightmost-trusted approach: walk backwards through
the chain and return the first IP not in trusted proxies
"""
ips = [ip.strip() for ip in header.split(",")]
if not self.trusted_proxies:
return ips[0]
for ip in reversed(ips):
if ip not in self.trusted_proxies:
return ip
return ips[0] if ips else (
request.client.host if request.client else "0.0.0.0"
)
def _is_trusted_proxy(self, ip: str) -> bool:
"""
Check if IP is in trusted proxies list
"""
return ip in self.trusted_proxies
def _normalize_ip(self, ip_str: str) -> str:
"""
Normalize IP address for rate limiting
IPv6 addresses are normalized to their /64 network prefix
since users typically control entire /64 blocks.
"""
try:
addr = ip_address(ip_str)
except ValueError:
return ip_str
if isinstance(addr, IPv6Address):
if addr.ipv4_mapped:
return str(addr.ipv4_mapped)
network = ip_network(
f"{ip_str}/{self.ipv6_prefix_length}",
strict = False,
)
return str(network.network_address)
return str(addr)
def is_ipv6(self, ip_str: str) -> bool:
"""
Check if IP string is IPv6
"""
try:
addr = ip_address(ip_str)
return isinstance(addr, IPv6Address)
except ValueError:
return False
def is_private(self, ip_str: str) -> bool:
"""
Check if IP is a private/internal address
"""
try:
addr = ip_address(ip_str)
return addr.is_private
except ValueError:
return False

View File

@ -0,0 +1,347 @@
"""
AngelaMos | 2025
limiter.py
"""
from __future__ import annotations
import asyncio
import logging
import functools
from typing import (
Any,
ParamSpec,
TypeVar,
TYPE_CHECKING,
)
from collections.abc import Callable
from starlette.requests import Request
from fastapi_420.algorithms import (
create_algorithm,
)
from fastapi_420.config import (
RateLimiterSettings,
get_settings,
)
from fastapi_420.exceptions import (
EnhanceYourCalm,
StorageError,
)
from fastapi_420.fingerprinting import (
CompositeFingerprinter,
)
from fastapi_420.storage import (
MemoryStorage,
RedisStorage,
create_storage,
)
from fastapi_420.types import (
Layer,
RateLimitKey,
RateLimitResult,
RateLimitRule,
)
if TYPE_CHECKING:
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.storage import Storage
logger = logging.getLogger("fastapi_420")
P = ParamSpec("P")
R = TypeVar("R")
class RateLimiter:
"""
Main rate limiter class for FastAPI applications.
Usage:
limiter = RateLimiter()
@app.get("/api/data")
@limiter.limit("100/minute", "1000/hour")
async def get_data(request: Request):
return {"data": "value"}
"""
def __init__(
self,
settings: RateLimiterSettings | None = None,
storage: Storage | None = None,
) -> None:
self._settings = settings or get_settings()
self._storage = storage
self._fallback_storage: MemoryStorage | None = None
self._algorithm: BaseAlgorithm | None = None
self._fingerprinter: CompositeFingerprinter | None = None
self._initialized = False
self._lock = asyncio.Lock()
async def init(self) -> None:
"""
Initialize storage, algorithm, and fingerprinter
"""
async with self._lock:
if self._initialized:
return
if self._storage is None:
self._storage = create_storage(self._settings.storage)
if isinstance(self._storage, RedisStorage):
await self._storage.connect()
if isinstance(self._storage, MemoryStorage):
await self._storage.start_cleanup_task()
if self._settings.storage.FALLBACK_TO_MEMORY:
self._fallback_storage = MemoryStorage.from_settings(
self._settings.storage
)
await self._fallback_storage.start_cleanup_task()
self._algorithm = create_algorithm(self._settings.ALGORITHM)
self._fingerprinter = CompositeFingerprinter.from_settings(
self._settings.fingerprint
)
self._initialized = True
logger.info(
"Rate limiter initialized",
extra = {
"algorithm":
self._settings.ALGORITHM.value,
"storage":
self._storage.storage_type.value,
"fingerprint_level":
self._settings.fingerprint.LEVEL.value,
},
)
async def close(self) -> None:
"""
Close storage connections
"""
if self._storage:
await self._storage.close()
if self._fallback_storage:
await self._fallback_storage.close()
self._initialized = False
def limit(
self,
*rules: str,
key_func: Callable[[Request],
str] | None = None,
) -> Callable[[Callable[P,
R]],
Callable[P,
R]]:
"""
Decorator to apply rate limits to an endpoint
Args:
rules: Rate limit strings like "100/minute", "1000/hour"
key_func: Optional custom function to generate rate limit key
"""
parsed_rules = [RateLimitRule.parse(rule) for rule in rules]
def decorator(func: Callable[P, R]) -> Callable[P, R]:
@functools.wraps(func)
async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
request = self._extract_request(args, kwargs)
if request is None:
return await func(*args, **kwargs) # type: ignore[misc, no-any-return]
await self._check_rate_limits(
request = request,
rules = parsed_rules,
key_func = key_func,
)
return await func(*args, **kwargs) # type: ignore[misc, no-any-return]
return wrapper # type: ignore[return-value]
return decorator
async def check(
self,
request: Request,
*rules: str,
key_func: Callable[[Request],
str] | None = None,
raise_on_limit: bool = True,
) -> RateLimitResult:
"""
Manually check rate limit without decorator
Returns the result of the strictest (most restrictive) rule
"""
parsed_rules = [RateLimitRule.parse(rule) for rule in rules]
if not parsed_rules:
parsed_rules = self._settings.get_default_rules()
return await self._check_rate_limits(
request = request,
rules = parsed_rules,
key_func = key_func,
raise_on_limit = raise_on_limit,
)
async def _check_rate_limits(
self,
request: Request,
rules: list[RateLimitRule],
key_func: Callable[[Request],
str] | None = None,
raise_on_limit: bool = True,
) -> RateLimitResult:
"""
Check all rules and return/raise for the most restrictive failure
"""
if not self._initialized:
await self.init()
storage = await self._get_active_storage()
if storage is None:
if self._settings.FAIL_OPEN:
return RateLimitResult(
allowed = True,
limit = 0,
remaining = 0,
reset_after = 0,
)
raise StorageError(operation = "check", backend = None)
fingerprint = await self._fingerprinter.extract(request) # type: ignore[union-attr]
endpoint = self._get_endpoint(request)
if key_func:
identifier = key_func(request)
else:
identifier = fingerprint.to_composite_key(
self._settings.fingerprint.LEVEL
)
worst_result: RateLimitResult | None = None
for rule in rules:
key = RateLimitKey(
prefix = self._settings.KEY_PREFIX,
version = self._settings.KEY_VERSION,
layer = Layer.USER,
endpoint = endpoint,
identifier = identifier,
window = rule.window_seconds,
).build()
result = await self._algorithm.check( # type: ignore[union-attr]
storage = storage,
key = key,
rule = rule,
)
if not result.allowed: # noqa: SIM102
if worst_result is None or result.retry_after > (worst_result.retry_after or 0): # type: ignore[operator]
worst_result = result
if worst_result is not None:
if self._settings.LOG_VIOLATIONS:
logger.warning(
"Rate limit exceeded",
extra = {
"endpoint": endpoint,
"identifier": identifier[: 16],
"remaining": worst_result.remaining,
"reset_after": worst_result.reset_after,
},
)
if raise_on_limit:
raise EnhanceYourCalm(
result = worst_result,
message = self._settings.HTTP_420_MESSAGE,
detail = self._settings.HTTP_420_DETAIL,
)
return worst_result
best_result = result
return best_result
async def _get_active_storage(self) -> Storage | None:
"""
Get active storage, falling back to memory if primary fails
"""
if self._storage is None:
return self._fallback_storage
try:
is_healthy = await self._storage.health_check()
if is_healthy:
return self._storage
except Exception: # noqa: S110
pass
if self._fallback_storage:
logger.warning(
"Primary storage unavailable, using memory fallback",
extra = {
"primary_storage": self._storage.storage_type.value
},
)
return self._fallback_storage
return None
def _extract_request(
self,
args: tuple[Any,
...],
kwargs: dict[str,
Any],
) -> Request | None:
"""
Extract Request object from function arguments
"""
for arg in args:
if isinstance(arg, Request):
return arg
for value in kwargs.values():
if isinstance(value, Request):
return value
return None
def _get_endpoint(self, request: Request) -> str:
"""
Get endpoint identifier from request
"""
route = request.scope.get("route")
if route:
return f"{request.method}:{route.path}"
return f"{request.method}:{request.url.path}"
@property
def settings(self) -> RateLimiterSettings:
"""
Get current settings
"""
return self._settings
@property
def is_initialized(self) -> bool:
"""
Check if limiter is initialized
"""
return self._initialized

View File

@ -0,0 +1,210 @@
"""
AngelaMos | 2025
middleware.py
"""
from __future__ import annotations
import re
import logging
import asyncio
from typing import (
TYPE_CHECKING,
)
from collections.abc import Callable
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import (
JSONResponse,
Response,
)
from fastapi_420.exceptions import (
HTTP_420_ENHANCE_YOUR_CALM,
EnhanceYourCalm,
)
from fastapi_420.limiter import RateLimiter
if TYPE_CHECKING:
from starlette.types import ASGIApp
logger = logging.getLogger("fastapi_420")
class RateLimitMiddleware(BaseHTTPMiddleware):
"""
ASGI middleware for automatic rate limiting on all routes
Usage:
from fastapi import FastAPI
from fastapi_420.middleware import RateLimitMiddleware
from fastapi_420.limiter import RateLimiter
app = FastAPI()
limiter = RateLimiter()
app.add_middleware(
RateLimitMiddleware,
limiter=limiter,
default_limit="100/minute",
)
"""
def __init__(
self,
app: ASGIApp,
limiter: RateLimiter,
default_limit: str = "100/minute",
exclude_paths: list[str] | None = None,
exclude_patterns: list[str] | None = None,
include_paths: list[str] | None = None,
path_limits: dict[str,
str] | None = None,
key_func: Callable[[Request],
str] | None = None,
) -> None:
super().__init__(app)
self.limiter = limiter
self.default_limit = default_limit
self.exclude_paths = set(exclude_paths or [])
self.exclude_patterns = [
re.compile(p) for p in (exclude_patterns or [])
]
self.include_paths = set(include_paths) if include_paths else None
self.path_limits = path_limits or {}
self.key_func = key_func
self.exclude_paths.add("/health")
self.exclude_paths.add("/healthz")
self.exclude_paths.add("/ready")
self.exclude_paths.add("/metrics")
async def dispatch(self, request: Request, call_next: Callable) -> Response: # type: ignore[type-arg]
"""
Process request and apply rate limiting
"""
if not await self._should_limit(request):
return await call_next(request) # type: ignore[no-any-return]
try:
limit = self._get_limit_for_path(request.url.path)
await self.limiter.check(
request,
limit,
key_func = self.key_func,
raise_on_limit = True,
)
response = await call_next(request)
if self.limiter.settings.INCLUDE_HEADERS:
result = await self.limiter.check(
request,
limit,
key_func = self.key_func,
raise_on_limit = False,
)
for header_name, header_value in result.headers.items():
response.headers[header_name] = header_value
return response # type: ignore[no-any-return]
except EnhanceYourCalm as exc:
return self._create_420_response(exc)
async def _should_limit(self, request: Request) -> bool:
"""
Determine if request should be rate limited
"""
path = request.url.path
if path in self.exclude_paths:
return False
for pattern in self.exclude_patterns:
if pattern.match(path):
return False
if self.include_paths is not None: # noqa: SIM102
if path not in self.include_paths:
for include_path in self.include_paths:
if path.startswith(include_path):
break
else:
return False
return True
def _get_limit_for_path(self, path: str) -> str:
"""
Get rate limit for specific path
"""
if path in self.path_limits:
return self.path_limits[path]
for pattern_path, limit in self.path_limits.items():
if path.startswith(pattern_path):
return limit
return self.default_limit
def _create_420_response(self, exc: EnhanceYourCalm) -> JSONResponse:
"""
Create HTTP 420 response
"""
headers = {}
if exc.result:
headers.update(exc.result.headers)
return JSONResponse(
status_code = HTTP_420_ENHANCE_YOUR_CALM,
content = exc.detail,
headers = headers,
)
class SlowDownMiddleware(BaseHTTPMiddleware):
"""
Alternative middleware that adds delays instead of blocking
Useful for gradual throttling rather than hard limits
"""
def __init__(
self,
app: ASGIApp,
limiter: RateLimiter,
threshold_limit: str = "50/minute",
max_delay_seconds: float = 5.0,
delay_increment: float = 0.5,
) -> None:
super().__init__(app)
self.limiter = limiter
self.threshold_limit = threshold_limit
self.max_delay_seconds = max_delay_seconds
self.delay_increment = delay_increment
async def dispatch(self, request: Request, call_next: Callable) -> Response: # type: ignore[type-arg]
"""
Process request with potential delays
"""
result = await self.limiter.check(
request,
self.threshold_limit,
raise_on_limit = False,
)
if result.remaining <= 0:
delay = min(
self.max_delay_seconds,
result.retry_after or self.delay_increment,
)
await asyncio.sleep(delay)
response = await call_next(request)
if self.limiter.settings.INCLUDE_HEADERS:
for header_name, header_value in result.headers.items():
response.headers[header_name] = header_value
return response # type: ignore[no-any-return]

View File

@ -0,0 +1,35 @@
"""
AngelaMos | 2025
__init__.py
"""
from __future__ import annotations
from typing import TYPE_CHECKING, TypeAlias
from fastapi_420.storage.memory import MemoryStorage
from fastapi_420.storage.redis_backend import RedisStorage
from fastapi_420.types import StorageType
if TYPE_CHECKING:
from fastapi_420.config import StorageSettings
Storage: TypeAlias = MemoryStorage | RedisStorage # noqa: UP040
def create_storage(settings: StorageSettings) -> Storage:
"""
Factory function to create appropriate storage backend.
"""
if settings.REDIS_URL is not None:
return RedisStorage.from_settings(settings)
return MemoryStorage.from_settings(settings)
__all__ = [
"MemoryStorage",
"RedisStorage",
"Storage",
"StorageType",
"create_storage",
]

View File

@ -0,0 +1,32 @@
--[[
AngelaMos | 2025
fixed_window.lua
Atomic fixed window counter rate limiting.
Simple but has boundary burst problem - use sliding_window for production.
Returns: {allowed (0/1), remaining, reset_after, retry_after}
--]]
local key = KEYS[1]
local window_seconds = tonumber(ARGV[1])
local limit = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local current_window = math.floor(now / window_seconds)
local window_key = key .. ":" .. current_window
local count = tonumber(redis.call('GET', window_key)) or 0
local reset_after = window_seconds - (now % window_seconds)
if count >= limit then
return {0, 0, reset_after, reset_after}
end
local new_count = redis.call('INCR', window_key)
if new_count == 1 then
redis.call('EXPIRE', window_key, window_seconds)
end
local remaining = math.max(0, limit - new_count)
return {1, remaining, reset_after, 0}

View File

@ -0,0 +1,37 @@
--[[
AngelaMos | 2025
sliding_window.lua
Atomic sliding window counter rate limiting.
Returns: {allowed (0/1), remaining, reset_after}
--]]
local key = KEYS[1]
local window_seconds = tonumber(ARGV[1])
local limit = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local current_window = math.floor(now / window_seconds)
local previous_window = current_window - 1
local elapsed_ratio = (now % window_seconds) / window_seconds
local current_key = key .. ":" .. current_window
local previous_key = key .. ":" .. previous_window
local current_count = tonumber(redis.call('GET', current_key)) or 0
local previous_count = tonumber(redis.call('GET', previous_key)) or 0
local weighted_count = math.floor(previous_count * (1 - elapsed_ratio) + current_count)
local reset_after = window_seconds - (now % window_seconds)
if weighted_count >= limit then
return {0, 0, reset_after, reset_after}
end
redis.call('INCR', current_key)
redis.call('EXPIRE', current_key, window_seconds * 2)
local new_weighted = math.floor(previous_count * (1 - elapsed_ratio) + current_count + 1)
local remaining = math.max(0, limit - new_weighted)
return {1, remaining, reset_after, 0}

View File

@ -0,0 +1,49 @@
--[[
AngelaMos | 2025
token_bucket.lua
Atomic token bucket rate limiting.
Returns: {allowed (0/1), remaining, reset_after, retry_after}
--]]
local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refill_rate = tonumber(ARGV[2])
local tokens_to_consume = tonumber(ARGV[3])
local now = tonumber(ARGV[4])
local bucket_key = key .. ":bucket"
local data = redis.call('HMGET', bucket_key, 'tokens', 'last_refill')
local tokens = tonumber(data[1])
local last_refill = tonumber(data[2])
if tokens == nil then
tokens = capacity
last_refill = now
end
local elapsed = now - last_refill
local tokens_to_add = elapsed * refill_rate
tokens = math.min(capacity, tokens + tokens_to_add)
if tokens >= tokens_to_consume then
tokens = tokens - tokens_to_consume
redis.call('HMSET', bucket_key, 'tokens', tokens, 'last_refill', now)
redis.call('EXPIRE', bucket_key, 3600)
local time_to_full = 0
if refill_rate > 0 then
time_to_full = (capacity - tokens) / refill_rate
end
return {1, math.floor(tokens), time_to_full, 0}
end
redis.call('HMSET', bucket_key, 'tokens', tokens, 'last_refill', now)
redis.call('EXPIRE', bucket_key, 3600)
local tokens_needed = tokens_to_consume - tokens
local wait_time = tokens_needed / refill_rate
return {0, 0, wait_time, wait_time}

View File

@ -0,0 +1,294 @@
"""
AngelaMos | 2025
memory.py
"""
from __future__ import annotations
import asyncio
import contextlib
import time
from collections import OrderedDict
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from fastapi_420.types import (
RateLimitResult,
StorageType,
TokenBucketState,
WindowState,
)
if TYPE_CHECKING:
from fastapi_420.config import StorageSettings
@dataclass
class WindowEntry:
"""
Storage entry for sliding window counter.
"""
count: int = 0
window_start: int = 0
expires_at: float = 0.0
@dataclass
class MemoryStorage:
"""
In memory storage backend for rate limiting.
Thread safe through asyncio locks. Suitable for single instance
deployments, development, and as a fallback when Redis is unavailable
"""
max_keys: int = 100_000
cleanup_interval: int = 60
_windows: OrderedDict[str,
WindowEntry] = field(
default_factory = OrderedDict
)
_buckets: dict[str, TokenBucketState] = field(default_factory = dict)
_lock: asyncio.Lock = field(default_factory = asyncio.Lock)
_cleanup_task: asyncio.Task[None] | None = field(
default = None,
repr = False
)
_closed: bool = field(default = False)
@classmethod
def from_settings(cls, settings: StorageSettings) -> MemoryStorage:
"""
Create storage instance from settings
"""
return cls(
max_keys = settings.MEMORY_MAX_KEYS,
cleanup_interval = settings.MEMORY_CLEANUP_INTERVAL,
)
async def start_cleanup_task(self) -> None:
"""
Start background cleanup task for expired entries
"""
if self._cleanup_task is None:
self._cleanup_task = asyncio.create_task(self._cleanup_loop())
async def _cleanup_loop(self) -> None:
"""
Background loop to clean up expired entries
"""
while not self._closed:
await asyncio.sleep(self.cleanup_interval)
await self._cleanup_expired()
async def _cleanup_expired(self) -> None:
"""
Remove expired entries from storage
"""
now = time.time()
async with self._lock:
expired_keys = [
key for key, entry in self._windows.items()
if entry.expires_at < now
]
for key in expired_keys:
del self._windows[key]
expired_buckets = [
key for key, state in self._buckets.items()
if state.last_refill + 3600 < now
]
for key in expired_buckets:
del self._buckets[key]
async def _enforce_max_keys(self) -> None:
"""
Evict oldest entries when max_keys is exceeded
"""
while len(self._windows) > self.max_keys:
self._windows.popitem(last = False)
async def get_window_state(
self,
key: str,
window_seconds: int,
) -> WindowState:
"""
Get current sliding window state
"""
now = time.time()
current_window = int(now // window_seconds)
previous_window = current_window - 1
current_key = f"{key}:{current_window}"
previous_key = f"{key}:{previous_window}"
async with self._lock:
current_entry = self._windows.get(current_key)
previous_entry = self._windows.get(previous_key)
return WindowState(
current_count = current_entry.count
if current_entry else 0,
previous_count = previous_entry.count
if previous_entry else 0,
current_window = current_window,
window_seconds = window_seconds,
)
async def increment(
self,
key: str,
window_seconds: int,
limit: int,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Atomically check and increment counter using sliding window algorithm
"""
now = timestamp if timestamp is not None else time.time()
current_window = int(now // window_seconds)
previous_window = current_window - 1
elapsed_ratio = (now % window_seconds) / window_seconds
current_key = f"{key}:{current_window}"
previous_key = f"{key}:{previous_window}"
async with self._lock:
current_entry = self._windows.get(current_key)
previous_entry = self._windows.get(previous_key)
current_count = current_entry.count if current_entry else 0
previous_count = previous_entry.count if previous_entry else 0
weighted_count = int(
previous_count * (1 - elapsed_ratio) + current_count
)
if weighted_count >= limit:
reset_after = window_seconds - (now % window_seconds)
return RateLimitResult(
allowed = False,
limit = limit,
remaining = 0,
reset_after = reset_after,
retry_after = reset_after,
)
if current_entry:
current_entry.count += 1
else:
self._windows[current_key] = WindowEntry(
count = 1,
window_start = current_window,
expires_at = now + (window_seconds * 2),
)
self._windows.move_to_end(current_key)
await self._enforce_max_keys()
new_weighted = int(
previous_count * (1 - elapsed_ratio) + current_count + 1
)
remaining = max(0, limit - new_weighted)
reset_after = window_seconds - (now % window_seconds)
return RateLimitResult(
allowed = True,
limit = limit,
remaining = remaining,
reset_after = reset_after,
)
async def get_token_bucket_state(
self,
key: str,
) -> TokenBucketState | None:
"""
Get token bucket state if it exists
"""
async with self._lock:
return self._buckets.get(key)
async def consume_token(
self,
key: str,
capacity: int,
refill_rate: float,
tokens_to_consume: int = 1,
) -> RateLimitResult:
"""
Attempt to consume tokens from bucket
"""
now = time.time()
async with self._lock:
state = self._buckets.get(key)
if state is None:
state = TokenBucketState(
tokens = float(capacity),
last_refill = now,
capacity = capacity,
refill_rate = refill_rate,
)
self._buckets[key] = state
elapsed = now - state.last_refill
tokens_to_add = elapsed * refill_rate
state.tokens = min(
float(capacity),
state.tokens + tokens_to_add
)
state.last_refill = now
if state.tokens >= tokens_to_consume:
state.tokens -= tokens_to_consume
time_to_full = (
capacity - state.tokens
) / refill_rate if refill_rate > 0 else 0
return RateLimitResult(
allowed = True,
limit = capacity,
remaining = int(state.tokens),
reset_after = time_to_full,
)
tokens_needed = tokens_to_consume - state.tokens
wait_time = tokens_needed / refill_rate if refill_rate > 0 else float(
"inf"
)
return RateLimitResult(
allowed = False,
limit = capacity,
remaining = 0,
reset_after = wait_time,
retry_after = wait_time,
)
async def close(self) -> None:
"""
Close storage and cleanup resources
"""
self._closed = True
if self._cleanup_task:
self._cleanup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await self._cleanup_task
self._cleanup_task = None
async with self._lock:
self._windows.clear()
self._buckets.clear()
async def health_check(self) -> bool:
"""
Check if storage is healthy
"""
return not self._closed
@property
def storage_type(self) -> StorageType:
"""
Return storage type identifier
"""
return StorageType.MEMORY

View File

@ -0,0 +1,408 @@
"""
AngelaMos | 2025
redis_backend.py
"""
from __future__ import annotations
import time
from dataclasses import (
dataclass,
field,
)
from pathlib import Path
from typing import TYPE_CHECKING, Any
import redis.asyncio as redis
from redis.asyncio.connection import ConnectionPool
from redis.exceptions import (
ConnectionError as RedisConnectionError,
)
from redis.exceptions import (
ResponseError,
TimeoutError,
)
from fastapi_420.exceptions import (
StorageConnectionError,
StorageError,
)
from fastapi_420.types import (
RateLimitResult,
StorageType,
TokenBucketState,
WindowState,
)
if TYPE_CHECKING:
from fastapi_420.config import StorageSettings
LUA_SCRIPTS_DIR = Path(__file__).parent / "lua"
@dataclass
class RedisStorage:
"""
Redis storage backend with atomic Lua script operations
Uses EVALSHA with cached script hashes for optimal performance
All rate limiting operations are atomic to prevent race conditions
"""
url: str
max_connections: int = 100
socket_timeout: float = 5.0
retry_on_timeout: bool = True
decode_responses: bool = True
_client: redis.Redis[Any] | None = field(default = None, repr = False)
_pool: ConnectionPool[Any] | None = field(default = None, repr = False)
_script_shas: dict[str,
str] = field(
default_factory = dict,
repr = False
)
_scripts_loaded: bool = field(default = False, repr = False)
@classmethod
def from_settings(cls, settings: StorageSettings) -> RedisStorage:
"""
Create storage instance from settings
"""
if settings.REDIS_URL is None:
raise StorageConnectionError(
backend = StorageType.REDIS,
original_error = ValueError("REDIS_URL is required"),
)
return cls(
url = str(settings.REDIS_URL),
max_connections = settings.REDIS_MAX_CONNECTIONS,
socket_timeout = settings.REDIS_SOCKET_TIMEOUT,
retry_on_timeout = settings.REDIS_RETRY_ON_TIMEOUT,
decode_responses = settings.REDIS_DECODE_RESPONSES,
)
async def connect(self) -> None:
"""
Establish Redis connection and load Lua scripts.
"""
try:
self._pool = ConnectionPool.from_url(
self.url,
max_connections = self.max_connections,
socket_timeout = self.socket_timeout,
retry_on_timeout = self.retry_on_timeout,
decode_responses = self.decode_responses,
)
self._client = redis.Redis(connection_pool = self._pool)
await self._client.ping()
await self._load_scripts()
except (RedisConnectionError, TimeoutError, OSError) as e:
raise StorageConnectionError(
backend = StorageType.REDIS,
original_error = e,
) from e
async def _load_scripts(self) -> None:
"""
Load Lua scripts into Redis and cache their SHA1 hashes
"""
if self._client is None:
raise StorageError(
operation = "load_scripts",
backend = StorageType.REDIS,
original_error = RuntimeError("Client not connected"),
)
script_files = {
"sliding_window": LUA_SCRIPTS_DIR / "sliding_window.lua",
"token_bucket": LUA_SCRIPTS_DIR / "token_bucket.lua",
"fixed_window": LUA_SCRIPTS_DIR / "fixed_window.lua",
}
for name, path in script_files.items():
script_content = path.read_text()
sha = await self._client.script_load(script_content) # type: ignore[no-untyped-call]
self._script_shas[name] = sha
self._scripts_loaded = True
async def _ensure_connected(self) -> redis.Redis[Any]:
"""
Ensure client is connected and scripts are loaded
"""
if self._client is None:
await self.connect()
if self._client is None:
raise StorageError(
operation = "ensure_connected",
backend = StorageType.REDIS,
original_error = RuntimeError(
"Failed to establish connection"
),
)
return self._client
async def _execute_script(
self,
script_name: str,
keys: list[str],
args: list[str | int | float],
) -> list[int | float]:
"""
Execute a Lua script using EVALSHA
"""
client = await self._ensure_connected()
if script_name not in self._script_shas:
raise StorageError(
operation = "execute_script",
backend = StorageType.REDIS,
original_error = ValueError(
f"Unknown script: {script_name}"
),
)
sha = self._script_shas[script_name]
try:
result = await client.evalsha(sha, len(keys), *keys, *args) # type: ignore[no-untyped-call]
return result # type: ignore[no-any-return]
except ResponseError as e:
if "NOSCRIPT" in str(e):
await self._load_scripts()
sha = self._script_shas[script_name]
result = await client.evalsha(sha, len(keys), *keys, *args) # type: ignore[no-untyped-call]
return result # type: ignore[no-any-return]
raise StorageError(
operation = "execute_script",
backend = StorageType.REDIS,
original_error = e,
) from e
except (RedisConnectionError, TimeoutError) as e:
raise StorageError(
operation = "execute_script",
backend = StorageType.REDIS,
original_error = e,
) from e
async def get_window_state(
self,
key: str,
window_seconds: int,
) -> WindowState:
"""
Get current sliding window state
"""
client = await self._ensure_connected()
now = time.time()
current_window = int(now // window_seconds)
previous_window = current_window - 1
current_key = f"{key}:{current_window}"
previous_key = f"{key}:{previous_window}"
try:
pipeline = client.pipeline()
pipeline.get(current_key)
pipeline.get(previous_key)
results = await pipeline.execute()
current_count = int(results[0]) if results[0] else 0
previous_count = int(results[1]) if results[1] else 0
return WindowState(
current_count = current_count,
previous_count = previous_count,
current_window = current_window,
window_seconds = window_seconds,
)
except (RedisConnectionError, TimeoutError) as e:
raise StorageError(
operation = "get_window_state",
backend = StorageType.REDIS,
original_error = e,
) from e
async def increment(
self,
key: str,
window_seconds: int,
limit: int,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Atomically check and increment counter using sliding window algorithm
"""
now = timestamp if timestamp is not None else time.time()
result = await self._execute_script(
"sliding_window",
keys = [key],
args = [window_seconds,
limit,
now],
)
allowed = bool(result[0])
remaining = int(result[1])
reset_after = float(result[2])
retry_after = float(result[3]) if result[3] else None
return RateLimitResult(
allowed = allowed,
limit = limit,
remaining = remaining,
reset_after = reset_after,
retry_after = retry_after if not allowed else None,
)
async def increment_fixed_window(
self,
key: str,
window_seconds: int,
limit: int,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Atomically check and increment counter using fixed window algorithm
"""
now = timestamp if timestamp is not None else time.time()
result = await self._execute_script(
"fixed_window",
keys = [key],
args = [window_seconds,
limit,
now],
)
allowed = bool(result[0])
remaining = int(result[1])
reset_after = float(result[2])
retry_after = float(result[3]) if result[3] else None
return RateLimitResult(
allowed = allowed,
limit = limit,
remaining = remaining,
reset_after = reset_after,
retry_after = retry_after if not allowed else None,
)
async def get_token_bucket_state(
self,
key: str,
) -> TokenBucketState | None:
"""
Get token bucket state if it exists
"""
client = await self._ensure_connected()
bucket_key = f"{key}:bucket"
try:
data = await client.hmget(
bucket_key,
"tokens",
"last_refill",
"capacity",
"refill_rate"
)
if data[0] is None:
return None
return TokenBucketState(
tokens = float(data[0]),
last_refill = float(data[1]), # type: ignore[arg-type]
capacity = int(data[2]), # type: ignore[arg-type]
refill_rate = float(data[3]), # type: ignore[arg-type]
)
except (RedisConnectionError, TimeoutError) as e:
raise StorageError(
operation = "get_token_bucket_state",
backend = StorageType.REDIS,
original_error = e,
) from e
async def consume_token(
self,
key: str,
capacity: int,
refill_rate: float,
tokens_to_consume: int = 1,
) -> RateLimitResult:
"""
Attempt to consume tokens from bucket atomically
"""
now = time.time()
result = await self._execute_script(
"token_bucket",
keys = [key],
args = [capacity,
refill_rate,
tokens_to_consume,
now],
)
allowed = bool(result[0])
remaining = int(result[1])
reset_after = float(result[2])
retry_after = float(result[3]) if result[3] else None
return RateLimitResult(
allowed = allowed,
limit = capacity,
remaining = remaining,
reset_after = reset_after,
retry_after = retry_after if not allowed else None,
)
async def close(self) -> None:
"""
Close Redis connection.
"""
if self._client:
await self._client.close()
self._client = None
if self._pool:
await self._pool.disconnect()
self._pool = None
self._scripts_loaded = False
self._script_shas.clear()
async def health_check(self) -> bool:
"""
Check if Redis connection is healthy
"""
try:
client = await self._ensure_connected()
await client.ping()
return True
except (StorageError, RedisConnectionError, TimeoutError):
return False
@property
def storage_type(self) -> StorageType:
"""
Return storage type identifier
"""
return StorageType.REDIS
@property
def is_connected(self) -> bool:
"""
Check if client is connected
"""
return self._client is not None and self._scripts_loaded

View File

@ -0,0 +1,401 @@
"""
AngelaMos | 2025
types.py
"""
# pylint: disable=unnecessary-ellipsis
from __future__ import annotations
from enum import StrEnum
from dataclasses import dataclass, field
from typing import (
TYPE_CHECKING,
Protocol,
runtime_checkable,
)
if TYPE_CHECKING:
from starlette.requests import Request
class Algorithm(StrEnum):
"""
Rate limiting algorithm selection
"""
SLIDING_WINDOW = "sliding_window"
TOKEN_BUCKET = "token_bucket"
FIXED_WINDOW = "fixed_window"
LEAKY_BUCKET = "leaky_bucket"
class FingerprintLevel(StrEnum):
"""
Preset fingerprinting intensity levels
"""
STRICT = "strict"
NORMAL = "normal"
RELAXED = "relaxed"
CUSTOM = "custom"
class DefenseMode(StrEnum):
"""
Global circuit breaker defense strategies
"""
ADAPTIVE = "adaptive"
LOCKDOWN = "lockdown"
CHALLENGE = "challenge"
DISABLED = "disabled"
class StorageType(StrEnum):
"""
Backend storage type selection
"""
REDIS = "redis"
MEMORY = "memory"
class Layer(StrEnum):
"""
Rate limiting defense layers
"""
USER = "user"
ENDPOINT = "endpoint"
GLOBAL = "global"
@dataclass(frozen = True, slots = True)
class RateLimitResult:
"""
Result of a rate limit check operation
"""
allowed: bool
limit: int
remaining: int
reset_after: float
retry_after: float | None = None
@property
def headers(self) -> dict[str, str]:
"""
Generate IETF draft-compliant rate limit headers
"""
hdrs = {
"RateLimit-Limit": str(self.limit),
"RateLimit-Remaining": str(max(0,
self.remaining)),
"RateLimit-Reset": str(int(self.reset_after)),
}
if self.retry_after is not None:
hdrs["Retry-After"] = str(int(self.retry_after))
return hdrs
@dataclass(frozen = True, slots = True)
class RateLimitRule:
"""
A single rate limit rule defining requests allowed per time window
"""
requests: int
window_seconds: int
def __post_init__(self) -> None:
if self.requests <= 0:
raise ValueError("requests must be positive")
if self.window_seconds <= 0:
raise ValueError("window_seconds must be positive")
@classmethod
def parse(cls, rule_string: str) -> RateLimitRule:
"""
Parse rate limit string like '100/minute' or '1000/hour'
"""
parts = rule_string.strip().lower().split("/")
if len(parts) != 2:
raise ValueError(f"Invalid rate limit format: {rule_string}")
try:
requests = int(parts[0])
except ValueError as e:
raise ValueError(f"Invalid request count: {parts[0]}") from e
window_map = {
"second": 1,
"seconds": 1,
"sec": 1,
"s": 1,
"minute": 60,
"minutes": 60,
"min": 60,
"m": 60,
"hour": 3600,
"hours": 3600,
"hr": 3600,
"h": 3600,
"day": 86400,
"days": 86400,
"d": 86400,
}
window_str = parts[1].strip()
if window_str not in window_map:
raise ValueError(f"Unknown time unit: {window_str}")
return cls(
requests = requests,
window_seconds = window_map[window_str]
)
def __str__(self) -> str:
if self.window_seconds == 1:
unit = "second"
elif self.window_seconds == 60:
unit = "minute"
elif self.window_seconds == 3600:
unit = "hour"
elif self.window_seconds == 86400:
unit = "day"
else:
unit = f"{self.window_seconds}s"
return f"{self.requests}/{unit}"
@dataclass(slots = True)
class FingerprintData:
"""
Collected fingerprint data from a request
"""
ip: str
ip_normalized: str
user_agent: str | None = None
accept_language: str | None = None
accept_encoding: str | None = None
headers_hash: str | None = None
auth_identifier: str | None = None
tls_fingerprint: str | None = None
geo_asn: str | None = None
def to_composite_key(self, level: FingerprintLevel) -> str:
"""
Generate composite fingerprint key based on level.
"""
if level == FingerprintLevel.RELAXED:
components = [self.ip_normalized]
if self.auth_identifier:
components.append(self.auth_identifier)
return ":".join(filter(None, components))
if level == FingerprintLevel.NORMAL:
components = [
self.ip_normalized,
self.user_agent or "",
self.auth_identifier or "",
]
return ":".join(components)
components = [
self.ip_normalized,
self.user_agent or "",
self.accept_language or "",
self.accept_encoding or "",
self.headers_hash or "",
self.auth_identifier or "",
self.tls_fingerprint or "",
self.geo_asn or "",
]
return ":".join(components)
@dataclass(slots = True)
class WindowState:
"""
State for sliding window counter algorithm
"""
current_count: int = 0
previous_count: int = 0
current_window: int = 0
window_seconds: int = 60
def weighted_count(self, elapsed_ratio: float) -> float:
"""
Calculate weighted count using sliding window interpolation.
"""
return self.previous_count * (
1 - elapsed_ratio
) + self.current_count
@dataclass(slots = True)
class TokenBucketState:
"""
State for token bucket algorithm
"""
tokens: float
last_refill: float
capacity: int
refill_rate: float
@dataclass(slots = True)
class CircuitState:
"""
Global circuit breaker state
"""
is_open: bool = False
failure_count: int = 0
last_failure_time: float = 0.0
half_open_requests: int = 0
total_requests_in_window: int = 0
@dataclass(slots = True)
class DefenseContext:
"""
Context passed to defense strategies
"""
fingerprint: FingerprintData
endpoint: str
method: str
is_authenticated: bool = False
reputation_score: float = 1.0
request_count_last_minute: int = 0
@dataclass(slots = True)
class RateLimitKey:
"""
Structured rate limit key components
"""
prefix: str = "ratelimit"
version: str = "v1"
layer: Layer = Layer.USER
endpoint: str = ""
identifier: str = ""
window: int = 0
def build(self) -> str:
"""
Build the full Redis key string.
"""
return f"{self.prefix}:{self.version}:{self.layer.value}:{self.endpoint}:{self.identifier}:{self.window}"
@runtime_checkable
class StorageBackend(Protocol):
"""
Protocol for rate limit storage backends
"""
async def get_window_state(
self,
key: str,
window_seconds: int,
) -> WindowState:
"""
Get current sliding window state
"""
...
async def increment(
self,
key: str,
window_seconds: int,
limit: int,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Atomically check and increment counter, returning result
"""
...
async def get_token_bucket_state(
self,
key: str,
) -> TokenBucketState | None:
"""
Get token bucket state if it exists
"""
...
async def consume_token(
self,
key: str,
capacity: int,
refill_rate: float,
tokens_to_consume: int = 1,
) -> RateLimitResult:
"""
Attempt to consume tokens from bucket
"""
...
async def close(self) -> None:
"""
Close storage connections.
"""
...
async def health_check(self) -> bool:
"""
Check if storage backend is healthy
"""
...
@runtime_checkable
class Fingerprinter(Protocol):
"""
Protocol for request fingerprinting strategies.
"""
async def extract(self, request: Request) -> FingerprintData:
"""
Extract fingerprint data from request.
"""
...
@runtime_checkable
class RateLimitAlgorithm(Protocol):
"""
Protocol for rate limiting algorithms
"""
async def check(
self,
storage: StorageBackend,
key: str,
rule: RateLimitRule,
timestamp: float | None = None,
) -> RateLimitResult:
"""
Check if request is allowed under rate limit
"""
...
@dataclass
class LimiterConfig:
"""
Main configuration for the rate limiter.
"""
default_rules: list[RateLimitRule] = field(default_factory = list)
algorithm: Algorithm = Algorithm.SLIDING_WINDOW
fingerprint_level: FingerprintLevel = FingerprintLevel.NORMAL
defense_mode: DefenseMode = DefenseMode.ADAPTIVE
fail_open: bool = True
key_prefix: str = "ratelimit"
include_headers: bool = True
log_violations: bool = True
global_limit: RateLimitRule | None = None
endpoint_limits: dict[str,
list[RateLimitRule]] = field(
default_factory = dict
)
def __post_init__(self) -> None:
if not self.default_rules:
self.default_rules = [
RateLimitRule.parse("100/minute"),
RateLimitRule.parse("1000/hour"),
]

View File

@ -0,0 +1,4 @@
"""
AngelaMos | 2025
__init__.py
"""

View File

@ -0,0 +1,984 @@
"""
AngelaMos | 2025
conftest.py
"""
from __future__ import annotations
import asyncio
import hashlib
import time
from collections.abc import AsyncGenerator, Generator
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import MagicMock
import pytest
from fastapi import FastAPI, Request
from httpx import ASGITransport, AsyncClient
from starlette.testclient import TestClient
from fastapi_420.algorithms import create_algorithm
from fastapi_420.algorithms.base import BaseAlgorithm
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
from fastapi_420.config import (
DefenseSettings,
FingerprintSettings,
RateLimiterSettings,
StorageSettings,
)
from fastapi_420.defense.circuit_breaker import CircuitBreaker
from fastapi_420.defense.layers import LayeredDefense
from fastapi_420.dependencies import set_global_limiter
from fastapi_420.exceptions import HTTP_420_ENHANCE_YOUR_CALM
from fastapi_420.fingerprinting.auth import AuthExtractor
from fastapi_420.fingerprinting.composite import CompositeFingerprinter
from fastapi_420.fingerprinting.headers import HeadersExtractor
from fastapi_420.fingerprinting.ip import IPExtractor
from fastapi_420.limiter import RateLimiter
from fastapi_420.middleware import RateLimitMiddleware
from fastapi_420.storage import MemoryStorage
from fastapi_420.types import (
Algorithm,
CircuitState,
DefenseContext,
DefenseMode,
FingerprintData,
FingerprintLevel,
Layer,
RateLimitKey,
RateLimitResult,
RateLimitRule,
StorageType,
TokenBucketState,
WindowState,
)
WINDOW_SECOND = 1
WINDOW_MINUTE = 60
WINDOW_HOUR = 3600
WINDOW_DAY = 86400
DEFAULT_LIMIT_REQUESTS = 100
DEFAULT_LIMIT_WINDOW = WINDOW_MINUTE
STRICT_LIMIT_REQUESTS = 10
STRICT_LIMIT_WINDOW = WINDOW_MINUTE
TEST_IP_V4 = "192.168.1.100"
TEST_IP_V4_PRIVATE = "10.0.0.1"
TEST_IP_V6 = "2001:0db8:85a3:0000:0000:8a2e:0370:7334"
TEST_IP_V6_NORMALIZED = "2001:db8:85a3::"
TEST_IP_LOCALHOST = "127.0.0.1"
TEST_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
TEST_ACCEPT_LANGUAGE = "en-US,en;q=0.9"
TEST_ACCEPT_ENCODING = "gzip, deflate, br"
TEST_JWT_TOKEN = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ1c2VyXzEyMyIsImV4cCI6OTk5OTk5OTk5OX0.signature"
TEST_JWT_SUBJECT = "user_123"
TEST_API_KEY = "sk-test-api-key-12345"
TEST_SESSION_ID = "sess_abc123def456"
TEST_ENDPOINT = "/api/v1/test"
TEST_METHOD = "GET"
CIRCUIT_THRESHOLD = 1000
CIRCUIT_WINDOW = WINDOW_MINUTE
CIRCUIT_RECOVERY = 30
KEY_PREFIX = "ratelimit"
KEY_VERSION = "v1"
@dataclass
class MockScope:
"""
Mock ASGI scope for request creation
"""
type: str = "http"
method: str = "GET"
path: str = "/"
query_string: bytes = b""
headers: list[tuple[bytes, bytes]] = field(default_factory=list)
client: tuple[str, int] | None = None
route: Any = None
class MockRoute:
"""
Mock route object for endpoint extraction
"""
def __init__(self, path: str = TEST_ENDPOINT) -> None:
self.path = path
class RequestFactory:
"""
Factory for creating mock Starlette Request objects
"""
@staticmethod
def create(
method: str = TEST_METHOD,
path: str = TEST_ENDPOINT,
client_ip: str = TEST_IP_V4,
client_port: int = 12345,
headers: dict[str, str] | None = None,
query_params: dict[str, str] | None = None,
cookies: dict[str, str] | None = None,
include_route: bool = True,
) -> Request:
"""
Create a mock Request with configurable parameters
"""
header_list: list[tuple[bytes, bytes]] = []
default_headers = {
"host": "localhost",
"user-agent": TEST_USER_AGENT,
"accept": "*/*",
"accept-language": TEST_ACCEPT_LANGUAGE,
"accept-encoding": TEST_ACCEPT_ENCODING,
}
if headers:
default_headers.update(headers)
for key, value in default_headers.items():
header_list.append((key.lower().encode(), value.encode()))
query_string = b""
if query_params:
query_string = "&".join(
f"{k}={v}" for k, v in query_params.items()
).encode()
scope = {
"type": "http",
"method": method,
"path": path,
"query_string": query_string,
"headers": header_list,
"client": (client_ip, client_port),
}
if include_route:
scope["route"] = MockRoute(path)
request = Request(scope)
if cookies:
request._cookies = cookies
return request
@staticmethod
def with_auth(
auth_type: str = "bearer",
token: str = TEST_JWT_TOKEN,
**kwargs: Any,
) -> Request:
"""
Create request with authentication header
"""
headers = kwargs.pop("headers", {}) or {}
if auth_type == "bearer":
headers["authorization"] = f"Bearer {token}"
elif auth_type == "api_key":
headers["x-api-key"] = token
elif auth_type == "basic":
headers["authorization"] = f"Basic {token}"
return RequestFactory.create(headers=headers, **kwargs)
@staticmethod
def with_forwarded_for(
forwarded_ips: list[str],
real_ip: str | None = None,
**kwargs: Any,
) -> Request:
"""
Create request with proxy headers
"""
headers = kwargs.pop("headers", {}) or {}
headers["x-forwarded-for"] = ", ".join(forwarded_ips)
if real_ip:
headers["x-real-ip"] = real_ip
return RequestFactory.create(headers=headers, **kwargs)
class FingerprintFactory:
"""
Factory for creating FingerprintData instances
"""
@staticmethod
def create(
ip: str = TEST_IP_V4,
ip_normalized: str | None = None,
user_agent: str | None = TEST_USER_AGENT,
accept_language: str | None = TEST_ACCEPT_LANGUAGE,
accept_encoding: str | None = TEST_ACCEPT_ENCODING,
headers_hash: str | None = None,
auth_identifier: str | None = None,
tls_fingerprint: str | None = None,
geo_asn: str | None = None,
) -> FingerprintData:
"""
Create FingerprintData with sensible defaults
"""
return FingerprintData(
ip=ip,
ip_normalized=ip_normalized or ip,
user_agent=user_agent,
accept_language=accept_language,
accept_encoding=accept_encoding,
headers_hash=headers_hash,
auth_identifier=auth_identifier,
tls_fingerprint=tls_fingerprint,
geo_asn=geo_asn,
)
@staticmethod
def authenticated(
auth_id: str = "user_123",
hash_id: bool = True,
**kwargs: Any,
) -> FingerprintData:
"""
Create authenticated fingerprint
"""
identifier = auth_id
if hash_id:
identifier = hashlib.sha256(auth_id.encode()).hexdigest()[:16]
return FingerprintFactory.create(auth_identifier=identifier, **kwargs)
@staticmethod
def anonymous(**kwargs: Any) -> FingerprintData:
"""
Create anonymous fingerprint (no auth)
"""
return FingerprintFactory.create(auth_identifier=None, **kwargs)
@staticmethod
def minimal(ip: str = TEST_IP_V4) -> FingerprintData:
"""
Create minimal fingerprint (IP only)
"""
return FingerprintData(
ip=ip,
ip_normalized=ip,
)
class RuleFactory:
"""
Factory for creating RateLimitRule instances
"""
@staticmethod
def create(
requests: int = DEFAULT_LIMIT_REQUESTS,
window_seconds: int = DEFAULT_LIMIT_WINDOW,
) -> RateLimitRule:
"""
Create RateLimitRule with defaults
"""
return RateLimitRule(requests=requests, window_seconds=window_seconds)
@staticmethod
def per_second(requests: int = 10) -> RateLimitRule:
return RateLimitRule(requests=requests, window_seconds=WINDOW_SECOND)
@staticmethod
def per_minute(requests: int = 100) -> RateLimitRule:
return RateLimitRule(requests=requests, window_seconds=WINDOW_MINUTE)
@staticmethod
def per_hour(requests: int = 1000) -> RateLimitRule:
return RateLimitRule(requests=requests, window_seconds=WINDOW_HOUR)
@staticmethod
def per_day(requests: int = 10000) -> RateLimitRule:
return RateLimitRule(requests=requests, window_seconds=WINDOW_DAY)
@staticmethod
def strict() -> RateLimitRule:
return RateLimitRule(
requests=STRICT_LIMIT_REQUESTS,
window_seconds=STRICT_LIMIT_WINDOW,
)
@staticmethod
def parse(rule_string: str) -> RateLimitRule:
return RateLimitRule.parse(rule_string)
class ResultFactory:
"""
Factory for creating RateLimitResult instances
"""
@staticmethod
def allowed(
limit: int = DEFAULT_LIMIT_REQUESTS,
remaining: int | None = None,
reset_after: float = 60.0,
) -> RateLimitResult:
"""
Create allowed result
"""
return RateLimitResult(
allowed=True,
limit=limit,
remaining=remaining if remaining is not None else limit - 1,
reset_after=reset_after,
)
@staticmethod
def denied(
limit: int = DEFAULT_LIMIT_REQUESTS,
reset_after: float = 60.0,
retry_after: float | None = None,
) -> RateLimitResult:
"""
Create denied result
"""
return RateLimitResult(
allowed=False,
limit=limit,
remaining=0,
reset_after=reset_after,
retry_after=retry_after or reset_after,
)
@staticmethod
def near_limit(
limit: int = DEFAULT_LIMIT_REQUESTS,
remaining: int = 1,
reset_after: float = 30.0,
) -> RateLimitResult:
"""
Create result near the limit
"""
return RateLimitResult(
allowed=True,
limit=limit,
remaining=remaining,
reset_after=reset_after,
)
class KeyFactory:
"""
Factory for creating RateLimitKey instances
"""
@staticmethod
def create(
prefix: str = KEY_PREFIX,
version: str = KEY_VERSION,
layer: Layer = Layer.USER,
endpoint: str = TEST_ENDPOINT,
identifier: str = TEST_IP_V4,
window: int = WINDOW_MINUTE,
) -> RateLimitKey:
"""
Create RateLimitKey with defaults
"""
return RateLimitKey(
prefix=prefix,
version=version,
layer=layer,
endpoint=endpoint,
identifier=identifier,
window=window,
)
@staticmethod
def user_key(
endpoint: str = TEST_ENDPOINT,
identifier: str = TEST_IP_V4,
window: int = WINDOW_MINUTE,
) -> RateLimitKey:
return KeyFactory.create(
layer=Layer.USER,
endpoint=endpoint,
identifier=identifier,
window=window,
)
@staticmethod
def endpoint_key(
endpoint: str = TEST_ENDPOINT,
window: int = WINDOW_MINUTE,
) -> RateLimitKey:
return KeyFactory.create(
layer=Layer.ENDPOINT,
endpoint=endpoint,
identifier="global",
window=window,
)
@staticmethod
def global_key(window: int = WINDOW_MINUTE) -> RateLimitKey:
return KeyFactory.create(
layer=Layer.GLOBAL,
endpoint="",
identifier="global",
window=window,
)
class WindowStateFactory:
"""
Factory for creating WindowState instances
"""
@staticmethod
def create(
current_count: int = 0,
previous_count: int = 0,
current_window: int | None = None,
window_seconds: int = WINDOW_MINUTE,
) -> WindowState:
"""
Create WindowState with defaults
"""
if current_window is None:
current_window = int(time.time() // window_seconds)
return WindowState(
current_count=current_count,
previous_count=previous_count,
current_window=current_window,
window_seconds=window_seconds,
)
@staticmethod
def empty() -> WindowState:
return WindowStateFactory.create()
@staticmethod
def with_usage(current: int, previous: int = 0) -> WindowState:
return WindowStateFactory.create(
current_count=current,
previous_count=previous,
)
class TokenBucketStateFactory:
"""
Factory for creating TokenBucketState instances
"""
@staticmethod
def create(
tokens: float = 100.0,
last_refill: float | None = None,
capacity: int = 100,
refill_rate: float = 1.67,
) -> TokenBucketState:
"""
Create TokenBucketState with defaults
"""
return TokenBucketState(
tokens=tokens,
last_refill=last_refill or time.time(),
capacity=capacity,
refill_rate=refill_rate,
)
@staticmethod
def full(capacity: int = 100) -> TokenBucketState:
return TokenBucketStateFactory.create(
tokens=float(capacity),
capacity=capacity,
)
@staticmethod
def empty(capacity: int = 100) -> TokenBucketState:
return TokenBucketStateFactory.create(tokens=0.0, capacity=capacity)
class DefenseContextFactory:
"""
Factory for creating DefenseContext instances
"""
@staticmethod
def create(
fingerprint: FingerprintData | None = None,
endpoint: str = TEST_ENDPOINT,
method: str = TEST_METHOD,
is_authenticated: bool = False,
reputation_score: float = 1.0,
request_count_last_minute: int = 0,
) -> DefenseContext:
"""
Create DefenseContext with defaults
"""
return DefenseContext(
fingerprint=fingerprint or FingerprintFactory.create(),
endpoint=endpoint,
method=method,
is_authenticated=is_authenticated,
reputation_score=reputation_score,
request_count_last_minute=request_count_last_minute,
)
@staticmethod
def authenticated(**kwargs: Any) -> DefenseContext:
fp = FingerprintFactory.authenticated()
return DefenseContextFactory.create(
fingerprint=fp,
is_authenticated=True,
**kwargs,
)
@staticmethod
def suspicious(
reputation_score: float = 0.3,
request_count: int = 500,
) -> DefenseContext:
return DefenseContextFactory.create(
reputation_score=reputation_score,
request_count_last_minute=request_count,
)
class CircuitStateFactory:
"""
Factory for creating CircuitState instances
"""
@staticmethod
def create(
is_open: bool = False,
failure_count: int = 0,
last_failure_time: float = 0.0,
half_open_requests: int = 0,
total_requests_in_window: int = 0,
) -> CircuitState:
"""
Create CircuitState with defaults
"""
return CircuitState(
is_open=is_open,
failure_count=failure_count,
last_failure_time=last_failure_time,
half_open_requests=half_open_requests,
total_requests_in_window=total_requests_in_window,
)
@staticmethod
def closed() -> CircuitState:
return CircuitStateFactory.create()
@staticmethod
def open(failure_time: float | None = None) -> CircuitState:
return CircuitStateFactory.create(
is_open=True,
failure_count=1,
last_failure_time=failure_time or time.time(),
)
@pytest.fixture
def storage_settings() -> StorageSettings:
"""
Create test storage settings (memory backend)
"""
return StorageSettings(
REDIS_URL=None,
MEMORY_MAX_KEYS=10000,
MEMORY_CLEANUP_INTERVAL=60,
FALLBACK_TO_MEMORY=True,
)
@pytest.fixture
def fingerprint_settings() -> FingerprintSettings:
"""
Create test fingerprint settings
"""
return FingerprintSettings(
LEVEL=FingerprintLevel.NORMAL,
USE_IP=True,
USE_USER_AGENT=True,
USE_ACCEPT_HEADERS=False,
USE_HEADER_ORDER=False,
USE_AUTH=True,
USE_TLS=False,
USE_GEO=False,
IPV6_PREFIX_LENGTH=64,
TRUSTED_PROXIES=[],
TRUST_X_FORWARDED_FOR=False,
)
@pytest.fixture
def defense_settings() -> DefenseSettings:
"""
Create test defense settings
"""
return DefenseSettings(
MODE=DefenseMode.ADAPTIVE,
GLOBAL_LIMIT="50000/minute",
CIRCUIT_THRESHOLD=CIRCUIT_THRESHOLD,
CIRCUIT_WINDOW=CIRCUIT_WINDOW,
CIRCUIT_RECOVERY_TIME=CIRCUIT_RECOVERY,
ADAPTIVE_REDUCTION_FACTOR=0.5,
ENDPOINT_LIMIT_MULTIPLIER=10,
LOCKDOWN_ALLOW_AUTHENTICATED=True,
LOCKDOWN_ALLOW_KNOWN_GOOD=True,
)
@pytest.fixture
def rate_limiter_settings(
storage_settings: StorageSettings,
fingerprint_settings: FingerprintSettings,
defense_settings: DefenseSettings,
) -> RateLimiterSettings:
"""
Create test rate limiter settings
"""
return RateLimiterSettings(
ENABLED=True,
ALGORITHM=Algorithm.SLIDING_WINDOW,
DEFAULT_LIMIT="100/minute",
DEFAULT_LIMITS=["100/minute", "1000/hour"],
FAIL_OPEN=True,
KEY_PREFIX=KEY_PREFIX,
KEY_VERSION=KEY_VERSION,
INCLUDE_HEADERS=True,
LOG_VIOLATIONS=False,
ENVIRONMENT="development",
HTTP_420_MESSAGE="Enhance your calm",
HTTP_420_DETAIL="Rate limit exceeded. Take a breather.",
storage=storage_settings,
fingerprint=fingerprint_settings,
defense=defense_settings,
)
@pytest.fixture
async def memory_storage() -> AsyncGenerator[MemoryStorage, None]:
"""
Create and manage MemoryStorage instance
"""
storage = MemoryStorage(max_keys=10000, cleanup_interval=60)
await storage.start_cleanup_task()
yield storage
await storage.close()
@pytest.fixture
def sliding_window_algorithm() -> SlidingWindowAlgorithm:
"""
Create sliding window algorithm instance
"""
return SlidingWindowAlgorithm()
@pytest.fixture
def token_bucket_algorithm() -> TokenBucketAlgorithm:
"""
Create token bucket algorithm instance
"""
return TokenBucketAlgorithm()
@pytest.fixture
def fixed_window_algorithm() -> FixedWindowAlgorithm:
"""
Create fixed window algorithm instance
"""
return FixedWindowAlgorithm()
@pytest.fixture
def ip_extractor() -> IPExtractor:
"""
Create IP extractor instance
"""
return IPExtractor(
ipv6_prefix_length=64,
trusted_proxies=[],
trust_x_forwarded_for=False,
)
@pytest.fixture
def headers_extractor() -> HeadersExtractor:
"""
Create headers extractor instance
"""
return HeadersExtractor(use_header_order=False, hash_length=16)
@pytest.fixture
def auth_extractor() -> AuthExtractor:
"""
Create auth extractor instance
"""
return AuthExtractor(
jwt_secret=None,
jwt_algorithms=["HS256"],
api_key_header="X-API-Key",
api_key_query_param="api_key",
session_cookie="session_id",
hash_identifiers=True,
hash_length=16,
)
@pytest.fixture
def composite_fingerprinter(
ip_extractor: IPExtractor,
headers_extractor: HeadersExtractor,
auth_extractor: AuthExtractor,
) -> CompositeFingerprinter:
"""
Create composite fingerprinter instance
"""
return CompositeFingerprinter(
level=FingerprintLevel.NORMAL,
ip_extractor=ip_extractor,
headers_extractor=headers_extractor,
auth_extractor=auth_extractor,
)
@pytest.fixture
async def circuit_breaker() -> CircuitBreaker:
"""
Create circuit breaker instance
"""
return CircuitBreaker(
threshold=CIRCUIT_THRESHOLD,
window_seconds=CIRCUIT_WINDOW,
recovery_time=CIRCUIT_RECOVERY,
defense_mode=DefenseMode.ADAPTIVE,
)
@pytest.fixture
async def rate_limiter(
rate_limiter_settings: RateLimiterSettings,
memory_storage: MemoryStorage,
) -> AsyncGenerator[RateLimiter, None]:
"""
Create and manage RateLimiter instance
"""
limiter = RateLimiter(
settings=rate_limiter_settings,
storage=memory_storage,
)
await limiter.init()
yield limiter
await limiter.close()
@pytest.fixture
async def layered_defense(
memory_storage: MemoryStorage,
rate_limiter_settings: RateLimiterSettings,
circuit_breaker: CircuitBreaker,
) -> LayeredDefense:
"""
Create layered defense instance
"""
return LayeredDefense(
storage=memory_storage,
settings=rate_limiter_settings,
circuit_breaker=circuit_breaker,
)
@pytest.fixture
def test_request() -> Request:
"""
Create a default test request
"""
return RequestFactory.create()
@pytest.fixture
def authenticated_request() -> Request:
"""
Create an authenticated test request
"""
return RequestFactory.with_auth()
@pytest.fixture
def test_fingerprint() -> FingerprintData:
"""
Create a default test fingerprint
"""
return FingerprintFactory.create()
@pytest.fixture
def test_rule() -> RateLimitRule:
"""
Create a default test rule
"""
return RuleFactory.per_minute()
@pytest.fixture
def strict_rule() -> RateLimitRule:
"""
Create a strict test rule
"""
return RuleFactory.strict()
def create_test_app(
limiter: RateLimiter | None = None,
with_middleware: bool = False,
default_limit: str = "100/minute",
) -> FastAPI:
"""
Create a test FastAPI application
"""
app = FastAPI(title="Test API")
@app.get("/")
async def root() -> dict[str, str]:
return {"message": "Hello World"}
@app.get("/health")
async def health() -> dict[str, str]:
return {"status": "healthy"}
@app.get(TEST_ENDPOINT)
async def test_endpoint() -> dict[str, str]:
return {"endpoint": "test"}
@app.post(TEST_ENDPOINT)
async def test_endpoint_post() -> dict[str, str]:
return {"created": "true"}
@app.get("/api/v1/protected")
async def protected_endpoint() -> dict[str, str]:
return {"protected": "true"}
if with_middleware and limiter:
app.add_middleware(
RateLimitMiddleware,
limiter=limiter,
default_limit=default_limit,
)
return app
@pytest.fixture
def test_app() -> FastAPI:
"""
Create a basic test app without middleware
"""
return create_test_app()
@pytest.fixture
async def test_app_with_limiter(
rate_limiter: RateLimiter,
) -> FastAPI:
"""
Create a test app with rate limiting middleware
"""
app = create_test_app(limiter=rate_limiter, with_middleware=True)
set_global_limiter(rate_limiter)
return app
@pytest.fixture
async def async_client(test_app: FastAPI) -> AsyncGenerator[AsyncClient, None]:
"""
Create async HTTP client for testing
"""
async with AsyncClient(
transport=ASGITransport(app=test_app),
base_url="http://test",
) as client:
yield client
@pytest.fixture
async def rate_limited_client(
test_app_with_limiter: FastAPI,
) -> AsyncGenerator[AsyncClient, None]:
"""
Create async HTTP client with rate limiting
"""
async with AsyncClient(
transport=ASGITransport(app=test_app_with_limiter),
base_url="http://test",
) as client:
yield client
@pytest.fixture
def sync_client(test_app: FastAPI) -> Generator[TestClient, None, None]:
"""
Create sync test client
"""
with TestClient(test_app) as client:
yield client
def assert_rate_limit_headers(
headers: dict[str, str],
expected_limit: int | None = None,
) -> None:
"""
Assert rate limit headers are present and valid
"""
assert "RateLimit-Limit" in headers
assert "RateLimit-Remaining" in headers
assert "RateLimit-Reset" in headers
if expected_limit is not None:
assert int(headers["RateLimit-Limit"]) == expected_limit
assert int(headers["RateLimit-Remaining"]) >= 0
assert int(headers["RateLimit-Reset"]) >= 0
def assert_420_response(
response: Any,
check_headers: bool = True,
) -> None:
"""
Assert response is HTTP 420 with proper structure
"""
assert response.status_code == HTTP_420_ENHANCE_YOUR_CALM
if check_headers:
assert "Retry-After" in response.headers or "RateLimit-Reset" in response.headers
async def exhaust_rate_limit(
storage: MemoryStorage,
key: str,
limit: int,
window_seconds: int = WINDOW_MINUTE,
) -> None:
"""
Helper to exhaust a rate limit by making requests
"""
for _ in range(limit):
await storage.increment(
key=key,
window_seconds=window_seconds,
limit=limit,
)
async def wait_for_window_reset(window_seconds: int = 1) -> None:
"""
Helper to wait for a window to reset
"""
await asyncio.sleep(window_seconds + 0.1)

View File

@ -0,0 +1,413 @@
"""
AngelaMos | 2025
test_algorithms.py
"""
from __future__ import annotations
import time
import pytest
from fastapi_420.algorithms import create_algorithm
from fastapi_420.algorithms.sliding_window import SlidingWindowAlgorithm
from fastapi_420.algorithms.token_bucket import TokenBucketAlgorithm
from fastapi_420.algorithms.fixed_window import FixedWindowAlgorithm
from fastapi_420.storage import MemoryStorage
from fastapi_420.types import Algorithm
from tests.conftest import (
WINDOW_MINUTE,
WINDOW_SECOND,
RuleFactory,
)
class TestAlgorithmFactory:
"""
Tests for algorithm factory function
"""
def test_create_sliding_window(self) -> None:
algo = create_algorithm(Algorithm.SLIDING_WINDOW)
assert isinstance(algo, SlidingWindowAlgorithm)
def test_create_token_bucket(self) -> None:
algo = create_algorithm(Algorithm.TOKEN_BUCKET)
assert isinstance(algo, TokenBucketAlgorithm)
def test_create_fixed_window(self) -> None:
algo = create_algorithm(Algorithm.FIXED_WINDOW)
assert isinstance(algo, FixedWindowAlgorithm)
def test_create_from_string(self) -> None:
algo = create_algorithm(Algorithm.SLIDING_WINDOW)
assert algo.name == Algorithm.SLIDING_WINDOW.value
class TestSlidingWindowAlgorithm:
"""
Tests for sliding window counter algorithm
"""
@pytest.mark.asyncio
async def test_algorithm_name(self) -> None:
algo = SlidingWindowAlgorithm()
assert algo.name == "sliding_window"
@pytest.mark.asyncio
async def test_first_request_allowed(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
result = await algo.check(
storage=storage,
key="test_key",
rule=rule,
)
assert result.allowed is True
assert result.remaining == 99
await storage.close()
@pytest.mark.asyncio
async def test_multiple_requests(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
for i in range(50):
result = await algo.check(storage, f"multi_key", rule)
assert result.allowed is True
assert result.remaining == 50
await storage.close()
@pytest.mark.asyncio
async def test_limit_exceeded(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
result = await algo.check(storage, "limit_key", rule)
assert result.allowed is True
result = await algo.check(storage, "limit_key", rule)
assert result.allowed is False
assert result.retry_after is not None
await storage.close()
@pytest.mark.asyncio
async def test_different_keys_independent(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
await algo.check(storage, "key_a", rule)
result_a = await algo.check(storage, "key_a", rule)
result_b = await algo.check(storage, "key_b", rule)
assert result_a.allowed is False
assert result_b.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_empty(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute()
usage = await algo.get_current_usage(storage, "empty_key", rule)
assert usage == 0
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_after_requests(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
for _ in range(25):
await algo.check(storage, "usage_key", rule)
usage = await algo.get_current_usage(storage, "usage_key", rule)
assert usage >= 24
await storage.close()
@pytest.mark.asyncio
async def test_explicit_timestamp(self) -> None:
algo = SlidingWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
fixed_time = 1000000.0
result = await algo.check(
storage=storage,
key="timestamp_key",
rule=rule,
timestamp=fixed_time,
)
assert result.allowed is True
await storage.close()
class TestTokenBucketAlgorithm:
"""
Tests for token bucket algorithm
"""
@pytest.mark.asyncio
async def test_algorithm_name(self) -> None:
algo = TokenBucketAlgorithm()
assert algo.name == "token_bucket"
@pytest.mark.asyncio
async def test_first_request_allowed(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
result = await algo.check(
storage=storage,
key="bucket_test",
rule=rule,
)
assert result.allowed is True
assert result.remaining == 99
await storage.close()
@pytest.mark.asyncio
async def test_burst_consumption(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=10, window_seconds=WINDOW_MINUTE)
for i in range(10):
result = await algo.check(storage, "burst_key", rule)
assert result.allowed is True
assert result.remaining == 10 - (i + 1)
await storage.close()
@pytest.mark.asyncio
async def test_bucket_exhausted(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
await algo.check(storage, "exhaust_key", rule)
result = await algo.check(storage, "exhaust_key", rule)
assert result.allowed is False
await storage.close()
@pytest.mark.asyncio
async def test_different_keys_independent(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
await algo.check(storage, "bucket_a", rule)
result_a = await algo.check(storage, "bucket_a", rule)
result_b = await algo.check(storage, "bucket_b", rule)
assert result_a.allowed is False
assert result_b.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_empty(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
usage = await algo.get_current_usage(storage, "empty_bucket", rule)
assert usage == 0
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_after_consumption(self) -> None:
algo = TokenBucketAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
for _ in range(25):
await algo.check(storage, "usage_bucket", rule)
usage = await algo.get_current_usage(storage, "usage_bucket", rule)
assert usage >= 24
await storage.close()
class TestFixedWindowAlgorithm:
"""
Tests for fixed window counter algorithm
"""
@pytest.mark.asyncio
async def test_algorithm_name(self) -> None:
algo = FixedWindowAlgorithm()
assert algo.name == "fixed_window"
@pytest.mark.asyncio
async def test_first_request_allowed(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
result = await algo.check(
storage=storage,
key="fixed_test",
rule=rule,
)
assert result.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_multiple_requests(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
for _ in range(50):
result = await algo.check(storage, "multi_fixed", rule)
assert result.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_limit_exceeded(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
result = await algo.check(storage, "limit_fixed", rule)
assert result.allowed is True
result = await algo.check(storage, "limit_fixed", rule)
assert result.allowed is False
await storage.close()
@pytest.mark.asyncio
async def test_different_keys_independent(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for _ in range(5):
await algo.check(storage, "fixed_a", rule)
result_a = await algo.check(storage, "fixed_a", rule)
result_b = await algo.check(storage, "fixed_b", rule)
assert result_a.allowed is False
assert result_b.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_empty(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
usage = await algo.get_current_usage(storage, "empty_fixed", rule)
assert usage == 0
await storage.close()
@pytest.mark.asyncio
async def test_get_current_usage_after_requests(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
for _ in range(25):
await algo.check(storage, "usage_fixed", rule)
usage = await algo.get_current_usage(storage, "usage_fixed", rule)
assert usage >= 0
await storage.close()
@pytest.mark.asyncio
async def test_explicit_timestamp(self) -> None:
algo = FixedWindowAlgorithm()
storage = MemoryStorage()
rule = RuleFactory.per_minute(100)
fixed_time = 1000000.0
result = await algo.check(
storage=storage,
key="timestamp_fixed",
rule=rule,
timestamp=fixed_time,
)
assert result.allowed is True
await storage.close()
class TestAlgorithmComparison:
"""
Comparative tests between algorithms
"""
@pytest.mark.asyncio
async def test_all_algorithms_allow_first_request(self) -> None:
algorithms = [
SlidingWindowAlgorithm(),
TokenBucketAlgorithm(),
FixedWindowAlgorithm(),
]
rule = RuleFactory.per_minute(100)
for algo in algorithms:
storage = MemoryStorage()
result = await algo.check(storage, "compare_key", rule)
assert result.allowed is True, f"{algo.name} failed first request"
await storage.close()
@pytest.mark.asyncio
async def test_all_algorithms_enforce_limit(self) -> None:
algorithms = [
SlidingWindowAlgorithm(),
TokenBucketAlgorithm(),
FixedWindowAlgorithm(),
]
rule = RuleFactory.create(requests=5, window_seconds=WINDOW_MINUTE)
for algo in algorithms:
storage = MemoryStorage()
for _ in range(5):
await algo.check(storage, "enforce_key", rule)
result = await algo.check(storage, "enforce_key", rule)
assert result.allowed is False, f"{algo.name} didn't enforce limit"
await storage.close()
@pytest.mark.asyncio
async def test_algorithms_have_correct_names(self) -> None:
assert SlidingWindowAlgorithm().name == Algorithm.SLIDING_WINDOW.value
assert TokenBucketAlgorithm().name == Algorithm.TOKEN_BUCKET.value
assert FixedWindowAlgorithm().name == Algorithm.FIXED_WINDOW.value

View File

@ -0,0 +1,421 @@
"""
AngelaMos | 2025
test_storage.py
"""
from __future__ import annotations
import asyncio
import time
import pytest
from fastapi_420.storage import MemoryStorage, create_storage
from fastapi_420.config import StorageSettings
from fastapi_420.types import StorageType
from tests.conftest import (
WINDOW_MINUTE,
WINDOW_SECOND,
DEFAULT_LIMIT_REQUESTS,
)
class TestMemoryStorageBasic:
"""
Basic MemoryStorage creation and lifecycle tests
"""
@pytest.mark.asyncio
async def test_create_storage(self) -> None:
storage = MemoryStorage()
assert storage.storage_type == StorageType.MEMORY
assert storage.max_keys == 100_000
@pytest.mark.asyncio
async def test_create_storage_custom_settings(self) -> None:
storage = MemoryStorage(max_keys=5000, cleanup_interval=30)
assert storage.max_keys == 5000
assert storage.cleanup_interval == 30
@pytest.mark.asyncio
async def test_from_settings(self) -> None:
settings = StorageSettings(
MEMORY_MAX_KEYS=2000,
MEMORY_CLEANUP_INTERVAL=120,
)
storage = MemoryStorage.from_settings(settings)
assert storage.max_keys == 2000
assert storage.cleanup_interval == 120
@pytest.mark.asyncio
async def test_health_check_healthy(self) -> None:
storage = MemoryStorage()
assert await storage.health_check() is True
@pytest.mark.asyncio
async def test_health_check_after_close(self) -> None:
storage = MemoryStorage()
await storage.close()
assert await storage.health_check() is False
@pytest.mark.asyncio
async def test_close_clears_data(self) -> None:
storage = MemoryStorage()
await storage.increment("test_key", WINDOW_MINUTE, 100)
await storage.close()
assert len(storage._windows) == 0
assert len(storage._buckets) == 0
class TestMemoryStorageSlidingWindow:
"""
Tests for sliding window counter in MemoryStorage
"""
@pytest.mark.asyncio
async def test_increment_first_request(self) -> None:
storage = MemoryStorage()
result = await storage.increment(
key="test",
window_seconds=WINDOW_MINUTE,
limit=100,
)
assert result.allowed is True
assert result.limit == 100
assert result.remaining == 99
await storage.close()
@pytest.mark.asyncio
async def test_increment_multiple_requests(self) -> None:
storage = MemoryStorage()
key = "multi_test"
for i in range(10):
result = await storage.increment(key, WINDOW_MINUTE, 100)
assert result.allowed is True
assert result.remaining == 100 - (i + 1)
await storage.close()
@pytest.mark.asyncio
async def test_increment_reaches_limit(self) -> None:
storage = MemoryStorage()
key = "limit_test"
limit = 5
for _ in range(limit):
result = await storage.increment(key, WINDOW_MINUTE, limit)
assert result.allowed is True
result = await storage.increment(key, WINDOW_MINUTE, limit)
assert result.allowed is False
assert result.remaining == 0
assert result.retry_after is not None
assert result.retry_after > 0
await storage.close()
@pytest.mark.asyncio
async def test_increment_with_explicit_timestamp(self) -> None:
storage = MemoryStorage()
fixed_time = 1000000.0
result = await storage.increment(
key="timestamp_test",
window_seconds=WINDOW_MINUTE,
limit=100,
timestamp=fixed_time,
)
assert result.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_get_window_state_empty(self) -> None:
storage = MemoryStorage()
state = await storage.get_window_state("nonexistent", WINDOW_MINUTE)
assert state.current_count == 0
assert state.previous_count == 0
await storage.close()
@pytest.mark.asyncio
async def test_get_window_state_with_data(self) -> None:
storage = MemoryStorage()
key = "state_test"
for _ in range(5):
await storage.increment(key, WINDOW_MINUTE, 100)
state = await storage.get_window_state(key, WINDOW_MINUTE)
assert state.current_count == 5
await storage.close()
@pytest.mark.asyncio
async def test_sliding_window_weighted_count(self) -> None:
storage = MemoryStorage()
key = "weighted_test"
window = 2
limit = 100
base_time = 1000.0
current_window = int(base_time // window)
previous_window = current_window - 1
prev_key = f"{key}:{previous_window}"
curr_key = f"{key}:{current_window}"
storage._windows[prev_key] = storage._windows.__class__().__class__
from fastapi_420.storage.memory import WindowEntry
storage._windows[prev_key] = WindowEntry(
count=50,
window_start=previous_window,
expires_at=base_time + window * 2,
)
result = await storage.increment(
key=key,
window_seconds=window,
limit=limit,
timestamp=base_time + 1.0,
)
assert result.allowed is True
await storage.close()
class TestMemoryStorageTokenBucket:
"""
Tests for token bucket algorithm in MemoryStorage
"""
@pytest.mark.asyncio
async def test_consume_token_first_request(self) -> None:
storage = MemoryStorage()
result = await storage.consume_token(
key="bucket_test",
capacity=100,
refill_rate=1.67,
tokens_to_consume=1,
)
assert result.allowed is True
assert result.remaining == 99
await storage.close()
@pytest.mark.asyncio
async def test_consume_token_multiple(self) -> None:
storage = MemoryStorage()
key = "multi_bucket"
for i in range(10):
result = await storage.consume_token(
key=key,
capacity=100,
refill_rate=1.67,
)
assert result.allowed is True
assert result.remaining == 100 - (i + 1)
await storage.close()
@pytest.mark.asyncio
async def test_consume_token_exhausted(self) -> None:
storage = MemoryStorage()
key = "exhaust_bucket"
capacity = 5
for _ in range(capacity):
result = await storage.consume_token(
key=key,
capacity=capacity,
refill_rate=1.0,
)
assert result.allowed is True
result = await storage.consume_token(
key=key,
capacity=capacity,
refill_rate=1.0,
)
assert result.allowed is False
assert result.retry_after is not None
await storage.close()
@pytest.mark.asyncio
async def test_consume_token_refill(self) -> None:
storage = MemoryStorage()
key = "refill_test"
capacity = 10
refill_rate = 10.0
for _ in range(capacity):
await storage.consume_token(key, capacity, refill_rate)
result = await storage.consume_token(key, capacity, refill_rate)
assert result.allowed is False
await asyncio.sleep(0.15)
result = await storage.consume_token(key, capacity, refill_rate)
assert result.allowed is True
await storage.close()
@pytest.mark.asyncio
async def test_get_token_bucket_state_empty(self) -> None:
storage = MemoryStorage()
state = await storage.get_token_bucket_state("nonexistent")
assert state is None
await storage.close()
@pytest.mark.asyncio
async def test_get_token_bucket_state_exists(self) -> None:
storage = MemoryStorage()
key = "bucket_state_test"
await storage.consume_token(key, capacity=100, refill_rate=1.67)
state = await storage.get_token_bucket_state(key)
assert state is not None
assert state.tokens == 99.0
assert state.capacity == 100
await storage.close()
class TestMemoryStorageMaxKeys:
"""
Tests for key eviction when max_keys is exceeded
"""
@pytest.mark.asyncio
async def test_max_keys_eviction(self) -> None:
storage = MemoryStorage(max_keys=5)
for i in range(10):
await storage.increment(f"key_{i}", WINDOW_MINUTE, 100)
assert len(storage._windows) <= 6
await storage.close()
@pytest.mark.asyncio
async def test_lru_eviction_order(self) -> None:
storage = MemoryStorage(max_keys=3)
await storage.increment("key_a", WINDOW_MINUTE, 100)
await storage.increment("key_b", WINDOW_MINUTE, 100)
await storage.increment("key_c", WINDOW_MINUTE, 100)
await storage.increment("key_d", WINDOW_MINUTE, 100)
keys = list(storage._windows.keys())
assert not any("key_a" in k for k in keys)
await storage.close()
class TestMemoryStorageCleanup:
"""
Tests for automatic cleanup of expired entries
"""
@pytest.mark.asyncio
async def test_cleanup_expired_entries(self) -> None:
storage = MemoryStorage(cleanup_interval=60)
from fastapi_420.storage.memory import WindowEntry
storage._windows["expired_key"] = WindowEntry(
count=10,
window_start=1,
expires_at=time.time() - 100,
)
storage._windows["valid_key"] = WindowEntry(
count=10,
window_start=1,
expires_at=time.time() + 100,
)
await storage._cleanup_expired()
assert "expired_key" not in storage._windows
assert "valid_key" in storage._windows
await storage.close()
@pytest.mark.asyncio
async def test_cleanup_task_starts(self) -> None:
storage = MemoryStorage(cleanup_interval=1)
await storage.start_cleanup_task()
assert storage._cleanup_task is not None
assert not storage._cleanup_task.done()
await storage.close()
@pytest.mark.asyncio
async def test_cleanup_task_stops_on_close(self) -> None:
storage = MemoryStorage(cleanup_interval=1)
await storage.start_cleanup_task()
task = storage._cleanup_task
await storage.close()
assert storage._cleanup_task is None
class TestStorageFactory:
"""
Tests for create_storage factory function
"""
def test_create_memory_storage_no_redis(self) -> None:
settings = StorageSettings(REDIS_URL=None)
storage = create_storage(settings)
assert isinstance(storage, MemoryStorage)
def test_create_memory_storage_explicit(self) -> None:
settings = StorageSettings(
REDIS_URL=None,
MEMORY_MAX_KEYS=5000,
)
storage = create_storage(settings)
assert isinstance(storage, MemoryStorage)
assert storage.max_keys == 5000
class TestMemoryStorageConcurrency:
"""
Tests for concurrent access to MemoryStorage
"""
@pytest.mark.asyncio
async def test_concurrent_increments(self) -> None:
storage = MemoryStorage()
key = "concurrent_test"
limit = 1000
async def increment() -> bool:
result = await storage.increment(key, WINDOW_MINUTE, limit)
return result.allowed
tasks = [increment() for _ in range(100)]
results = await asyncio.gather(*tasks)
assert all(results)
state = await storage.get_window_state(key, WINDOW_MINUTE)
assert state.current_count == 100
await storage.close()
@pytest.mark.asyncio
async def test_concurrent_different_keys(self) -> None:
storage = MemoryStorage()
limit = 100
async def increment(key: str) -> int:
for _ in range(10):
await storage.increment(key, WINDOW_MINUTE, limit)
state = await storage.get_window_state(key, WINDOW_MINUTE)
return state.current_count
tasks = [increment(f"key_{i}") for i in range(10)]
results = await asyncio.gather(*tasks)
assert all(r == 10 for r in results)
await storage.close()

View File

@ -0,0 +1,547 @@
"""
AngelaMos | 2025
test_types.py
"""
from __future__ import annotations
import pytest
from fastapi_420.types import (
Algorithm,
CircuitState,
DefenseContext,
DefenseMode,
FingerprintData,
FingerprintLevel,
Layer,
RateLimitKey,
RateLimitResult,
RateLimitRule,
StorageType,
TokenBucketState,
WindowState,
)
from tests.conftest import (
KEY_PREFIX,
KEY_VERSION,
TEST_ENDPOINT,
TEST_IP_V4,
TEST_USER_AGENT,
WINDOW_DAY,
WINDOW_HOUR,
WINDOW_MINUTE,
WINDOW_SECOND,
CircuitStateFactory,
DefenseContextFactory,
FingerprintFactory,
KeyFactory,
ResultFactory,
RuleFactory,
TokenBucketStateFactory,
WindowStateFactory,
)
class TestAlgorithmEnum:
"""
Tests for Algorithm enum values
"""
def test_sliding_window_value(self) -> None:
assert Algorithm.SLIDING_WINDOW.value == "sliding_window"
def test_token_bucket_value(self) -> None:
assert Algorithm.TOKEN_BUCKET.value == "token_bucket"
def test_fixed_window_value(self) -> None:
assert Algorithm.FIXED_WINDOW.value == "fixed_window"
def test_leaky_bucket_value(self) -> None:
assert Algorithm.LEAKY_BUCKET.value == "leaky_bucket"
def test_enum_is_str_enum(self) -> None:
assert isinstance(Algorithm.SLIDING_WINDOW, str)
assert Algorithm.SLIDING_WINDOW == "sliding_window"
class TestFingerprintLevelEnum:
"""
Tests for FingerprintLevel enum values
"""
def test_strict_value(self) -> None:
assert FingerprintLevel.STRICT.value == "strict"
def test_normal_value(self) -> None:
assert FingerprintLevel.NORMAL.value == "normal"
def test_relaxed_value(self) -> None:
assert FingerprintLevel.RELAXED.value == "relaxed"
def test_custom_value(self) -> None:
assert FingerprintLevel.CUSTOM.value == "custom"
class TestDefenseModeEnum:
"""
Tests for DefenseMode enum values
"""
def test_adaptive_value(self) -> None:
assert DefenseMode.ADAPTIVE.value == "adaptive"
def test_lockdown_value(self) -> None:
assert DefenseMode.LOCKDOWN.value == "lockdown"
def test_challenge_value(self) -> None:
assert DefenseMode.CHALLENGE.value == "challenge"
def test_disabled_value(self) -> None:
assert DefenseMode.DISABLED.value == "disabled"
class TestStorageTypeEnum:
"""
Tests for StorageType enum values
"""
def test_redis_value(self) -> None:
assert StorageType.REDIS.value == "redis"
def test_memory_value(self) -> None:
assert StorageType.MEMORY.value == "memory"
class TestLayerEnum:
"""
Tests for Layer enum values
"""
def test_user_value(self) -> None:
assert Layer.USER.value == "user"
def test_endpoint_value(self) -> None:
assert Layer.ENDPOINT.value == "endpoint"
def test_global_value(self) -> None:
assert Layer.GLOBAL.value == "global"
class TestRateLimitRule:
"""
Tests for RateLimitRule dataclass and parsing
"""
def test_create_valid_rule(self) -> None:
rule = RateLimitRule(requests=100, window_seconds=60)
assert rule.requests == 100
assert rule.window_seconds == 60
def test_invalid_requests_zero(self) -> None:
with pytest.raises(ValueError, match="requests must be positive"):
RateLimitRule(requests=0, window_seconds=60)
def test_invalid_requests_negative(self) -> None:
with pytest.raises(ValueError, match="requests must be positive"):
RateLimitRule(requests=-1, window_seconds=60)
def test_invalid_window_zero(self) -> None:
with pytest.raises(ValueError, match="window_seconds must be positive"):
RateLimitRule(requests=100, window_seconds=0)
def test_invalid_window_negative(self) -> None:
with pytest.raises(ValueError, match="window_seconds must be positive"):
RateLimitRule(requests=100, window_seconds=-1)
def test_parse_per_second(self) -> None:
for unit in ["second", "seconds", "sec", "s"]:
rule = RateLimitRule.parse(f"10/{unit}")
assert rule.requests == 10
assert rule.window_seconds == WINDOW_SECOND
def test_parse_per_minute(self) -> None:
for unit in ["minute", "minutes", "min", "m"]:
rule = RateLimitRule.parse(f"100/{unit}")
assert rule.requests == 100
assert rule.window_seconds == WINDOW_MINUTE
def test_parse_per_hour(self) -> None:
for unit in ["hour", "hours", "hr", "h"]:
rule = RateLimitRule.parse(f"1000/{unit}")
assert rule.requests == 1000
assert rule.window_seconds == WINDOW_HOUR
def test_parse_per_day(self) -> None:
for unit in ["day", "days", "d"]:
rule = RateLimitRule.parse(f"10000/{unit}")
assert rule.requests == 10000
assert rule.window_seconds == WINDOW_DAY
def test_parse_with_whitespace(self) -> None:
rule = RateLimitRule.parse(" 100 / minute ")
assert rule.requests == 100
assert rule.window_seconds == WINDOW_MINUTE
def test_parse_case_insensitive(self) -> None:
rule = RateLimitRule.parse("100/MINUTE")
assert rule.requests == 100
assert rule.window_seconds == WINDOW_MINUTE
def test_parse_invalid_format_no_slash(self) -> None:
with pytest.raises(ValueError, match="Invalid rate limit format"):
RateLimitRule.parse("100minute")
def test_parse_invalid_format_multiple_slashes(self) -> None:
with pytest.raises(ValueError, match="Invalid rate limit format"):
RateLimitRule.parse("100/per/minute")
def test_parse_invalid_request_count(self) -> None:
with pytest.raises(ValueError, match="Invalid request count"):
RateLimitRule.parse("abc/minute")
def test_parse_unknown_time_unit(self) -> None:
with pytest.raises(ValueError, match="Unknown time unit"):
RateLimitRule.parse("100/fortnight")
def test_str_representation_second(self) -> None:
rule = RateLimitRule(requests=10, window_seconds=1)
assert str(rule) == "10/second"
def test_str_representation_minute(self) -> None:
rule = RateLimitRule(requests=100, window_seconds=60)
assert str(rule) == "100/minute"
def test_str_representation_hour(self) -> None:
rule = RateLimitRule(requests=1000, window_seconds=3600)
assert str(rule) == "1000/hour"
def test_str_representation_day(self) -> None:
rule = RateLimitRule(requests=10000, window_seconds=86400)
assert str(rule) == "10000/day"
def test_str_representation_custom_window(self) -> None:
rule = RateLimitRule(requests=50, window_seconds=120)
assert str(rule) == "50/120s"
def test_frozen_immutable(self) -> None:
rule = RateLimitRule(requests=100, window_seconds=60)
with pytest.raises(AttributeError):
rule.requests = 200
def test_factory_create(self) -> None:
rule = RuleFactory.create()
assert rule.requests == 100
assert rule.window_seconds == WINDOW_MINUTE
def test_factory_per_second(self) -> None:
rule = RuleFactory.per_second(10)
assert rule.requests == 10
assert rule.window_seconds == WINDOW_SECOND
def test_factory_strict(self) -> None:
rule = RuleFactory.strict()
assert rule.requests == 10
assert rule.window_seconds == WINDOW_MINUTE
class TestRateLimitResult:
"""
Tests for RateLimitResult dataclass and headers
"""
def test_allowed_result(self) -> None:
result = RateLimitResult(
allowed=True,
limit=100,
remaining=50,
reset_after=30.0,
)
assert result.allowed is True
assert result.limit == 100
assert result.remaining == 50
assert result.reset_after == 30.0
assert result.retry_after is None
def test_denied_result(self) -> None:
result = RateLimitResult(
allowed=False,
limit=100,
remaining=0,
reset_after=60.0,
retry_after=60.0,
)
assert result.allowed is False
assert result.remaining == 0
assert result.retry_after == 60.0
def test_headers_basic(self) -> None:
result = RateLimitResult(
allowed=True,
limit=100,
remaining=50,
reset_after=30.5,
)
headers = result.headers
assert headers["RateLimit-Limit"] == "100"
assert headers["RateLimit-Remaining"] == "50"
assert headers["RateLimit-Reset"] == "30"
assert "Retry-After" not in headers
def test_headers_with_retry_after(self) -> None:
result = RateLimitResult(
allowed=False,
limit=100,
remaining=0,
reset_after=60.0,
retry_after=45.5,
)
headers = result.headers
assert headers["Retry-After"] == "45"
def test_headers_remaining_never_negative(self) -> None:
result = RateLimitResult(
allowed=False,
limit=100,
remaining=-5,
reset_after=60.0,
)
headers = result.headers
assert headers["RateLimit-Remaining"] == "0"
def test_frozen_immutable(self) -> None:
result = ResultFactory.allowed()
with pytest.raises(AttributeError):
result.allowed = False
def test_factory_allowed(self) -> None:
result = ResultFactory.allowed(limit=200, remaining=150)
assert result.allowed is True
assert result.limit == 200
assert result.remaining == 150
def test_factory_denied(self) -> None:
result = ResultFactory.denied(retry_after=30.0)
assert result.allowed is False
assert result.remaining == 0
assert result.retry_after == 30.0
def test_factory_near_limit(self) -> None:
result = ResultFactory.near_limit(remaining=2)
assert result.allowed is True
assert result.remaining == 2
class TestFingerprintData:
"""
Tests for FingerprintData and composite key generation
"""
def test_create_full_fingerprint(self) -> None:
fp = FingerprintData(
ip=TEST_IP_V4,
ip_normalized=TEST_IP_V4,
user_agent=TEST_USER_AGENT,
accept_language="en-US",
accept_encoding="gzip",
headers_hash="abc123",
auth_identifier="user_456",
tls_fingerprint="ja3hash",
geo_asn="AS12345",
)
assert fp.ip == TEST_IP_V4
assert fp.user_agent == TEST_USER_AGENT
assert fp.auth_identifier == "user_456"
def test_create_minimal_fingerprint(self) -> None:
fp = FingerprintData(ip=TEST_IP_V4, ip_normalized=TEST_IP_V4)
assert fp.ip == TEST_IP_V4
assert fp.user_agent is None
assert fp.auth_identifier is None
def test_composite_key_relaxed_ip_only(self) -> None:
fp = FingerprintFactory.anonymous()
key = fp.to_composite_key(FingerprintLevel.RELAXED)
assert key == TEST_IP_V4
def test_composite_key_relaxed_with_auth(self) -> None:
fp = FingerprintFactory.create(auth_identifier="user_123")
key = fp.to_composite_key(FingerprintLevel.RELAXED)
assert "user_123" in key
assert TEST_IP_V4 in key
def test_composite_key_normal(self) -> None:
fp = FingerprintFactory.create(auth_identifier="user_123")
key = fp.to_composite_key(FingerprintLevel.NORMAL)
parts = key.split(":")
assert len(parts) == 3
assert parts[0] == TEST_IP_V4
assert TEST_USER_AGENT in parts[1]
assert parts[2] == "user_123"
def test_composite_key_strict(self) -> None:
fp = FingerprintData(
ip=TEST_IP_V4,
ip_normalized=TEST_IP_V4,
user_agent=TEST_USER_AGENT,
accept_language="en-US",
accept_encoding="gzip",
headers_hash="abc123",
auth_identifier="user_456",
tls_fingerprint="ja3hash",
geo_asn="AS12345",
)
key = fp.to_composite_key(FingerprintLevel.STRICT)
parts = key.split(":")
assert len(parts) == 8
assert parts[0] == TEST_IP_V4
def test_factory_authenticated(self) -> None:
fp = FingerprintFactory.authenticated(auth_id="testuser")
assert fp.auth_identifier is not None
assert len(fp.auth_identifier) == 16
def test_factory_minimal(self) -> None:
fp = FingerprintFactory.minimal()
assert fp.ip == TEST_IP_V4
assert fp.user_agent is None
class TestWindowState:
"""
Tests for WindowState and weighted count calculation
"""
def test_create_empty_state(self) -> None:
state = WindowState()
assert state.current_count == 0
assert state.previous_count == 0
def test_weighted_count_start_of_window(self) -> None:
state = WindowState(current_count=50, previous_count=100)
weighted = state.weighted_count(0.0)
assert weighted == 150.0
def test_weighted_count_end_of_window(self) -> None:
state = WindowState(current_count=50, previous_count=100)
weighted = state.weighted_count(1.0)
assert weighted == 50.0
def test_weighted_count_mid_window(self) -> None:
state = WindowState(current_count=50, previous_count=100)
weighted = state.weighted_count(0.5)
assert weighted == 100.0
def test_weighted_count_quarter_window(self) -> None:
state = WindowState(current_count=40, previous_count=80)
weighted = state.weighted_count(0.25)
assert weighted == 80 * 0.75 + 40
def test_factory_empty(self) -> None:
state = WindowStateFactory.empty()
assert state.current_count == 0
assert state.previous_count == 0
def test_factory_with_usage(self) -> None:
state = WindowStateFactory.with_usage(current=50, previous=100)
assert state.current_count == 50
assert state.previous_count == 100
class TestTokenBucketState:
"""
Tests for TokenBucketState
"""
def test_create_state(self) -> None:
state = TokenBucketState(
tokens=50.0,
last_refill=1000.0,
capacity=100,
refill_rate=1.67,
)
assert state.tokens == 50.0
assert state.capacity == 100
assert state.refill_rate == 1.67
def test_factory_full(self) -> None:
state = TokenBucketStateFactory.full(capacity=200)
assert state.tokens == 200.0
assert state.capacity == 200
def test_factory_empty(self) -> None:
state = TokenBucketStateFactory.empty()
assert state.tokens == 0.0
class TestCircuitState:
"""
Tests for CircuitState
"""
def test_default_closed(self) -> None:
state = CircuitState()
assert state.is_open is False
assert state.failure_count == 0
def test_factory_closed(self) -> None:
state = CircuitStateFactory.closed()
assert state.is_open is False
def test_factory_open(self) -> None:
state = CircuitStateFactory.open()
assert state.is_open is True
assert state.failure_count == 1
assert state.last_failure_time > 0
class TestDefenseContext:
"""
Tests for DefenseContext
"""
def test_create_context(self) -> None:
fp = FingerprintFactory.create()
context = DefenseContext(
fingerprint=fp,
endpoint=TEST_ENDPOINT,
method="GET",
is_authenticated=True,
reputation_score=0.9,
)
assert context.endpoint == TEST_ENDPOINT
assert context.is_authenticated is True
assert context.reputation_score == 0.9
def test_factory_authenticated(self) -> None:
context = DefenseContextFactory.authenticated()
assert context.is_authenticated is True
assert context.fingerprint.auth_identifier is not None
def test_factory_suspicious(self) -> None:
context = DefenseContextFactory.suspicious(reputation_score=0.2)
assert context.reputation_score == 0.2
assert context.request_count_last_minute == 500
class TestRateLimitKey:
"""
Tests for RateLimitKey and key building
"""
def test_build_key(self) -> None:
key = RateLimitKey(
prefix=KEY_PREFIX,
version=KEY_VERSION,
layer=Layer.USER,
endpoint=TEST_ENDPOINT,
identifier=TEST_IP_V4,
window=WINDOW_MINUTE,
)
built = key.build()
expected = f"{KEY_PREFIX}:{KEY_VERSION}:user:{TEST_ENDPOINT}:{TEST_IP_V4}:{WINDOW_MINUTE}"
assert built == expected
def test_factory_user_key(self) -> None:
key = KeyFactory.user_key()
built = key.build()
assert ":user:" in built
assert TEST_ENDPOINT in built
def test_factory_endpoint_key(self) -> None:
key = KeyFactory.endpoint_key()
built = key.build()
assert ":endpoint:" in built
assert ":global:" in built
def test_factory_global_key(self) -> None:
key = KeyFactory.global_key()
built = key.build()
assert ":global:" in built

File diff suppressed because it is too large Load Diff