import asyncio import os import random import re from collections import defaultdict from functools import wraps import nonebot import pytz from nonebot.command import SwitchException, _FinishException, _PauseException from nonebot.message import CanceledException import hoshino from hoshino import log, priv, trigger from hoshino.typing import * try: import ujson as json except: import json # service management _loaded_services: Dict[str, "Service"] = {} # {name: service} _service_bundle: Dict[str, List["Service"]] = defaultdict(list) _re_illegal_char = re.compile(r'[\\/:*?"<>|\.]') _service_config_dir = os.path.expanduser('~/.hoshino/service_config/') os.makedirs(_service_config_dir, exist_ok=True) def _load_service_config(service_name): config_file = os.path.join(_service_config_dir, f'{service_name}.json') if not os.path.exists(config_file): return {} # config file not found, return default config. try: with open(config_file, encoding='utf8') as f: config = json.load(f) return config except Exception as e: hoshino.logger.exception(e) return {} def _save_service_config(service): config_file = os.path.join(_service_config_dir, f'{service.name}.json') with open(config_file, 'w', encoding='utf8') as f: json.dump( { "name": service.name, "use_priv": service.use_priv, "manage_priv": service.manage_priv, "enable_on_default": service.enable_on_default, "visible": service.visible, "enable_group": list(service.enable_group), "disable_group": list(service.disable_group) }, f, ensure_ascii=False, indent=2) class ServiceFunc: def __init__(self, sv: "Service", func: Callable, only_to_me: bool, normalize_text: bool=False): self.sv = sv self.func = func self.only_to_me = only_to_me self.normalize_text = normalize_text self.__name__ = func.__name__ def __call__(self, *args, **kwargs): return self.func(*args, **kwargs) class Service: """将一组功能包装为服务, 提供增强的触发条件与分群权限管理. 支持的触发条件: `on_message`, `on_prefix`, `on_fullmatch`, `on_suffix`, `on_keyword`, `on_rex`, `on_command`, `on_natural_language` 提供接口： `scheduled_job`, `broadcast` 服务的配置文件格式为： { "name": "ServiceName", "use_priv": priv.NORMAL, "manage_priv": priv.ADMIN, "enable_on_default": true/false, "visible": true/false, "enable_group": [], "disable_group": [] } 储存位置： `~/.hoshino/service_config/{ServiceName}.json` """ def __init__(self, name, use_priv=None, manage_priv=None, enable_on_default=None, visible=None, help_=None, bundle=None): """ 定义一个服务 配置的优先级别：配置文件 > 程序指定 > 缺省值 """ assert not _re_illegal_char.search(name), r'Service name cannot contain character in `\/:*?"<>|.`' config = _load_service_config(name) self.name = name self.use_priv = config.get('use_priv') or use_priv or priv.NORMAL self.manage_priv = config.get( 'manage_priv') or manage_priv or priv.ADMIN self.enable_on_default = config.get('enable_on_default') if self.enable_on_default is None: self.enable_on_default = enable_on_default if self.enable_on_default is None: self.enable_on_default = True self.visible = config.get('visible') if self.visible is None: self.visible = visible if self.visible is None: self.visible = True self.help = help_ self.enable_group = set(config.get('enable_group', [])) self.disable_group = set(config.get('disable_group', [])) self.logger = log.new_logger(name, hoshino.config.DEBUG) assert self.name not in _loaded_services, f'Service name "{self.name}" already exist!' _loaded_services[self.name] = self _service_bundle[bundle or "通用"].append(self) @property def bot(self): return hoshino.get_bot() @staticmethod def get_loaded_services() -> Dict[str, "Service"]: return _loaded_services @staticmethod def get_bundles(): return _service_bundle def set_enable(self, group_id): self.enable_group.add(group_id) self.disable_group.discard(group_id) _save_service_config(self) self.logger.info(f'Service {self.name} is enabled at group {group_id}') def set_disable(self, group_id): self.enable_group.discard(group_id) self.disable_group.add(group_id) _save_service_config(self) self.logger.info( f'Service {self.name} is disabled at group {group_id}') def check_enabled(self, group_id): return bool((group_id in self.enable_group) or (self.enable_on_default and group_id not in self.disable_group)) def _check_all(self, ev: CQEvent): if ev.detail_type == 'private': return True else: gid = ev.group_id return self.check_enabled(gid) and not priv.check_block_group(gid) and priv.check_priv(ev, self.use_priv) async def get_enable_groups(self) -> dict: """获取所有启用本服务的群 @return { group_id: [self_id1, self_id2] } """ gl = defaultdict(list) for sid in hoshino.get_self_ids(): try: sgl = await self.bot.get_group_list(self_id=sid) except CQHttpError: sgl = [] sgl = set(g['group_id'] for g in sgl) if self.enable_on_default: sgl = sgl - self.disable_group else: sgl = sgl & self.enable_group for g in sgl: gl[g].append(sid) return gl def on_message(self, *events) -> Callable: def deco(func) -> Callable: @wraps(func) async def wrapper(ctx): if self._check_all(ctx): try: return await func(self.bot, ctx) except Exception as e: self.logger.error(f'{type(e)} occured when {func.__name__} handling message {ctx["message_id"]}.') self.logger.exception(e) return return self.bot.on_message(*events)(wrapper) #return self.bot.on_message(event)(wrapper) if type(event) is str else self.bot.on_message(*events)(wrapper) return deco def on_prefix(self, *prefix, only_to_me=False) -> Callable: if len(prefix) == 1 and not isinstance(prefix[0], str): prefix = prefix[0] def deco(func) -> Callable: sf = ServiceFunc(self, func, only_to_me) for p in prefix: if isinstance(p, str): trigger.prefix.add(p, sf) else: self.logger.error(f'Failed to add prefix trigger `{p}`, expecting `str` but `{type(p)}` given!') return func return deco def on_fullmatch(self, *word, only_to_me=False) -> Callable: if len(word) == 1 and not isinstance(word[0], str): word = word[0] def deco(func) -> Callable: @wraps(func) async def wrapper(bot: HoshinoBot, event: CQEvent): if len(event.message) != 1 or event.message[0].data.get('text'): self.logger.info(f'Message {event.message_id} is ignored by fullmatch condition.') return return await func(bot, event) sf = ServiceFunc(self, wrapper, only_to_me) for w in word: if isinstance(w, str): trigger.prefix.add(w, sf) else: self.logger.error(f'Failed to add fullmatch trigger `{w}`, expecting `str` but `{type(w)}` given!') return func # func itself is still func, not wrapper. wrapper is a part of trigger. # so that we could use multi-trigger freely, regardless of the order of decorators. # ``` # """the order doesn't matter""" # @on_keyword(...) # @on_fullmatch(...) # async def func(...): # ... # ``` return deco def on_suffix(self, *suffix, only_to_me=False) -> Callable: if len(suffix) == 1 and not isinstance(suffix[0], str): suffix = suffix[0] def deco(func) -> Callable: sf = ServiceFunc(self, func, only_to_me) for s in suffix: if isinstance(s, str): trigger.suffix.add(s, sf) else: self.logger.error(f'Failed to add suffix trigger `{s}`, expecting `str` but `{type(s)}` given!') return func return deco def on_keyword(self, *keywords, only_to_me=False, normalize=True) -> Callable: if len(keywords) == 1 and not isinstance(keywords[0], str): keywords = keywords[0] def deco(func) -> Callable: sf = ServiceFunc(self, func, only_to_me, normalize) for kw in keywords: if isinstance(kw, str): trigger.keyword.add(kw, sf) else: self.logger.error(f'Failed to add keyword trigger `{kw}`, expecting `str` but `{type(kw)}` given!') return func return deco def on_rex(self, rex: Union[str, re.Pattern], only_to_me=False, normalize=True) -> Callable: if isinstance(rex, str): rex = re.compile(rex) def deco(func) -> Callable: sf = ServiceFunc(self, func, only_to_me, normalize) if isinstance(rex, re.Pattern): trigger.rex.add(rex, sf) else: self.logger.error(f'Failed to add rex trigger `{rex}`, expecting `str` or `re.Pattern` but `{type(rex)}` given!') return func return deco def on_command(self, name, *, only_to_me=False, deny_tip=None, **kwargs) -> Callable: kwargs['only_to_me'] = only_to_me def deco(func) -> Callable: @wraps(func) async def wrapper(session): #if session.ctx['message_type'] != 'group' or session.ctx['message_type']!='guild': # return if not self.check_enabled(session.ctx['group_id']): self.logger.debug( f'Message {session.ctx["message_id"]} is command of a disabled service, ignored.' ) if deny_tip: session.finish(deny_tip, at_sender=True) return if self._check_all(session.ctx): try: ret = await func(session) self.logger.info( f'Message {session.ctx["message_id"]} is handled as command by {func.__name__}.' ) return ret except CanceledException: raise _FinishException except (_PauseException, _FinishException, SwitchException) as e: raise e except Exception as e: self.logger.error(f'{type(e)} occured when {func.__name__} handling message {session.ctx["message_id"]}.') self.logger.exception(e) return nonebot.on_command(name, **kwargs)(wrapper) return deco def on_natural_language(self, keywords=None, **kwargs) -> Callable: def deco(func) -> Callable: @wraps(func) async def wrapper(session: nonebot.NLPSession): if self._check_all(session.ctx): try: ret = await func(session) self.logger.info( f'Message {session.ctx["message_id"]} is handled as natural language by {func.__name__}.' ) return ret except CanceledException: raise _FinishException except (_PauseException, _FinishException, SwitchException) as e: raise e except Exception as e: self.logger.error(f'{type(e)} occured when {func.__name__} handling message {session.ctx["message_id"]}.') self.logger.exception(e) return nonebot.on_natural_language(keywords, **kwargs)(wrapper) return deco def scheduled_job(self, *args, **kwargs) -> Callable: kwargs.setdefault('timezone', pytz.timezone('Asia/Shanghai')) kwargs.setdefault('misfire_grace_time', 60) kwargs.setdefault('coalesce', True) def deco(func: Callable[[], Any]) -> Callable: @wraps(func) async def wrapper(): try: self.logger.info(f'Scheduled job {func.__name__} start.') ret = await func() self.logger.info(f'Scheduled job {func.__name__} completed.') return ret except Exception as e: self.logger.error(f'{type(e)} occured when doing scheduled job {func.__name__}.') self.logger.exception(e) return nonebot.scheduler.scheduled_job(*args, **kwargs)(wrapper) return deco async def broadcast(self, msgs, TAG='', interval_time=0.5, randomiser=None): bot = self.bot if isinstance(msgs, (str, MessageSegment, Message)): msgs = (msgs, ) groups = await self.get_enable_groups() for gid, selfids in groups.items(): try: for msg in msgs: await asyncio.sleep(interval_time) msg = randomiser(msg) if randomiser else msg await bot.send_group_msg(self_id=random.choice(selfids), group_id=gid, message=msg) l = len(msgs) if l: self.logger.info(f"群{gid} 投递{TAG}成功 共{l}条消息") except Exception as e: self.logger.error(f"群{gid} 投递{TAG}失败：{type(e)}") self.logger.exception(e) def on_request(self, *events): def deco(func): @wraps(func) async def wrapper(session): if not self.check_enabled(session.event.group_id): return return await func(session) return nonebot.on_request(*events)(wrapper) return deco def on_notice(self, *events): def deco(func): @wraps(func) async def wrapper(session): if not self.check_enabled(session.event.group_id): return return await func(session) return nonebot.on_notice(*events)(wrapper) return deco sulogger = log.new_logger('sucmd', hoshino.config.DEBUG) def sucmd(name, force_private=True, **kwargs) -> Callable: kwargs['privileged'] = True kwargs['only_to_me'] = False def deco(func) -> Callable: @wraps(func) async def wrapper(session: CommandSession): if session.event.user_id not in hoshino.config.SUPERUSERS: return if force_private and session.event.detail_type != 'private': await session.send('> This command should only be used in private session.') return try: return await func(session) except CanceledException: raise _FinishException except (_PauseException, _FinishException, SwitchException): raise except Exception as e: sulogger.error(f'{type(e)} occured when {func.__name__} handling message {session.event.message_id}.') sulogger.exception(e) return nonebot.on_command(name, **kwargs)(wrapper) return deco