Kris
655eab7230
新增了完整的认证工具库模块，包含以下功能： 1. 指数退避重试组件 2. 敏感数据过滤与日志脱敏 3. 安全策略管理引擎 4. 认证健康监控模块 5. SSRF防护工具集 6. 多类型token管理系统 7. 密钥管理器加解密工具
134 lines
4.3 KiB
Python
134 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
|
|
class DmPolicy(Enum):
|
|
OPEN = "open"
|
|
PAIRING = "pairing"
|
|
ALLOWLIST = "allowlist"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
class GroupPolicy(Enum):
|
|
OPEN = "open"
|
|
ALLOWLIST = "allowlist"
|
|
DISABLED = "disabled"
|
|
|
|
|
|
class SecurityPolicy:
|
|
def __init__(
|
|
self,
|
|
dm_policy: DmPolicy = DmPolicy.OPEN,
|
|
group_policy: GroupPolicy = GroupPolicy.OPEN,
|
|
allowlist: set[str] | None = None,
|
|
blocklist: set[str] | None = None,
|
|
):
|
|
self.dm_policy = dm_policy
|
|
self.group_policy = group_policy
|
|
self._allowlist: set[str] = allowlist or set()
|
|
self._blocklist: set[str] = blocklist or set()
|
|
|
|
@classmethod
|
|
def from_config(cls, config: dict[str, Any]) -> SecurityPolicy:
|
|
dm_raw = config.get("dm_policy", config.get("dmPolicy", "open"))
|
|
group_raw = config.get("group_policy", config.get("groupPolicy", "open"))
|
|
|
|
dm_map = {
|
|
"open": DmPolicy.OPEN,
|
|
"pairing": DmPolicy.PAIRING,
|
|
"allowlist": DmPolicy.ALLOWLIST,
|
|
"disabled": DmPolicy.DISABLED,
|
|
}
|
|
group_map = {
|
|
"open": GroupPolicy.OPEN,
|
|
"allowlist": GroupPolicy.ALLOWLIST,
|
|
"disabled": GroupPolicy.DISABLED,
|
|
}
|
|
|
|
return cls(
|
|
dm_policy=dm_map.get(dm_raw, DmPolicy.OPEN),
|
|
group_policy=group_map.get(group_raw, GroupPolicy.OPEN),
|
|
allowlist=set(config.get("allowlist", []) or []),
|
|
blocklist=set(config.get("blocklist", []) or []),
|
|
)
|
|
|
|
def check_dm(self, sender_id: str) -> bool:
|
|
if self.dm_policy == DmPolicy.OPEN:
|
|
return True
|
|
if self.dm_policy == DmPolicy.DISABLED:
|
|
return False
|
|
if self.dm_policy == DmPolicy.ALLOWLIST:
|
|
return sender_id in self._allowlist
|
|
if self.dm_policy == DmPolicy.PAIRING:
|
|
return True
|
|
return False
|
|
|
|
def check_group(self, sender_id: str | None = None) -> bool:
|
|
if self.group_policy == GroupPolicy.OPEN:
|
|
return True
|
|
if self.group_policy == GroupPolicy.DISABLED:
|
|
return False
|
|
if self.group_policy == GroupPolicy.ALLOWLIST:
|
|
return sender_id is None or sender_id in self._allowlist
|
|
return False
|
|
|
|
def is_allowlisted(self, sender_id: str) -> bool:
|
|
return sender_id in self._allowlist
|
|
|
|
def is_blocklisted(self, sender_id: str) -> bool:
|
|
return sender_id in self._blocklist
|
|
|
|
def add_to_allowlist(self, sender_id: str) -> None:
|
|
self._allowlist.add(sender_id)
|
|
|
|
def remove_from_allowlist(self, sender_id: str) -> None:
|
|
self._allowlist.discard(sender_id)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {
|
|
"dm_policy": self.dm_policy.value,
|
|
"group_policy": self.group_policy.value,
|
|
"allowlist": list(self._allowlist),
|
|
"blocklist": list(self._blocklist),
|
|
}
|
|
|
|
|
|
class SecurityPolicyEngine:
|
|
def __init__(self):
|
|
self._policies: dict[str, SecurityPolicy] = {}
|
|
|
|
def register_policy(self, channel_id: str, policy: SecurityPolicy) -> None:
|
|
self._policies[channel_id] = policy
|
|
|
|
def get_policy(self, channel_id: str) -> SecurityPolicy:
|
|
if channel_id not in self._policies:
|
|
self._policies[channel_id] = SecurityPolicy()
|
|
return self._policies[channel_id]
|
|
|
|
def resolve_dm_policy(self, channel_id: str) -> DmPolicy:
|
|
return self.get_policy(channel_id).dm_policy
|
|
|
|
def resolve_group_policy(self, channel_id: str) -> GroupPolicy:
|
|
return self.get_policy(channel_id).group_policy
|
|
|
|
def check_allowlist(self, channel_id: str, sender_id: str) -> bool:
|
|
return self.get_policy(channel_id).is_allowlisted(sender_id)
|
|
|
|
def check_dm(self, channel_id: str, sender_id: str) -> bool:
|
|
return self.get_policy(channel_id).check_dm(sender_id)
|
|
|
|
def check_group(self, channel_id: str, sender_id: str | None = None) -> bool:
|
|
return self.get_policy(channel_id).check_group(sender_id)
|
|
|
|
|
|
_security_policy_engine: SecurityPolicyEngine | None = None
|
|
|
|
|
|
def get_security_policy_engine() -> SecurityPolicyEngine:
|
|
global _security_policy_engine
|
|
if _security_policy_engine is None:
|
|
_security_policy_engine = SecurityPolicyEngine()
|
|
return _security_policy_engine
|