initial commit

kurocore/main/handler.py

@@ -0,0 +1,256 @@
+import asyncio
+import traceback
+
+from aiohttp.web import run_app
+from sentry_sdk import init as sentry_init, capture_exception, set_user
+
+from .event import ChatEvent, Event
+from kurocore.utils.database.database import Database
+from kurocore.utils.vk.longpoll import VKLongPoll, VkBotEventType
+from kurocore.utils.vk.vk import VK
+from ..logger import BotLogger
+
+
+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()
+
+                        await p._process_command(command, msg, 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):
+        from kurocore import Message
+        if ((event.type == VkBotEventType.MESSAGE_NEW and 'action' not in event.obj)
+                or event.type == VkBotEventType.MESSAGE_EVENT):
+            msg = Message(self.session, 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.session, self.api, event.obj['action'])
+            msg = Message(self.session, 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.session, 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)