import asyncio import traceback from aiohttp.web import run_app from sentry_sdk import init as sentry_init, capture_exception, set_user from kurocore.logger import BotLogger from kurocore.vk import VK, VKLongPoll, VkBotEventType from .event import ChatEvent, Event from .message import MessageContext, Message from kurocore.database import Database class Handler: __slots__ = ( 'config', 'session', 'api', 'plugins', 'middlewares', 'loop', 'app' ) def __init__(self, config, middlewares: list): self.config = config self.session = None self.api = None self.app = None self.plugins = [] self.middlewares = middlewares self.loop = asyncio.get_event_loop() def init(self): if not self.config.tokens: BotLogger.log.error('No access token!') exit() self.session = VK(self.config.tokens[0]) self.api = self.session.get_api() if not self.config.debug and self.config.sentry_dsn: sentry_init( self.config.sentry_dsn, traces_sample_rate=1.0 ) BotLogger.log.info('Sentry initialized!') for p in self.config.plugins: for init in p.init_methods: BotLogger.log.debug(f'Init: {p.__class__.__name__}') init.call() self.plugins.append(p) def shutdown(self): for p in self.plugins: for method in p.shutdown_methods: method.call() BotLogger.log.info('Bot has been shutdown!') async def check_payload(self, msg): payload = msg.payload['command'] if 'command' in msg.payload else '' args = msg.payload['args'] if 'args' in msg.payload else [] for p in self.plugins: if p.custom_checker: try: for before_process in p.before_process_methods: before_process() await p.custom_checker(msg, p) return True except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) if payload in p.payloads.keys(): try: msg.meta.update({'args': args}) for before_process in p.before_process_methods: before_process.call() valid, args = await p._validate_payload_args(payload, args) if not valid: await msg.answer('Неверное количество или тип аргументов!') return False for before_process in p.before_process_methods: before_process.call() if valid and args is not None: await p._process_payload_with_args(payload, msg, args) else: await p._process_payload(payload, msg) return True except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) else: return False async def check_command(self, msg): text = msg.text for prefix in self.config.prefixes: if text.startswith(prefix): text = text[len(prefix):] msg.meta['prefix'] = prefix break else: if msg.is_chat: return for p in self.plugins: if p.custom_checker: try: for before_process in p.before_process_methods: before_process() await p.custom_checker(msg, p) return except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) for command in p.commands: if text.startswith(command): # if p._is_vip_command(command): # user = await get_user_or_none(msg.from_id) # if not user or not user.group.is_vip: # return msg.answer('Для доступа к этой команде требует VIP доступ!') # # if p._is_admin_command(command): # user = await get_user_or_none(msg.from_id) # if not user or not user.group.is_admin: # return msg.answer('Данная комманда доступна только для администраторов!') try: msg.meta['cmd'] = command args = text[len(command) + 1:].split() msg.meta['args'] = args args_valid, args = await p._validate_command_args(command, args) if not args_valid: return await msg.answer('Неверное количество или тип аргументов!') for before_process in p.before_process_methods: before_process.call() ctx = MessageContext(self.api, msg, None) if Database.db: ctx.db = Database await p._process_command(command, ctx, args) return except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) async def check(self, msg): db = Database.db if db: if not db.is_closed(): BotLogger.log.debug('Connection reused!') else: db.connect() BotLogger.log.debug('Connection reopened!') else: BotLogger.log.debug('No database') for m in self.middlewares: await m(msg) for plugin in self.plugins: for method in plugin.before_check_methods: await method(msg) if not await self.check_payload(msg): await self.check_command(msg) if db: db.close() BotLogger.log.debug('Connection closed!') async def check_event(self, event: (ChatEvent, Event), msg): event_type = event.type for plugin in self.plugins: if event_type in plugin.chat_events.keys(): try: for before_process in plugin.before_process_methods: before_process.call() return await plugin.chat_events[event_type](event, msg) except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) elif event_type in plugin.events.keys(): try: for before_process in plugin.before_process_methods: before_process.call() return await plugin.events[event_type](event) except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e) def run(self): if self.config.callback.enabled: run_app(self.app, port=self.config.port) try: # Register all plugins tasks for p in self.config.plugins: for task in p.tasks: BotLogger.log.debug(f'registered task: {task.__name__}') self.loop.create_task(task()) self.loop.run_until_complete(self._run()) except KeyboardInterrupt: self.session.shutdown() # self.shutdown() async def handle_event(self, event): if ((event.type == VkBotEventType.MESSAGE_NEW and 'action' not in event.obj) or event.type == VkBotEventType.MESSAGE_EVENT): msg = Message(self.api, event.obj) if msg.user_id > 0: set_user({'id': msg.user_id}) await self.check(msg) elif event.type == VkBotEventType.MESSAGE_NEW and 'action' in event.obj: e = ChatEvent(self.api, event.obj['action']) msg = Message(self.api, event.obj) if msg.user_id > 0: set_user({'id': msg.user_id}) await self.check_event(e, msg) else: e = Event(self.api, event.raw) await self.check_event(e, None) async def _run(self): group = (await self.api.groups.getById())['groups'][0] BotLogger.log.info(f'Login as {group["name"]} (https://vk.com/public{group["id"]})') lp = VKLongPoll(self.config, self.session) await lp.init_lp() async for event in lp.listen(): try: await self.loop.create_task(self.handle_event(event)) except Exception as e: if self.config.debug: BotLogger.log.error(traceback.format_exc()) else: capture_exception(e)