KuroCore/kurocore/main/handler.py

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)