Kris
849d21120f
新增一系列安全相关功能： 1. 新增二维码生成工具，支持自定义参数导出图片/Base64/字节流 2. 新增HTML内容安全处理工具，包括标签剥离、转义、URL校验 3. 新增日志敏感信息脱敏工具，支持配置、字典、日志记录脱敏 4. 新增密钥管理运行时，支持从环境变量/文件加载密钥 5. 新增身份链接管理，支持多渠道身份绑定与解析 6. 新增SSRF防护工具，支持域名/IP校验与固定主机 7. 新增安全权限修复工具，修复文件目录权限与配置项 8. 新增外部内容安全处理，支持LLM特殊令牌剥离与注入检测 9. 新增认证限流工具，支持多维度限流与本地回环豁免 10. 新增配对管理工具，支持安全配对码生成与校验 11. 新增白名单管理工具，支持DM/群组/来源白名单校验
334 lines
13 KiB
Python
334 lines
13 KiB
Python
import asyncio
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from enum import StrEnum
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AllowlistHitReason(StrEnum):
|
|
DM_WILDCARD = "dm_wildcard"
|
|
DM_EXACT = "dm_exact"
|
|
GROUP_WILDCARD = "group_wildcard"
|
|
GROUP_EXACT = "group_exact"
|
|
GROUP_ALLOW_FROM_EXACT = "group_allow_from_exact"
|
|
REJECTED = "rejected"
|
|
DM_OPEN = "dm_open"
|
|
DM_DISABLED = "dm_disabled"
|
|
GROUP_OPEN = "group_open"
|
|
GROUP_DISABLED = "group_disabled"
|
|
ALLOW_FROM_EMPTY = "allow_from_empty"
|
|
|
|
|
|
class DmPolicyMode(StrEnum):
|
|
OPEN = "open"
|
|
PAIRING = "pairing"
|
|
ALLOWLIST = "allowlist"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
class GroupPolicyMode(StrEnum):
|
|
OPEN = "open"
|
|
ALLOWLIST = "allowlist"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
class AllowlistCheckResult:
|
|
def __init__(
|
|
self,
|
|
allowed: bool,
|
|
reason: AllowlistHitReason = AllowlistHitReason.REJECTED,
|
|
matched_entry: str | None = None,
|
|
pairing_required: bool = False,
|
|
):
|
|
self.allowed = allowed
|
|
self.reason = reason
|
|
self.matched_entry = matched_entry
|
|
self.pairing_required = pairing_required
|
|
|
|
|
|
# ── AllowlistStore (abstract) ─────────────────────────────────────────
|
|
|
|
|
|
class AllowlistStore(ABC):
|
|
@abstractmethod
|
|
async def load_all(self, channel_id: str) -> dict[str, list[str]]: ...
|
|
|
|
@abstractmethod
|
|
async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: ...
|
|
|
|
@abstractmethod
|
|
async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None: ...
|
|
|
|
|
|
# ── InMemoryAllowlistStore ────────────────────────────────────────────
|
|
|
|
|
|
class InMemoryAllowlistStore(AllowlistStore):
|
|
def __init__(self):
|
|
self._entries: dict[str, dict[str, set[str]]] = {}
|
|
|
|
async def load_all(self, channel_id: str) -> dict[str, list[str]]:
|
|
channel_entries = self._entries.get(channel_id, {})
|
|
return {allow_type: list(ids) for allow_type, ids in channel_entries.items()}
|
|
|
|
async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None:
|
|
if channel_id not in self._entries:
|
|
self._entries[channel_id] = {}
|
|
if allow_type not in self._entries[channel_id]:
|
|
self._entries[channel_id][allow_type] = set()
|
|
self._entries[channel_id][allow_type].add(allow_id)
|
|
|
|
async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None:
|
|
if channel_id in self._entries and allow_type in self._entries[channel_id]:
|
|
self._entries[channel_id][allow_type].discard(allow_id)
|
|
|
|
|
|
# ── PostgresAllowlistStore ────────────────────────────────────────────
|
|
|
|
|
|
class PostgresAllowlistStore(AllowlistStore):
|
|
def __init__(self):
|
|
from yuxi.repositories.channel_allowlist_repo import ChannelAllowlistRepository
|
|
|
|
self._repo = ChannelAllowlistRepository()
|
|
|
|
async def load_all(self, channel_id: str) -> dict[str, list[str]]:
|
|
entries = await self._repo.list_by_channel(channel_id)
|
|
result: dict[str, list[str]] = {}
|
|
for entry in entries:
|
|
result.setdefault(entry.allow_type, []).append(entry.allow_id)
|
|
return result
|
|
|
|
async def add_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None:
|
|
await self._repo.add_entry(channel_id, allow_type, allow_id)
|
|
|
|
async def remove_entry(self, channel_id: str, allow_type: str, allow_id: str) -> None:
|
|
await self._repo.remove_entry(channel_id, allow_type, allow_id)
|
|
|
|
|
|
# ── AllowlistChecker ──────────────────────────────────────────────────
|
|
|
|
|
|
class AllowlistChecker:
|
|
def __init__(self, store: AllowlistStore | None = None):
|
|
self._store = store or InMemoryAllowlistStore()
|
|
self._dm_allowlist: list[str] = []
|
|
self._group_allowlist: list[str] = []
|
|
self._group_allow_from: list[str] = []
|
|
self._dm_policy: DmPolicyMode = DmPolicyMode.ALLOWLIST
|
|
self._group_policy: GroupPolicyMode = GroupPolicyMode.ALLOWLIST
|
|
self._loaded = False
|
|
self._loading = False
|
|
self._channel_id: str | None = None
|
|
self._load_lock = asyncio.Lock()
|
|
|
|
def load_dm_allowlist(self, entries: list[str]) -> None:
|
|
self._dm_allowlist = list(entries)
|
|
self._loaded = True
|
|
|
|
def load_group_allowlist(self, group_ids: list[str]) -> None:
|
|
self._group_allowlist = list(group_ids)
|
|
self._loaded = True
|
|
|
|
def load_group_allow_from(self, entries: list[str]) -> None:
|
|
self._group_allow_from = list(entries)
|
|
|
|
async def ensure_loaded(self, channel_id: str) -> None:
|
|
if self._loaded:
|
|
return
|
|
async with self._load_lock:
|
|
if self._loaded:
|
|
return
|
|
self._loading = True
|
|
try:
|
|
data = await self._store.load_all(channel_id)
|
|
self._dm_allowlist = data.get("user", [])
|
|
self._group_allowlist = data.get("group", [])
|
|
self._group_allow_from = data.get("allow_from", [])
|
|
self._loaded = True
|
|
self._channel_id = channel_id
|
|
except Exception:
|
|
logger.exception("Failed to load allowlist for channel %s", channel_id)
|
|
raise
|
|
finally:
|
|
self._loading = False
|
|
|
|
@property
|
|
def dm_policy(self) -> DmPolicyMode:
|
|
return self._dm_policy
|
|
|
|
@dm_policy.setter
|
|
def dm_policy(self, mode: DmPolicyMode) -> None:
|
|
self._dm_policy = mode
|
|
|
|
@property
|
|
def group_policy(self) -> GroupPolicyMode:
|
|
return self._group_policy
|
|
|
|
@group_policy.setter
|
|
def group_policy(self, mode: GroupPolicyMode) -> None:
|
|
self._group_policy = mode
|
|
|
|
def check_dm(self, peer_id: str) -> AllowlistCheckResult:
|
|
if self._dm_policy == DmPolicyMode.OPEN:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_OPEN)
|
|
|
|
if self._dm_policy == DmPolicyMode.DISABLED:
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.DM_DISABLED)
|
|
|
|
if self._dm_policy == DmPolicyMode.PAIRING:
|
|
if self._dm_allowlist and "*" in self._dm_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_WILDCARD, matched_entry="*")
|
|
if peer_id in self._dm_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_EXACT, matched_entry=peer_id)
|
|
return AllowlistCheckResult(
|
|
allowed=False,
|
|
reason=AllowlistHitReason.REJECTED,
|
|
pairing_required=True,
|
|
)
|
|
|
|
if not self._dm_allowlist:
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED)
|
|
if "*" in self._dm_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_WILDCARD, matched_entry="*")
|
|
if peer_id in self._dm_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.DM_EXACT, matched_entry=peer_id)
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED)
|
|
|
|
def check_group(self, group_id: str) -> AllowlistCheckResult:
|
|
if self._group_policy == GroupPolicyMode.OPEN:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_OPEN)
|
|
|
|
if self._group_policy == GroupPolicyMode.DISABLED:
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.GROUP_DISABLED)
|
|
|
|
if not self._group_allowlist:
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED)
|
|
if "*" in self._group_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_WILDCARD, matched_entry="*")
|
|
if group_id in self._group_allowlist:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_EXACT, matched_entry=group_id)
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED)
|
|
|
|
def check_group_allow_from(self, sender_id: str) -> AllowlistCheckResult:
|
|
if self._group_policy == GroupPolicyMode.OPEN:
|
|
return AllowlistCheckResult(allowed=True, reason=AllowlistHitReason.GROUP_OPEN)
|
|
|
|
if self._group_policy == GroupPolicyMode.DISABLED:
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.GROUP_DISABLED)
|
|
|
|
if not self._group_allow_from:
|
|
logger.warning(
|
|
"group_allow_from is empty under allowlist policy for channel %s — blocking all group senders. "
|
|
"Add explicit entries or use groupPolicy=open to allow.",
|
|
self._channel_id,
|
|
)
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.ALLOW_FROM_EMPTY)
|
|
if "*" in self._group_allow_from:
|
|
return AllowlistCheckResult(
|
|
allowed=True,
|
|
reason=AllowlistHitReason.GROUP_ALLOW_FROM_EXACT,
|
|
matched_entry="*",
|
|
)
|
|
if sender_id in self._group_allow_from:
|
|
return AllowlistCheckResult(
|
|
allowed=True,
|
|
reason=AllowlistHitReason.GROUP_ALLOW_FROM_EXACT,
|
|
matched_entry=sender_id,
|
|
)
|
|
return AllowlistCheckResult(allowed=False, reason=AllowlistHitReason.REJECTED)
|
|
|
|
def add_to_dm_allowlist(self, peer_id: str) -> None:
|
|
if peer_id not in self._dm_allowlist:
|
|
self._dm_allowlist.append(peer_id)
|
|
|
|
def remove_from_dm_allowlist(self, peer_id: str) -> None:
|
|
if peer_id in self._dm_allowlist:
|
|
self._dm_allowlist.remove(peer_id)
|
|
|
|
_VALID_SCOPES = frozenset({"dm", "group", "allow_from"})
|
|
|
|
async def add_entry(self, channel_type: str, scope: str, value: str) -> bool:
|
|
if scope not in self._VALID_SCOPES:
|
|
raise ValueError(f"Invalid scope: {scope!r}, must be one of {sorted(self._VALID_SCOPES)}")
|
|
allow_type = {"dm": "user", "group": "group", "allow_from": "allow_from"}[scope]
|
|
await self._store.add_entry(channel_type, allow_type, value)
|
|
if scope == "dm":
|
|
if value not in self._dm_allowlist:
|
|
self._dm_allowlist.append(value)
|
|
return True
|
|
if scope == "group":
|
|
if value not in self._group_allowlist:
|
|
self._group_allowlist.append(value)
|
|
return True
|
|
if scope == "allow_from":
|
|
if value not in self._group_allow_from:
|
|
self._group_allow_from.append(value)
|
|
return True
|
|
return False
|
|
|
|
async def remove_entry(self, channel_type: str, scope: str, value: str) -> bool:
|
|
if scope not in self._VALID_SCOPES:
|
|
raise ValueError(f"Invalid scope: {scope!r}, must be one of {sorted(self._VALID_SCOPES)}")
|
|
allow_type = {"dm": "user", "group": "group", "allow_from": "allow_from"}[scope]
|
|
await self._store.remove_entry(channel_type, allow_type, value)
|
|
if scope == "dm":
|
|
if value in self._dm_allowlist:
|
|
self._dm_allowlist.remove(value)
|
|
return True
|
|
if scope == "group":
|
|
if value in self._group_allowlist:
|
|
self._group_allowlist.remove(value)
|
|
return True
|
|
if scope == "allow_from":
|
|
if value in self._group_allow_from:
|
|
self._group_allow_from.remove(value)
|
|
return True
|
|
return False
|
|
|
|
def list_entries(self, channel_type: str, scope: str) -> list[str]:
|
|
if scope == "dm":
|
|
return list(self._dm_allowlist)
|
|
if scope == "group":
|
|
return list(self._group_allowlist)
|
|
if scope == "allow_from":
|
|
return list(self._group_allow_from)
|
|
return []
|
|
|
|
|
|
# ── Singleton ─────────────────────────────────────────────────────────
|
|
|
|
import contextvars
|
|
|
|
_allowlist_ctx: contextvars.ContextVar[AllowlistChecker | None] = contextvars.ContextVar(
|
|
"allowlist_checker", default=None
|
|
)
|
|
_allowlist_lock = asyncio.Lock()
|
|
|
|
|
|
def create_allowlist_checker() -> AllowlistChecker:
|
|
return AllowlistChecker(PostgresAllowlistStore())
|
|
|
|
|
|
async def get_allowlist_checker() -> AllowlistChecker:
|
|
checker = _allowlist_ctx.get()
|
|
if checker is None:
|
|
async with _allowlist_lock:
|
|
if _allowlist_ctx.get() is None:
|
|
_allowlist_ctx.set(create_allowlist_checker())
|
|
checker = _allowlist_ctx.get()
|
|
if checker is None:
|
|
checker = create_allowlist_checker()
|
|
_allowlist_ctx.set(checker)
|
|
return checker
|
|
|
|
|
|
def set_allowlist_checker(checker: AllowlistChecker) -> None:
|
|
_allowlist_ctx.set(checker)
|
|
|
|
|
|
async def reset_allowlist_checker() -> None:
|
|
async with _allowlist_lock:
|
|
_allowlist_ctx.set(None)
|