43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
|
from enum import StrEnum
|
||
|
|
|
||
|
|
from yuxi.channels.models import ChannelMessage, ChatType
|
||
|
|
|
||
|
|
|
||
|
|
class GroupChatMode(StrEnum):
|
||
|
|
MENTION_ONLY = "mention_only"
|
||
|
|
ALL = "all"
|
||
|
|
WHITELIST = "whitelist"
|
||
|
|
BLACKLIST = "blacklist"
|
||
|
|
|
||
|
|
|
||
|
|
class GroupChatPolicy:
|
||
|
|
DEFAULT_MODE = GroupChatMode.MENTION_ONLY
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._whitelist: set[str] = set()
|
||
|
|
self._blacklist: set[str] = set()
|
||
|
|
self._mode = self.DEFAULT_MODE
|
||
|
|
|
||
|
|
def configure(self, mode: GroupChatMode, whitelist=None, blacklist=None) -> None:
|
||
|
|
self._mode = mode
|
||
|
|
if whitelist is not None:
|
||
|
|
self._whitelist = set(whitelist)
|
||
|
|
if blacklist is not None:
|
||
|
|
self._blacklist = set(blacklist)
|
||
|
|
|
||
|
|
def should_respond(self, message: ChannelMessage, is_at_bot: bool) -> bool:
|
||
|
|
if message.chat_type == ChatType.DIRECT:
|
||
|
|
return True
|
||
|
|
|
||
|
|
chat_id = message.identity.channel_chat_id
|
||
|
|
|
||
|
|
if self._mode == GroupChatMode.MENTION_ONLY:
|
||
|
|
return is_at_bot
|
||
|
|
if self._mode == GroupChatMode.ALL:
|
||
|
|
return True
|
||
|
|
if self._mode == GroupChatMode.WHITELIST:
|
||
|
|
return chat_id in self._whitelist
|
||
|
|
if self._mode == GroupChatMode.BLACKLIST:
|
||
|
|
return chat_id not in self._blacklist
|
||
|
|
return True
|