129 lines
3.5 KiB
Python
129 lines
3.5 KiB
Python
"""
|
||
消息过滤器模块
|
||
|
||
提供消息过滤功能：
|
||
- 白名单/黑名单过滤
|
||
- 机器人自身消息过滤
|
||
- 系统消息放行
|
||
"""
|
||
|
||
from typing import Any, Dict, List, Optional, Set
|
||
|
||
from loguru import logger
|
||
|
||
|
||
class MessageFilter:
|
||
"""
|
||
消息过滤器
|
||
|
||
支持三种模式:
|
||
- None: 不过滤，处理所有消息
|
||
- Whitelist: 只处理白名单中的消息
|
||
- Blacklist: 过滤黑名单中的消息
|
||
"""
|
||
|
||
# 系统消息类型（始终放行）
|
||
SYSTEM_MESSAGE_TYPES = {11058}
|
||
|
||
def __init__(
|
||
self,
|
||
mode: str = "None",
|
||
whitelist: List[str] = None,
|
||
blacklist: List[str] = None,
|
||
bot_wxid: str = None,
|
||
):
|
||
"""
|
||
初始化过滤器
|
||
|
||
Args:
|
||
mode: 过滤模式 ("None", "Whitelist", "Blacklist")
|
||
whitelist: 白名单 wxid 列表
|
||
blacklist: 黑名单 wxid 列表
|
||
bot_wxid: 机器人自身 wxid（用于过滤自己的消息）
|
||
"""
|
||
self.mode = mode
|
||
self.whitelist: Set[str] = set(whitelist or [])
|
||
self.blacklist: Set[str] = set(blacklist or [])
|
||
self.bot_wxid = bot_wxid
|
||
|
||
def set_bot_wxid(self, wxid: str):
|
||
"""设置机器人 wxid"""
|
||
self.bot_wxid = wxid
|
||
|
||
def add_to_whitelist(self, wxid: str):
|
||
"""添加到白名单"""
|
||
self.whitelist.add(wxid)
|
||
|
||
def remove_from_whitelist(self, wxid: str):
|
||
"""从白名单移除"""
|
||
self.whitelist.discard(wxid)
|
||
|
||
def add_to_blacklist(self, wxid: str):
|
||
"""添加到黑名单"""
|
||
self.blacklist.add(wxid)
|
||
|
||
def remove_from_blacklist(self, wxid: str):
|
||
"""从黑名单移除"""
|
||
self.blacklist.discard(wxid)
|
||
|
||
def should_process(self, message: Dict[str, Any]) -> bool:
|
||
"""
|
||
判断消息是否应该被处理
|
||
|
||
Args:
|
||
message: 标准化后的消息字典
|
||
|
||
Returns:
|
||
True 表示应该处理，False 表示应该过滤
|
||
"""
|
||
from_wxid = message.get("FromWxid", "")
|
||
sender_wxid = message.get("SenderWxid", "")
|
||
msg_type = message.get("MsgType", 0)
|
||
|
||
# 系统消息始终放行
|
||
if msg_type in self.SYSTEM_MESSAGE_TYPES:
|
||
return True
|
||
|
||
# 过滤机器人自己的消息
|
||
if self.bot_wxid and (from_wxid == self.bot_wxid or sender_wxid == self.bot_wxid):
|
||
return False
|
||
|
||
# 根据模式过滤
|
||
return self._check_mode(from_wxid, sender_wxid)
|
||
|
||
def _check_mode(self, from_wxid: str, sender_wxid: str) -> bool:
|
||
"""根据模式检查是否放行"""
|
||
if self.mode == "None":
|
||
return True
|
||
|
||
if self.mode == "Whitelist":
|
||
return from_wxid in self.whitelist or sender_wxid in self.whitelist
|
||
|
||
if self.mode == "Blacklist":
|
||
return from_wxid not in self.blacklist and sender_wxid not in self.blacklist
|
||
|
||
return True
|
||
|
||
@classmethod
|
||
def from_config(cls, bot_config: Dict[str, Any]) -> "MessageFilter":
|
||
"""
|
||
从配置创建过滤器
|
||
|
||
Args:
|
||
bot_config: Bot 配置节
|
||
|
||
Returns:
|
||
MessageFilter 实例
|
||
"""
|
||
return cls(
|
||
mode=bot_config.get("ignore-mode", "None"),
|
||
whitelist=bot_config.get("whitelist", []),
|
||
blacklist=bot_config.get("blacklist", []),
|
||
bot_wxid=bot_config.get("wxid") or bot_config.get("bot_wxid"),
|
||
)
|
||
|
||
|
||
# ==================== 导出 ====================
|
||
|
||
__all__ = ['MessageFilter']
|