87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
"""P0-5 · 基于 Django cache 的通用限流/防爆破计数器。
|
||
|
||
设计：
|
||
- 后端无关：locmem（开发/测试）与 Redis（生产）都可用同一 API。
|
||
- 登录防爆破：按 `(用户名 + IP)` 记连续失败次数，达阈值锁定 M 分钟；
|
||
成功登录清零。失败计数只在"凭证错误"时 +1，避免把不相关 400 计入。
|
||
- API Key 限流：按 key prefix 记滑动分钟窗口计数，超 `rate_limit` 抛 Throttled。
|
||
- `rate_limit` 下界：<=0 视为 1（防止被设成 0/负数导致全拒绝或无限）。
|
||
|
||
常量：
|
||
- LOGIN_MAX_FAILURES=5, LOGIN_LOCK_SECONDS=900（15 分钟）
|
||
- API KEY 默认窗口 60s，限额取 `max(1, key.rate_limit)`。
|
||
"""
|
||
|
||
from django.core.cache import cache
|
||
|
||
LOGIN_MAX_FAILURES = 5
|
||
LOGIN_LOCK_SECONDS = 15 * 60
|
||
APIKEY_WINDOW_SECONDS = 60
|
||
|
||
|
||
def _login_fail_key(username: str, ip: str) -> str:
|
||
return f"dealerhub:login-fail:{username}:{ip}"
|
||
|
||
|
||
def _login_lock_key(username: str, ip: str) -> str:
|
||
return f"dealerhub:login-lock:{username}:{ip}"
|
||
|
||
|
||
def is_login_locked(username: str, ip: str) -> bool:
|
||
return bool(cache.get(_login_lock_key(username, ip)))
|
||
|
||
|
||
def record_login_failure(username: str, ip: str) -> int:
|
||
"""失败 +1；达阈值则加锁并返回当前失败数。"""
|
||
fails = (cache.get(_login_fail_key(username, ip)) or 0) + 1
|
||
cache.set(_login_fail_key(username, ip), fails, LOGIN_LOCK_SECONDS)
|
||
if fails >= LOGIN_MAX_FAILURES:
|
||
cache.set(_login_lock_key(username, ip), 1, LOGIN_LOCK_SECONDS)
|
||
return fails
|
||
|
||
|
||
def clear_login_failures(username: str, ip: str) -> None:
|
||
cache.delete(_login_fail_key(username, ip))
|
||
cache.delete(_login_lock_key(username, ip))
|
||
|
||
|
||
def effective_rate_limit(raw) -> int:
|
||
"""rate_limit 下界保护：非正数/非法值 → 1。"""
|
||
try:
|
||
v = int(raw)
|
||
except (TypeError, ValueError):
|
||
return 1
|
||
return v if v >= 1 else 1
|
||
|
||
|
||
def check_apikey_rate_limit(key_prefix: str, limit: int) -> None:
|
||
"""分钟窗口计数；超限抛 DRF Throttled（429 + Retry-After 由 DRF 补）。
|
||
|
||
注意：必须传 `wait=`，DRF 只在 `exc.wait` 非空时才回 `Retry-After` 头
|
||
（`rest_framework/views.py::exception_handler`）。不传则只有 429 无头，
|
||
客户端无法知道何时重试。
|
||
"""
|
||
from rest_framework.exceptions import Throttled
|
||
|
||
window_key = f"dealerhub:apikey:{key_prefix}"
|
||
count = cache.get(window_key)
|
||
if count is None:
|
||
cache.set(window_key, 1, APIKEY_WINDOW_SECONDS)
|
||
return
|
||
if int(count) >= effective_rate_limit(limit):
|
||
raise Throttled(
|
||
wait=APIKEY_WINDOW_SECONDS,
|
||
detail="API Key 请求频率超限，请稍后重试",
|
||
)
|
||
try:
|
||
cache.incr(window_key)
|
||
except ValueError:
|
||
cache.set(window_key, 1, APIKEY_WINDOW_SECONDS)
|
||
|
||
|
||
def client_ip(request) -> str:
|
||
xff = request.META.get("HTTP_X_FORWARDED_FOR", "")
|
||
if xff:
|
||
return xff.split(",")[0].strip()
|
||
return request.META.get("REMOTE_ADDR", "") or "unknown"
|