Merge branch 'main' of https://gitea.functen.cn/shihao/WechatHookBot

utils/message_filter.py

@@ -0,0 +1,128 @@
+"""
+消息过滤器模块
+
+提供消息过滤功能：
+- 白名单/黑名单过滤
+- 机器人自身消息过滤
+- 系统消息放行
+"""
+
+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']