pbf.utils.Register

  1import re
  2import time
  3from enum import unique, Enum
  4from functools import wraps
  5from typing import List
  6
  7from ..setup import logger, ListenerManager, pluginsManager
  8from ..error import LimitExceedError
  9
 10
 11def allPermission(*_):  # 默认权限，需要接收两个实参
 12    return True
 13
 14def adminPermission(event):  # 在装饰器中指定的权限，只需要接受`event`参数
 15    role: str = event.sender.get("role", "member")
 16    if role == "admin" or role == "owner":
 17        return True
 18    return False
 19
 20def ownerPermission(event):
 21    if event.sender.get("role", "member") == "owner":
 22        return True
 23    return False
 24
 25def default_callback(event, func):
 26    raise LimitExceedError(f"User {event.user_id} exceeded the limit of {func.__name__}.")
 27
 28
 29class Limit:
 30    def __init__(self, duration, times, user=None, callback=None):
 31        """
 32        限制指令执行频率
 33        :param duration: int 限制时间
 34        :param times: int 限制次数
 35        :param user: str/int/None (可选)用户ID
 36        :param callback: Callable (可选)回调函数
 37        """
 38        self.duration = duration
 39        self.times = times
 40        self.user = user
 41        self.call_count = {}
 42        self.callback = callback or default_callback
 43
 44    def __call__(self, origin_func):
 45        @wraps(origin_func)
 46        def limit_wrapper(*args, **kwargs):
 47            # 假设第一个参数是event，且event有user_id属性
 48            event = args[0] if args else None
 49            user_id = event.user_id if event else None
 50            current_time = time.time()
 51            if self.user is not None:
 52                user_id = self.user
 53
 54            if user_id not in self.call_count:
 55                self.call_count[user_id] = []
 56
 57            # 移除过期的调用记录
 58            self.call_count[user_id] = [
 59                timestamp for timestamp in self.call_count[user_id]
 60                if current_time - timestamp < self.duration
 61            ]
 62
 63            # 检查调用次数
 64            if len(self.call_count[user_id]) < self.times:
 65                self.call_count[user_id].append(current_time)
 66                return origin_func(*args, **kwargs)
 67            else:
 68                # logger.warning(f"User {user_id} exceeded the limit of {origin_func.__name__}.")
 69                return self.callback(event, origin_func)  # 调用回调函数
 70
 71        return limit_wrapper
 72
 73
 74@unique
 75class RegexMode(Enum):
 76    match = "match"
 77    search = "search"
 78    none = "none"
 79
 80
 81class Base:
 82    name: str = 'Command name'
 83    description: str = 'Command description'
 84    permission: callable = allPermission
 85    usage: str = 'Command usage'
 86    alias: List[str] = []
 87    hidden: bool = False
 88    enabled: bool = True
 89    type: str = 'command'
 90    func: callable = None
 91    regex_mode: RegexMode = RegexMode.none
 92    regex_flags = re.I
 93    regex_data = None
 94    origin_func: callable = None
 95
 96    def __init__(self, **kwargs):
 97        """
 98        Initialize the command information.
 99        :param kwargs: **dict
100        """
101        for key, value in kwargs.items():
102            setattr(self, key, value)
103
104    def __call__(self, origin_func):
105        if self.enabled:
106            logger.debug(f"Registering {self.type} `{self.name}`")
107            ListenerManager.add_listener(self.type, pluginsManager.plugin_name, self)
108
109        def try_except_wrapper(*args, **kwargs):
110            try:
111                return origin_func(*args, **kwargs)
112            except Exception as e:
113                logger.error(f"Error in {self.type} `{self.name}`: {e}")
114
115        func = try_except_wrapper
116        self.origin_func = origin_func
117        self.func = func
118
119        @wraps(func)
120        def wrapper(*args, **kwargs):
121            logger.debug(f"<{self.type}> {self.name} called")
122            return func(*args, **kwargs)
123
124        return wrapper
125
126    def __str__(self):
127        return f"<{self.type}Listener \"{self.name}\" call:{self.origin_func.__name__}>"
128
129
130class Command(Base):
131    def __init__(self, **kwargs):
132        super().__init__(**kwargs)
133        self.type = "command"
134
135
136class Message(Base):
137    def __init__(self, **kwargs):
138        super().__init__(**kwargs)
139        self.type = "message"
140
141
142class Notice(Base):
143    def __init__(self, **kwargs):
144        super().__init__(**kwargs)
145        self.type = "notice"
146
147
148class Request(Base):
149    def __init__(self, **kwargs):
150        super().__init__(**kwargs)
151        self.type = "request"
152
153
154class Meta(Base):
155    def __init__(self, **kwargs):
156        super().__init__(**kwargs)
157        self.type = "meta"