Kris
849d21120f
新增一系列安全相关功能： 1. 新增二维码生成工具，支持自定义参数导出图片/Base64/字节流 2. 新增HTML内容安全处理工具，包括标签剥离、转义、URL校验 3. 新增日志敏感信息脱敏工具，支持配置、字典、日志记录脱敏 4. 新增密钥管理运行时，支持从环境变量/文件加载密钥 5. 新增身份链接管理，支持多渠道身份绑定与解析 6. 新增SSRF防护工具，支持域名/IP校验与固定主机 7. 新增安全权限修复工具，修复文件目录权限与配置项 8. 新增外部内容安全处理，支持LLM特殊令牌剥离与注入检测 9. 新增认证限流工具，支持多维度限流与本地回环豁免 10. 新增配对管理工具，支持安全配对码生成与校验 11. 新增白名单管理工具，支持DM/群组/来源白名单校验
235 lines
9.0 KiB
Python
235 lines
9.0 KiB
Python
import asyncio
|
|
import hmac
|
|
import logging
|
|
import secrets
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from yuxi.repositories.channel_pairing_repo import ChannelPairingRecordRepository
|
|
|
|
PAIRING_CODE_LENGTH = 8
|
|
PAIRING_CODE_ALPHABET = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
|
|
PAIRING_CODE_MAX_ATTEMPTS = 500
|
|
PAIRING_PENDING_TTL_SECONDS = 600
|
|
MAX_PENDING_PER_ACCOUNT = 3
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _generate_code(alphabet: str, length: int) -> str:
|
|
return "".join(secrets.choice(alphabet) for _ in range(length))
|
|
|
|
|
|
def _generate_token() -> str:
|
|
return secrets.token_hex(32)
|
|
|
|
|
|
@dataclass
|
|
class PendingPairingRequest:
|
|
channel_type: str
|
|
peer_id: str
|
|
account_id: str
|
|
code: str
|
|
created_at: float
|
|
_record_id: str | None = None
|
|
|
|
|
|
class PairingError(Exception):
|
|
pass
|
|
|
|
|
|
class PairingStore(ABC):
|
|
@abstractmethod
|
|
async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]: ...
|
|
@abstractmethod
|
|
async def upsert_pending(self, req: PendingPairingRequest) -> None: ...
|
|
@abstractmethod
|
|
async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None: ...
|
|
|
|
async def expire_pending(self, channel_type: str, peer_id: str, account_id: str) -> None:
|
|
await self.remove_pending(channel_type, peer_id, account_id)
|
|
|
|
|
|
class InMemoryPairingStore(PairingStore):
|
|
def __init__(self):
|
|
self._pending: dict[str, dict[str, PendingPairingRequest]] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
@staticmethod
|
|
def _account_key(channel_type: str, account_id: str) -> str:
|
|
return f"{channel_type}:{account_id}"
|
|
|
|
async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]:
|
|
async with self._lock:
|
|
now = time.monotonic()
|
|
account_key = self._account_key(channel_type, account_id)
|
|
bucket = self._pending.get(account_key)
|
|
if bucket is None:
|
|
return []
|
|
result: list[PendingPairingRequest] = []
|
|
expired_peers: list[str] = []
|
|
for peer_id, v in bucket.items():
|
|
if now - v.created_at >= PAIRING_PENDING_TTL_SECONDS:
|
|
expired_peers.append(peer_id)
|
|
continue
|
|
result.append(v)
|
|
for peer_id in expired_peers:
|
|
del bucket[peer_id]
|
|
if not bucket:
|
|
del self._pending[account_key]
|
|
return result
|
|
|
|
async def upsert_pending(self, req: PendingPairingRequest) -> None:
|
|
account_key = self._account_key(req.channel_type, req.account_id)
|
|
async with self._lock:
|
|
bucket = self._pending.setdefault(account_key, {})
|
|
bucket[req.peer_id] = req
|
|
|
|
async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None:
|
|
account_key = self._account_key(channel_type, account_id)
|
|
async with self._lock:
|
|
bucket = self._pending.get(account_key)
|
|
if bucket is None:
|
|
return
|
|
bucket.pop(peer_id, None)
|
|
if not bucket:
|
|
del self._pending[account_key]
|
|
|
|
|
|
class PostgresPairingStore(PairingStore):
|
|
def __init__(self, repo: ChannelPairingRecordRepository | None = None):
|
|
self._repo = repo or ChannelPairingRecordRepository()
|
|
|
|
async def load_pending(self, channel_type: str, account_id: str) -> list[PendingPairingRequest]:
|
|
records = await self._repo.find_pending_by_account(channel_type, account_id)
|
|
result: list[PendingPairingRequest] = []
|
|
for record in records:
|
|
result.append(
|
|
PendingPairingRequest(
|
|
channel_type=record.channel_type,
|
|
peer_id=record.peer_id,
|
|
account_id=record.account_id,
|
|
code=record.pairing_code,
|
|
created_at=record.created_at.timestamp(),
|
|
)
|
|
)
|
|
return result
|
|
|
|
async def upsert_pending(self, req: PendingPairingRequest) -> None:
|
|
expires_at = datetime.now(UTC).replace(tzinfo=None) + timedelta(seconds=PAIRING_PENDING_TTL_SECONDS)
|
|
token = _generate_token()
|
|
record = await self._repo.upsert_pending(
|
|
channel_type=req.channel_type,
|
|
account_id=req.account_id,
|
|
peer_id=req.peer_id,
|
|
code=req.code,
|
|
token=token,
|
|
expires_at=expires_at,
|
|
)
|
|
req._record_id = record.id
|
|
|
|
async def remove_pending(self, channel_type: str, peer_id: str, account_id: str) -> None:
|
|
await self._repo.remove_pending(channel_type, account_id, peer_id, target_status="paired")
|
|
|
|
async def expire_pending(self, channel_type: str, peer_id: str, account_id: str) -> None:
|
|
await self._repo.remove_pending(channel_type, account_id, peer_id, target_status="expired")
|
|
|
|
|
|
class PairingManager:
|
|
def __init__(self, store: PairingStore | None = None, allowlist=None):
|
|
self._store = store or PostgresPairingStore()
|
|
self._allowlist = allowlist
|
|
self._verify_attempts: dict[str, list[float]] = {}
|
|
self._attempts_lock = asyncio.Lock()
|
|
|
|
def _attempts_key(self, channel_type: str, peer_id: str, account_id: str) -> str:
|
|
return f"{channel_type}:{account_id}:{peer_id}"
|
|
|
|
async def _check_rate_limit(self, key: str) -> bool:
|
|
async with self._attempts_lock:
|
|
now = time.monotonic()
|
|
attempts = self._verify_attempts.get(key, [])
|
|
attempts = [t for t in attempts if now - t < 300]
|
|
if len(attempts) >= 10:
|
|
self._verify_attempts[key] = attempts
|
|
return False
|
|
attempts.append(now)
|
|
self._verify_attempts[key] = attempts
|
|
self._cleanup_stale_attempts(now)
|
|
return True
|
|
|
|
def _cleanup_stale_attempts(self, now: float) -> None:
|
|
stale = [k for k, v in self._verify_attempts.items() if not v or all(now - t >= 300 for t in v)]
|
|
for k in stale:
|
|
del self._verify_attempts[k]
|
|
|
|
async def upsert_code(self, channel_type: str, peer_id: str, account_id: str = "default") -> PendingPairingRequest:
|
|
now = time.monotonic()
|
|
existing = await self._store.load_pending(channel_type, account_id)
|
|
|
|
existing_req = next((r for r in existing if r.peer_id == peer_id), None)
|
|
if existing_req is not None:
|
|
return existing_req
|
|
|
|
if len(existing) >= MAX_PENDING_PER_ACCOUNT:
|
|
raise RuntimeError(f"Too many pending pairings for channel '{channel_type}' account '{account_id}'")
|
|
|
|
existing_codes = {r.code for r in existing}
|
|
for _ in range(PAIRING_CODE_MAX_ATTEMPTS):
|
|
code = _generate_code(PAIRING_CODE_ALPHABET, PAIRING_CODE_LENGTH)
|
|
if code not in existing_codes:
|
|
break
|
|
else:
|
|
raise RuntimeError("Failed to generate unique pairing code")
|
|
|
|
req = PendingPairingRequest(
|
|
channel_type=channel_type,
|
|
peer_id=peer_id,
|
|
account_id=account_id,
|
|
code=code,
|
|
created_at=now,
|
|
)
|
|
await self._store.upsert_pending(req)
|
|
return req
|
|
|
|
async def verify(self, channel_type: str, peer_id: str, code: str, account_id: str = "default") -> bool:
|
|
key = self._attempts_key(channel_type, peer_id, account_id)
|
|
if not await self._check_rate_limit(key):
|
|
logger.warning("Pairing verify rate limit exceeded for %s", key)
|
|
return False
|
|
|
|
now = time.monotonic()
|
|
existing = await self._store.load_pending(channel_type, account_id)
|
|
for record in existing:
|
|
if record.peer_id != peer_id:
|
|
continue
|
|
if now - record.created_at > PAIRING_PENDING_TTL_SECONDS:
|
|
await self._store.expire_pending(channel_type, peer_id, account_id)
|
|
return False
|
|
if not hmac.compare_digest(record.code, code.upper()):
|
|
return False
|
|
try:
|
|
await self._store.remove_pending(channel_type, peer_id, account_id)
|
|
except Exception:
|
|
logger.exception("Failed to remove pending pairing after successful verification")
|
|
return True
|
|
return False
|
|
|
|
async def list_pending(self, channel_type: str, account_id: str = "default") -> list[PendingPairingRequest]:
|
|
return await self._store.load_pending(channel_type, account_id)
|
|
|
|
async def approve(self, channel_type: str, peer_id: str, account_id: str = "default") -> None:
|
|
if self._allowlist is not None:
|
|
await self._allowlist.add_entry(channel_type, "dm", peer_id)
|
|
await self._store.remove_pending(channel_type, peer_id, account_id)
|
|
|
|
async def reject(self, channel_type: str, peer_id: str, account_id: str = "default") -> None:
|
|
if self._allowlist is not None:
|
|
await self._allowlist.remove_entry(channel_type, "dm", peer_id)
|
|
await self._store.expire_pending(channel_type, peer_id, account_id)
|
|
|
|
|
|
|