Kris
849d21120f
新增一系列安全相关功能： 1. 新增二维码生成工具，支持自定义参数导出图片/Base64/字节流 2. 新增HTML内容安全处理工具，包括标签剥离、转义、URL校验 3. 新增日志敏感信息脱敏工具，支持配置、字典、日志记录脱敏 4. 新增密钥管理运行时，支持从环境变量/文件加载密钥 5. 新增身份链接管理，支持多渠道身份绑定与解析 6. 新增SSRF防护工具，支持域名/IP校验与固定主机 7. 新增安全权限修复工具，修复文件目录权限与配置项 8. 新增外部内容安全处理，支持LLM特殊令牌剥离与注入检测 9. 新增认证限流工具，支持多维度限流与本地回环豁免 10. 新增配对管理工具，支持安全配对码生成与校验 11. 新增白名单管理工具，支持DM/群组/来源白名单校验
237 lines
8.4 KiB
Python
237 lines
8.4 KiB
Python
import asyncio
|
||
import contextlib
|
||
import time
|
||
from collections import defaultdict
|
||
from collections.abc import Callable
|
||
from dataclasses import dataclass, field
|
||
from enum import StrEnum
|
||
from functools import wraps
|
||
|
||
|
||
class RateLimitScope(StrEnum):
|
||
LOGIN = "login"
|
||
API = "api"
|
||
SIGNUP = "signup"
|
||
PASSWORD_RESET = "password_reset"
|
||
TOKEN_REFRESH = "token_refresh"
|
||
|
||
|
||
@dataclass
|
||
class ScopeConfig:
|
||
max_attempts: int = 10
|
||
window_seconds: int = 60
|
||
cooldown_seconds: int = 300
|
||
loopback_exempt: bool = True
|
||
|
||
|
||
DEFAULT_SCOPE_CONFIGS = {
|
||
RateLimitScope.LOGIN: ScopeConfig(max_attempts=10, window_seconds=60, cooldown_seconds=300),
|
||
RateLimitScope.API: ScopeConfig(max_attempts=100, window_seconds=60, cooldown_seconds=60),
|
||
RateLimitScope.SIGNUP: ScopeConfig(max_attempts=3, window_seconds=3600, cooldown_seconds=7200),
|
||
RateLimitScope.PASSWORD_RESET: ScopeConfig(max_attempts=5, window_seconds=300, cooldown_seconds=600),
|
||
RateLimitScope.TOKEN_REFRESH: ScopeConfig(max_attempts=30, window_seconds=60, cooldown_seconds=120),
|
||
}
|
||
|
||
_LOCAL_IPS = frozenset({"127.0.0.1", "::1", "localhost"})
|
||
|
||
|
||
@dataclass
|
||
class _BucketState:
|
||
timestamps: list[float] = field(default_factory=list)
|
||
cooldown_until: float = 0.0
|
||
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||
|
||
|
||
class AuthRateLimiter:
|
||
"""认证级限流器：内存滑动窗口 + 多 Scope + 本地回环豁免 + 锁定冷却。
|
||
|
||
特性：
|
||
- 多 Scope 独立配置（login/api/signup/password_reset/token_refresh）
|
||
- 滑动窗口计数，窗口外记录自动淘汰
|
||
- 本地回环 (127.0.0.1, ::1) 可豁免
|
||
- 锁定冷却期 5min（超限后拒绝所有尝试）
|
||
- asyncio.Lock 串行化防竞态
|
||
- 后台定期清理失效 bucket，防止 DDoS 下内存持续增长
|
||
"""
|
||
|
||
_CLEANUP_INTERVAL = 300 # seconds between stale bucket cleanups
|
||
|
||
def __init__(self, scope_configs: dict[RateLimitScope, ScopeConfig] | None = None):
|
||
self._configs = scope_configs or DEFAULT_SCOPE_CONFIGS
|
||
self._buckets: dict[str, dict[RateLimitScope, _BucketState]] = defaultdict(dict)
|
||
self._lock = asyncio.Lock()
|
||
self._last_cleanup: float = 0.0
|
||
self._cleanup_task: asyncio.Task | None = None
|
||
self._start_background_cleanup()
|
||
|
||
def _is_loopback(self, identifier: str) -> bool:
|
||
if identifier in _LOCAL_IPS:
|
||
return True
|
||
host = identifier.split(":")[0]
|
||
if host.startswith("127."):
|
||
return True
|
||
try:
|
||
import ipaddress
|
||
|
||
ip = ipaddress.ip_address(host)
|
||
return ip.is_loopback
|
||
except ValueError:
|
||
return False
|
||
|
||
async def _get_bucket(self, identifier: str, scope: RateLimitScope) -> _BucketState:
|
||
async with self._lock:
|
||
scope_buckets = self._buckets[identifier]
|
||
if scope not in scope_buckets:
|
||
scope_buckets[scope] = _BucketState()
|
||
return scope_buckets[scope]
|
||
|
||
async def _cleanup_stale_buckets(self, now: float) -> None:
|
||
async with self._lock:
|
||
stale_identifiers: list[str] = []
|
||
for identifier, scopes in self._buckets.items():
|
||
active_scopes: dict[RateLimitScope, _BucketState] = {}
|
||
for scope, bucket in scopes.items():
|
||
config = self._configs.get(scope)
|
||
window = config.window_seconds if config else 60
|
||
bucket.timestamps = [t for t in bucket.timestamps if t > now - window]
|
||
has_activity = bucket.timestamps or bucket.cooldown_until > now
|
||
if has_activity:
|
||
active_scopes[scope] = bucket
|
||
if active_scopes:
|
||
self._buckets[identifier] = active_scopes
|
||
else:
|
||
stale_identifiers.append(identifier)
|
||
for identifier in stale_identifiers:
|
||
del self._buckets[identifier]
|
||
|
||
def _start_background_cleanup(self) -> None:
|
||
async def _loop() -> None:
|
||
while True:
|
||
await asyncio.sleep(self._CLEANUP_INTERVAL)
|
||
await self._cleanup_stale_buckets(time.monotonic())
|
||
|
||
try:
|
||
loop = asyncio.get_running_loop()
|
||
self._cleanup_task = loop.create_task(_loop())
|
||
except RuntimeError:
|
||
pass
|
||
|
||
async def shutdown(self) -> None:
|
||
if self._cleanup_task and not self._cleanup_task.done():
|
||
self._cleanup_task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError):
|
||
await self._cleanup_task
|
||
|
||
async def check_and_record(
|
||
self,
|
||
identifier: str,
|
||
scope: RateLimitScope,
|
||
) -> bool:
|
||
"""检查是否允许请求，并记录本次尝试。
|
||
|
||
返回 True 表示允许，False 表示被限流。
|
||
"""
|
||
config = self._configs.get(scope)
|
||
if config is None:
|
||
return True
|
||
|
||
if config.loopback_exempt and self._is_loopback(identifier):
|
||
return True
|
||
|
||
now = time.monotonic()
|
||
|
||
if now - self._last_cleanup > self._CLEANUP_INTERVAL:
|
||
self._last_cleanup = now
|
||
await self._cleanup_stale_buckets(now)
|
||
|
||
bucket = await self._get_bucket(identifier, scope)
|
||
now = time.monotonic()
|
||
|
||
async with bucket.lock:
|
||
if bucket.cooldown_until > now:
|
||
return False
|
||
|
||
window_start = now - config.window_seconds
|
||
bucket.timestamps = [t for t in bucket.timestamps if t > window_start]
|
||
|
||
if len(bucket.timestamps) >= config.max_attempts:
|
||
bucket.cooldown_until = now + config.cooldown_seconds
|
||
return False
|
||
|
||
bucket.timestamps.append(now)
|
||
return True
|
||
|
||
async def check_and_record_or_raise(
|
||
self,
|
||
identifier: str,
|
||
scope: RateLimitScope,
|
||
) -> None:
|
||
if not await self.check_and_record(identifier, scope):
|
||
config = self._configs.get(scope)
|
||
retry_after = config.cooldown_seconds if config else 60
|
||
raise RateLimitExceeded(scope=scope, retry_after=retry_after)
|
||
|
||
async def clear_cooldown(self, identifier: str, scope: RateLimitScope):
|
||
bucket = await self._get_bucket(identifier, scope)
|
||
async with bucket.lock:
|
||
bucket.cooldown_until = 0.0
|
||
|
||
async def reset(self, identifier: str, scope: RateLimitScope | None = None):
|
||
async with self._lock:
|
||
if scope:
|
||
if identifier in self._buckets and scope in self._buckets[identifier]:
|
||
async with self._buckets[identifier][scope].lock:
|
||
self._buckets[identifier].pop(scope, None)
|
||
else:
|
||
self._buckets.pop(identifier, None)
|
||
|
||
def get_stats(self) -> dict:
|
||
now = time.monotonic()
|
||
stats: dict[str, dict] = {}
|
||
for identifier, scopes in self._buckets.items():
|
||
stats[identifier] = {}
|
||
for scope, bucket in scopes.items():
|
||
config = self._configs.get(scope)
|
||
window = config.window_seconds if config else 60
|
||
active = len([t for t in bucket.timestamps if t > now - window])
|
||
stats[identifier][scope.value] = {
|
||
"attempts_in_window": active,
|
||
"cooldown_active": bucket.cooldown_until > now,
|
||
"cooldown_remaining": max(0, bucket.cooldown_until - now),
|
||
}
|
||
return stats
|
||
|
||
def wrap(
|
||
self,
|
||
scope: RateLimitScope,
|
||
identifier_fn: Callable[..., str] | None = None,
|
||
):
|
||
"""装饰器：将限流包装到异步函数上。
|
||
|
||
使用方式：
|
||
@auth_rate_limiter.wrap(RateLimitScope.LOGIN)
|
||
async def login(request: Request):
|
||
...
|
||
"""
|
||
|
||
def decorator(fn):
|
||
@wraps(fn)
|
||
async def wrapper(*args, **kwargs):
|
||
ident = identifier_fn(*args, **kwargs) if identifier_fn else "default"
|
||
await self.check_and_record_or_raise(ident, scope)
|
||
return await fn(*args, **kwargs)
|
||
|
||
return wrapper
|
||
|
||
return decorator
|
||
|
||
|
||
class RateLimitExceeded(Exception):
|
||
def __init__(self, scope: RateLimitScope, retry_after: int = 60):
|
||
self.scope = scope
|
||
self.retry_after = retry_after
|
||
super().__init__(f"Rate limit exceeded for scope '{scope.value}', retry after {retry_after}s")
|
||
|
||
|
||
auth_rate_limiter = AuthRateLimiter()
|