add sandbox filter

qg_botsdk/__init__.py

@@ -34,6 +34,7 @@
 from .proto import Proto
 from .proto.wh_backend import SigningKey
 from .qg_bot import BOT
+from .sandbox import SandBox
 from .version import __version__
 
 json.JSONEncoder.default = lambda self, obj: (
@@ -68,4 +69,5 @@
     "IdTokenError",
     "WaitTimeoutError",
     "Proto",
+    "SandBox",
 )

qg_botsdk/qg_bot.py

@@ -24,6 +24,7 @@
 from .plugins import Plugins
 from .proto import Proto
 from .qg_bot_proto import BotProto as _BotWs
+from .sandbox import SandBox
 from .session import AbstractSessionManager, SessionPatcher
 from .version import __version__
 
@@ -48,6 +49,7 @@ def __init__(
         api_max_concurrency: int = 0,
         api_timeout: int = 20,
         protocol: Proto = Proto.websocket(),
+        sandbox: Optional[SandBox] = None,
     ):
         """
         机器人主体，输入BotAppID和密钥，并绑定函数后即可快速使用
@@ -65,6 +67,7 @@ def __init__(
         :param api_max_concurrency: API允许的最大并发数，超过此并发数将进入队列，如此数值<=0代表不开启任何队列，默认0
         :param api_timeout: API请求的超时设置。默认20
         :param protocol: 机器人连接协议，默认为Proto.websocket()
+        :param sandbox: 沙箱模式配置项，当 is_sandbox=True 时，只有指定的频道、群、用户可以接收到消息；否则当非沙箱环境时，过滤掉指定频道、群、用户的消息。
         """
         try:
             self._loop = get_event_loop()
@@ -145,6 +148,13 @@ def __init__(
         ] = {1 << x: [] for x in range(CommandValidScenes.ALL.bit_length())}
         self.session: AbstractSessionManager = SessionPatcher()
         self.protocol = protocol
+        self.sandbox = sandbox
+        if isinstance(self.sandbox, SandBox):
+            self.sandbox.set_logger(self.logger)
+            self.sandbox.set_is_sandbox(is_sandbox)
+            self.logger.info("已加载沙箱配置项")
+        elif self.sandbox:
+            raise TypeError("传入的沙箱配置项不是SandBox类")
 
     def __repr__(self):
         return f"<qg_botsdk.BOT object [id: {self.bot_id}, token: {self.bot_token}]>"
@@ -804,6 +814,7 @@ def start(self, is_blocking: bool = True):
                     self._preprocessors,
                     self.__session_manager,
                     self.protocol,
+                    self.sandbox,
                 )
                 self.__task = self._loop.create_task(self._bot_class.start())
                 if is_blocking and not self._loop.is_running():

qg_botsdk/qg_bot_proto.py

@@ -17,6 +17,7 @@
 from .logger import Logger
 from .model import BotCommandObject, CommandValidScenes, Model
 from .proto import proto
+from .sandbox import SandBox
 
 
 class BotProto:
@@ -49,6 +50,7 @@ def __init__(
         ],
         session_manager: SessionManager,
         protocol: proto.Proto,
+        sandbox: SandBox,
     ):
         """
         此为SDK内部使用类，注册机器人请使用from qg_botsdk.qg_bot import BOT
@@ -131,6 +133,7 @@ def __init__(
             loop=loop,
             dispatch_func=self.dispatch_events,
         )
+        self.sandbox = sandbox
         self.seq_cache = SeqCache()
 
     @exception_processor
@@ -368,21 +371,35 @@ async def data_process(self, data: Dict):
                 treated_msg = ""
             # distribute_commands return True when short circuit
             if t in EVENTS.MESSAGE_CREATE:
+                if self.sandbox and not self.sandbox.checker(
+                    EVENTS.MESSAGE_CREATE, data
+                ):
+                    return
                 if not await self.distribute_commands(
                     CommandValidScenes.GUILD, data, treated_msg
                 ):
                     await self.distribute(self.func_registers["on_msg"], data)
             elif t in EVENTS.DM_CREATE:
+                if self.sandbox and not self.sandbox.checker(EVENTS.DM_CREATE, data):
+                    return
                 if not await self.distribute_commands(
                     CommandValidScenes.DM, data, treated_msg
                 ):
                     await self.distribute(self.func_registers["on_dm"], data)
             elif t in EVENTS.C2C_MESSAGE_CREATE:
+                if self.sandbox and not self.sandbox.checker(
+                    EVENTS.C2C_MESSAGE_CREATE, data
+                ):
+                    return
                 if not await self.distribute_commands(
                     CommandValidScenes.C2C, data, treated_msg
                 ):
                     await self.distribute(self.func_registers["on_friend_msg"], data)
             else:  # t in EVENTS.GROUP_AT_MESSAGE_CREATE
+                if self.sandbox and not self.sandbox.checker(
+                    EVENTS.GROUP_AT_MESSAGE_CREATE, data
+                ):
+                    return
                 if not await self.distribute_commands(
                     CommandValidScenes.GROUP, data, treated_msg
                 ):

qg_botsdk/sandbox.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+from typing import Any, List, Optional, Iterable
+
+from ._statics import EVENTS
+from ._utils import exception_handler
+from .logger import Logger
+
+
+class SandBox:
+    def __init__(
+        self,
+        guilds: Optional[List[str]] = None,
+        guild_users: Optional[List[str]] = None,
+        groups: Optional[List[str]] = None,
+        q_users: Optional[List[str]] = None,
+    ):
+        """ "
+        沙箱模式配置项，当BOT(..., is_sandbox=True)时，只有列表中指定的频道、群、用户可以接收到消息；
+        当为False（非沙箱环境）时，过滤掉列表中指定频道、群、用户的消息。
+
+        :param guilds: 设置为沙箱的频道ID列表
+        :param guild_users: 设置为沙箱的频道私信用户ID列表
+        :param groups: 设置为沙箱的群ID列表
+        :param q_users: 设置为沙箱的QQ私信用户ID列表
+        """
+        self.guilds = self.__to_set(guilds)
+        self.guild_users = self.__to_set(guild_users)
+        self.groups = self.__to_set(groups)
+        self.q_users = self.__to_set(q_users)
+
+        self.is_sandbox = False
+        self.logger: Optional[Logger] = None
+
+    @staticmethod
+    def __to_set(data: Optional[List[str]]) -> set:
+        if isinstance(data, Iterable):
+            return set(data)
+        return set()
+
+    def set_logger(self, logger):
+        self.logger = logger
+
+    def set_is_sandbox(self, is_sandbox: bool):
+        self.is_sandbox = is_sandbox
+
+    def checker(self, event: EVENTS, data: dict) -> bool:
+        try:
+            if event == EVENTS.MESSAGE_CREATE:
+                return (data["d"]["guild_id"] not in self.guilds) ^ self.is_sandbox
+            elif event == EVENTS.DM_CREATE:
+                return (data["d"]["author"]["id"] not in self.guild_users) ^ self.is_sandbox
+            elif event == EVENTS.GROUP_AT_MESSAGE_CREATE:
+                return (data["d"]["group_openid"] not in self.groups) ^ self.is_sandbox
+            elif event == EVENTS.C2C_MESSAGE_CREATE:
+                return (data["d"]["author"]["id"] not in self.q_users) ^ self.is_sandbox
+        except Exception as e:
+            if self.logger:
+                self.logger.error(f"沙箱模式检查失败，已放行：{repr(e)}")
+                self.logger.debug(f"事件：{event}, 数据：{data}")
+                self.logger.error(exception_handler(e))
+            return True

qg_botsdk/version.py

@@ -1 +1 @@
-__version__ = "4.3.0"
+__version__ = "4.3.1"

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = qg-botsdk
-version = 4.3.0
+version = 4.3.1
 author = GDLY
 author_email = [email protected]
 description = easy-to-use SDK for Tencent QQ guild robot