120 lines
4.8 KiB
Python
120 lines
4.8 KiB
Python
"""
|
||
安全工具模块。
|
||
|
||
职责：
|
||
- 邮箱格式校验（统一规则，各路由复用）
|
||
- 客户端 IP 提取（兼容 Nginx X-Forwarded-For）
|
||
- 内存限流器：登录失败锁定、验证码发送/校验频率控制
|
||
|
||
说明：
|
||
- 本项目为单进程 uvicorn 部署，内存限流器即可满足需求；
|
||
若未来改为多进程部署，可将实现替换为 Redis 等共享存储。
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import time
|
||
from collections import defaultdict, deque
|
||
from threading import Lock
|
||
from typing import Deque, Dict
|
||
|
||
from fastapi import Request
|
||
|
||
|
||
def _env_int(name: str, default: int) -> int:
|
||
"""读取整数型环境变量，缺失或非法时返回默认值。"""
|
||
try:
|
||
return int(os.getenv(name, "").strip() or default)
|
||
except ValueError:
|
||
return default
|
||
|
||
|
||
# ---------- 邮箱校验 ----------
|
||
|
||
EMAIL_PATTERN = r"^[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}$"
|
||
|
||
|
||
def is_valid_email(email: str) -> bool:
|
||
"""校验邮箱基本格式，返回是否合法。"""
|
||
return bool(re.match(EMAIL_PATTERN, email or ""))
|
||
|
||
|
||
# ---------- 客户端 IP ----------
|
||
|
||
def get_client_ip(request: Request) -> str:
|
||
"""获取客户端 IP：优先信任 Nginx 写入的 X-Forwarded-For，否则取直连地址。"""
|
||
forwarded = request.headers.get("x-forwarded-for")
|
||
if forwarded:
|
||
return forwarded.split(",")[0].strip()
|
||
return request.client.host if request.client else "unknown"
|
||
|
||
|
||
# ---------- 内存限流器 ----------
|
||
|
||
class RateLimiter:
|
||
"""滑动窗口限流器：记录 key 在窗口内的访问次数，超过阈值判定为阻塞。"""
|
||
|
||
def __init__(self, max_events: int, window_seconds: int):
|
||
self.max_events = max_events
|
||
self.window_seconds = window_seconds
|
||
self._records: Dict[str, Deque[float]] = defaultdict(deque)
|
||
self._lock = Lock()
|
||
|
||
def _prune(self, key: str, now: float) -> None:
|
||
"""清理窗口外的历史记录，避免内存无限增长。"""
|
||
queue = self._records[key]
|
||
while queue and now - queue[0] > self.window_seconds:
|
||
queue.popleft()
|
||
|
||
def hit(self, key: str) -> int:
|
||
"""记录一次事件，返回窗口内的总次数。"""
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
self._prune(key, now)
|
||
self._records[key].append(now)
|
||
return len(self._records[key])
|
||
|
||
def is_blocked(self, key: str) -> bool:
|
||
"""判断当前是否已达阈值（阻塞）。"""
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
self._prune(key, now)
|
||
return len(self._records[key]) >= self.max_events
|
||
|
||
def reset(self, key: str) -> None:
|
||
"""清空指定 key 的记录（如登录成功后重置失败计数）。"""
|
||
with self._lock:
|
||
self._records.pop(key, None)
|
||
|
||
|
||
# ---------- 限流策略（阈值与窗口，均可通过 .env 调整，默认值见下） ----------
|
||
|
||
# 登录失败：同一邮箱在窗口内失败达到上限后临时锁定，防止暴力破解博主账号
|
||
LOGIN_MAX_FAILURES = _env_int("LOGIN_MAX_FAILURES", 5)
|
||
LOGIN_WINDOW_SECONDS = _env_int("LOGIN_WINDOW_SECONDS", 15 * 60)
|
||
login_failure_limiter = RateLimiter(LOGIN_MAX_FAILURES, LOGIN_WINDOW_SECONDS)
|
||
|
||
# 验证码发送：同一邮箱每小时最多 N 封、同一 IP 每小时最多 N 封，防止被当作垃圾邮件中继
|
||
SEND_CODE_MAX_PER_EMAIL = _env_int("SEND_CODE_MAX_PER_EMAIL", 5)
|
||
SEND_CODE_MAX_PER_IP = _env_int("SEND_CODE_MAX_PER_IP", 10)
|
||
SEND_CODE_WINDOW_SECONDS = _env_int("SEND_CODE_WINDOW_SECONDS", 60 * 60)
|
||
send_code_email_limiter = RateLimiter(SEND_CODE_MAX_PER_EMAIL, SEND_CODE_WINDOW_SECONDS)
|
||
send_code_ip_limiter = RateLimiter(SEND_CODE_MAX_PER_IP, SEND_CODE_WINDOW_SECONDS)
|
||
|
||
# 验证码校验：同一邮箱在验证码有效期内错误尝试达到上限后作废验证码，防 6 位验证码被暴力枚举
|
||
VERIFY_CODE_MAX_ATTEMPTS = _env_int("VERIFY_CODE_MAX_ATTEMPTS", 5)
|
||
VERIFY_CODE_WINDOW_SECONDS = _env_int("VERIFY_CODE_WINDOW_SECONDS", 10 * 60)
|
||
verify_code_limiter = RateLimiter(VERIFY_CODE_MAX_ATTEMPTS, VERIFY_CODE_WINDOW_SECONDS)
|
||
# 注册限流：同一 IP 每小时最多 3 次注册，防止批量注册垃圾账号（可通过 .env 调整）
|
||
REGISTER_MAX_PER_IP = _env_int("REGISTER_MAX_PER_IP", 3)
|
||
REGISTER_WINDOW_SECONDS = _env_int("REGISTER_WINDOW_SECONDS", 60 * 60)
|
||
register_ip_limiter = RateLimiter(REGISTER_MAX_PER_IP, REGISTER_WINDOW_SECONDS)
|
||
|
||
# 登录限流（IP 维度）：同一 IP 在登录窗口内最多尝试 30 次，配合邮箱维度防分布式爆破
|
||
LOGIN_MAX_PER_IP = _env_int("LOGIN_MAX_PER_IP", 30)
|
||
login_ip_limiter = RateLimiter(LOGIN_MAX_PER_IP, LOGIN_WINDOW_SECONDS)
|
||
|
||
# 写操作限流（评论/点赞）：同一 IP 每分钟最多 30 次，防止好友账号刷屏
|
||
ACTION_MAX_PER_MINUTE = _env_int("ACTION_MAX_PER_MINUTE", 30)
|
||
action_limiter = RateLimiter(ACTION_MAX_PER_MINUTE, 60)
|