89 lines
3.0 KiB
Python
89 lines
3.0 KiB
Python
|
"""登录防爆破：内存滑动窗口限速（公网暴露前的最低防线）。
|
||
|
|
|
||
|
|
- 按 (IP, 用户名) 组合计数，失败 5 次 / 10 分钟 → 锁定 15 分钟；
|
||
|
|
- 同一 IP 全局失败 30 次 / 10 分钟 → 锁定 15 分钟（防换用户名扫号）；
|
||
|
|
- 成功登录清零该组合的计数。
|
||
|
|
内存实现，单进程部署足够；重启即清零（可接受的测试期取舍）。
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
|
||
|
|
from fastapi import HTTPException
|
||
|
|
|
||
|
|
FAIL_WINDOW_SECONDS = 600 # 10 分钟滑动窗口
|
||
|
|
FAIL_LIMIT_PER_KEY = 5 # 同 IP+用户名 失败上限
|
||
|
|
FAIL_LIMIT_PER_IP = 30 # 同 IP 全局失败上限
|
||
|
|
LOCK_SECONDS = 900 # 锁定时长
|
||
|
|
|
||
|
|
_lock = threading.Lock()
|
||
|
|
_fails: dict[str, list[float]] = {} # key -> 失败时间戳列表
|
||
|
|
_locked_until: dict[str, float] = {} # key -> 锁定截止时间
|
||
|
|
|
||
|
|
|
||
|
|
def _prune(now: float) -> None:
|
||
|
|
cutoff = now - FAIL_WINDOW_SECONDS
|
||
|
|
for key in list(_fails):
|
||
|
|
rows = [t for t in _fails[key] if t > cutoff]
|
||
|
|
if rows:
|
||
|
|
_fails[key] = rows
|
||
|
|
else:
|
||
|
|
_fails.pop(key, None)
|
||
|
|
for key in list(_locked_until):
|
||
|
|
if _locked_until[key] < now:
|
||
|
|
_locked_until.pop(key, None)
|
||
|
|
|
||
|
|
|
||
|
|
def _guard_locked(key: str, now: float) -> None:
|
||
|
|
until = _locked_until.get(key)
|
||
|
|
if until and until > now:
|
||
|
|
remain = int((until - now) / 60) + 1
|
||
|
|
raise HTTPException(429, f"尝试次数过多，请约 {remain} 分钟后再试")
|
||
|
|
|
||
|
|
|
||
|
|
def check_login_allowed(ip: str, username: str) -> None:
|
||
|
|
"""登录前调用：命中锁定则抛 429。"""
|
||
|
|
now = time.time()
|
||
|
|
with _lock:
|
||
|
|
_prune(now)
|
||
|
|
_guard_locked(f"key:{ip}:{username}", now)
|
||
|
|
_guard_locked(f"ip:{ip}", now)
|
||
|
|
|
||
|
|
|
||
|
|
def record_login_failure(ip: str, username: str) -> None:
|
||
|
|
"""登录失败后调用：达到阈值即锁定。"""
|
||
|
|
now = time.time()
|
||
|
|
with _lock:
|
||
|
|
_prune(now)
|
||
|
|
key = f"key:{ip}:{username}"
|
||
|
|
_fails.setdefault(key, []).append(now)
|
||
|
|
if len(_fails[key]) >= FAIL_LIMIT_PER_KEY:
|
||
|
|
_locked_until[key] = now + LOCK_SECONDS
|
||
|
|
ip_key = f"ip:{ip}"
|
||
|
|
_fails.setdefault(ip_key, []).append(now)
|
||
|
|
if len(_fails[ip_key]) >= FAIL_LIMIT_PER_IP:
|
||
|
|
_locked_until[ip_key] = now + LOCK_SECONDS
|
||
|
|
|
||
|
|
|
||
|
|
def record_login_success(ip: str, username: str) -> None:
|
||
|
|
"""登录成功后清零该组合计数（IP 级计数保留）。"""
|
||
|
|
with _lock:
|
||
|
|
_fails.pop(f"key:{ip}:{username}", None)
|
||
|
|
_locked_until.pop(f"key:{ip}:{username}", None)
|
||
|
|
|
||
|
|
|
||
|
|
def client_ip(request) -> str:
|
||
|
|
"""优先取反代/隧道传递的真实 IP。"""
|
||
|
|
forwarded = request.headers.get("x-forwarded-for") if request else None
|
||
|
|
if forwarded:
|
||
|
|
return forwarded.split(",")[0].strip()
|
||
|
|
real = request.headers.get("x-real-ip") if request else None
|
||
|
|
if real:
|
||
|
|
return real.strip()
|
||
|
|
client = request.client if request else None
|
||
|
|
return client.host if client else "unknown"
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = ["check_login_allowed", "record_login_failure", "record_login_success", "client_ip"]
|