Exemple #1
0
class ParseBot(object):
    bot = None
    states = None
    oops_handler = None

    def __init__(self):
        self.states = {}

    def init(self):
        self.bot = Bot(api_token=config.TELEGRAM_TOKEN)
        from bot.conversation import init_dialogs
        init_dialogs(self)
        self.bot.default(self.get_default_handler())
        yield_data.add_observer('pbot', self.data_receive_observer)

    async def data_receive_observer(self, tg_id, data):
        log.debug("DATA RECEIVED: to {}, data: {}", tg_id, data)
        chat = self.bot.private(tg_id)
        return chat.send_text('Look what I found!\n{}'.format(data))

    def add_state_handler(self, state, hlr):
        self.states[state] = hlr

    def add_default_state_handler(self, hlr):
        self.oops_handler = hlr

    def add_callback(self, regexp, fn):
        return self.bot.add_callback(regexp, fn)

    def add_command(self, regexp, fn):
        return self.bot.add_command(regexp, fn)

    def add_inline(self, regexp, fn):
        return self.bot.add_inline(regexp, fn)

    def get_default_handler(self):
        @with_user_data
        async def def_handler(chat, message, user_data):
            state = user_data.get('state')
            if state:
                handler = self.states.get(state)
                if handler:
                    return await handler(chat, message, user_data)
            if self.oops_handler:
                return await self.oops_handler(chat, message, user_data)
            return await chat.send_text('Oops! Something went wrong, for now we can start again and see how it goes')
        return def_handler

    def run(self, webhook=config.WEBHOOK_URL):
        if webhook:
            self.bot.run_webhook(webhook)
        else:
            self.bot.run(config.DEBUG)
Exemple #2
0
class CamBot:
    def __init__(self, agent: 'gphotos.GooglePhotosManager',
                 manager: vkmanager.VKManager):
        self._bot = Bot(conf.bot_token, proxy=conf.tele_proxy)
        self.session = self._bot.session
        self.loop = self._bot.loop
        self.menu_markup = Menu()
        self.init_handlers()
        self.agent = agent
        self.vk_manager = manager

    def init_handlers(self):
        self._bot.add_command(r'/mov (.+) (.+)', self.mov)
        self._bot.add_command(r'/push_vk (.+) (.+)', self.push_vk)
        self._bot.add_command(r'/check (.+) (.+)', self.check_album)
        self._bot.add_command(r'/full_check (.+)', self.full_check)
        self._bot.add_command(r'/clear (.+)', self.clear_command)
        self._bot.add_command(r'/reg', reg)
        self._bot.add_command(r'/ch', self.reg_channel)
        self._bot.add_command(r'/photo_reg', self.reg_photo_channel)
        self._bot.add_command(r'/menu', self.menu)
        self._bot.add_command(r'/all', self.img_all_cams)
        self._bot.add_command(r'/stats (.+)', self.stats_command)
        self._bot.add_command(r'/stats', self.stats_command)
        self._bot.add_command(r'/lstats (.+)', self.lstats_command)
        self._bot.add_command(r'/lstats', self.lstats_command)
        self._bot.add_command(r'/dbdata', self.db_data)
        self._bot.add_command(r'/daily', self.daily_movie_group_command)
        self._bot.add_command(r'/push_on', self.push_vk_on)
        self._bot.add_command(r'/push_off', self.push_vk_off)
        self._bot.add_callback(r'regular (.+)', regular)
        self._bot.add_callback(r'today (.+)', today)
        self._bot.add_callback(r'weekly (.+)', weekly)
        self._bot.add_callback(r'select (.+)', self.select)
        self._bot.add_callback(r'back', self.back)
        self._bot.add_callback(r'img (.+)', self.img_callback)
        self._bot.add_callback(r'choose_cam (.+)', self.choose_cam_callback)
        self._bot.add_callback(r'choose_photo_cam (.+)',
                               self.choose_photo_cam_callback)
        self._bot.add_callback(r'sync (.+)', self.sync_gphotos)
        self._bot.add_callback(r'gsnc (.+)', self.run_sync_gphotos)
        self._bot.add_callback(r'remove (.+)', self.remove_folder)
        self._bot.add_callback(r'post (.+) (.+)', self.post_photo)
        self._bot.add_callback(r'clear_cb (.+)', self.clear_callback)
        self._bot.add_callback(r'check_cb (.+)', self.full_check_callback)
        self._bot.callback(unhandled_callbacks)

    def stop(self):
        self._bot.stop()

    @ThreadSwitcherWithDB.optimized
    async def daily_stats(self):
        await self.stats_request(pendulum.yesterday(), self.notify_admins)

    @ThreadSwitcherWithDB.optimized
    async def daily_movie(self, cam: Cam):
        day = datetime.datetime.now() - datetime.timedelta(days=1)
        day = day.strftime('%d_%m_%Y')
        loop = asyncio.get_event_loop()
        with concurrent.futures.ThreadPoolExecutor() as pool:
            try:
                clip = await loop.run_in_executor(pool,
                                                  lambda: make_movie(cam, day))
            except FileNotFoundError as exc:
                logger.exception(exc)
                await self.notify_admins(
                    f'File {exc.filename} not found for daily movie {cam.name}: {day}'
                )
                return
            except Exception as exc:
                logger.exception(exc)
                await self.notify_admins(
                    f'Error during making daily movie for {cam.name}: {day}')
                return
        if cam.update_channel:
            async with db_in_thread():
                channels = db.query(Channel).filter(
                    Channel.cam == cam.name).all()
            for channel in channels:
                await send_video(Chat(self._bot, channel.chat_id), clip)
        await self.notify_admins(f'Daily movie for {cam.name}: {day} ready!')
        for chat in await self.admin_chats():
            await send_video(Chat(self._bot, chat.chat_id), clip)

    async def push_vk(self, chat, match):
        cam = await get_cam(match.group(1), chat)
        if not cam:
            return
        day = match.group(2)
        path = Path(conf.root_dir
                    ) / 'data' / cam.name / 'regular' / 'clips' / f'{day}.mp4'
        if not path.exists():
            await chat.send_text('Movie file does not exist!')
            return
        try:
            await self.vk_manager.new_post(cam.name, str(path),
                                           day.replace('_', ' '),
                                           day.replace('_', '/'))
        except vkmanager.VKManagerError as exc:
            logger.exception('Error during pushing video to vk')
            await chat.send_text(exc.detail)
        except Exception:
            logger.exception('Unhandled exception during pushing video to vk')
            await chat.send_text('Unhandled error!')
        await chat.send_text('Movie successfully published')

    async def mov(self, chat, match):
        """
        Make movie for specified cam and day. Example: /mov favcam 25_04_2019
        :param chat:
        :param match:
        :return:
        """
        cam = await get_cam(match.group(1), chat)
        if not cam:
            return
        day = match.group(2)
        loop = asyncio.get_event_loop()
        with concurrent.futures.ThreadPoolExecutor() as pool:
            try:
                clip = await loop.run_in_executor(pool,
                                                  lambda: make_movie(cam, day))
            except Exception:
                logger.exception('Error during movie request')
                await self.notify_admins(
                    f'Error during movie request {day} {cam.name}')
                return
        await self.notify_admins(f'Video ready. Uploading..')
        with open(clip.path, 'rb') as clip:
            await chat.send_video(clip)

    async def daily_movie_group(self):
        for cam in sorted(conf.cameras_list, key=lambda k: k.offset):
            if cam.render_daily:
                await self.daily_movie(cam)
        await self.daily_stats()

    async def daily_photo_group(self):
        for cam in conf.cameras_list:
            image = await CamHandler(cam,
                                     self._bot.session).get_img(regular=False)
            if not image:
                await self.notify_admins(
                    f'Error during image request for {cam.name}')
                continue
            path = image.original_path if cam.resize else image.path
            await self._post_photo(cam, path)

    async def daily_movie_group_command(self, chat, match):
        logger.info('Forced daily movie group command')
        await self.daily_movie_group()

    async def push_vk_on(self, chat: Chat, match):
        loop = asyncio.get_event_loop()
        cmd = f'systemctl --user start {conf.vk_service}'.split()
        try:
            await loop.run_in_executor(None, lambda: subprocess_call(cmd))
        except Exception:
            msg = 'Error during starting vk push service'
            logger.exception(msg)
            await chat.send_text(msg)
            return
        await chat.send_text('vk push service started')

    async def push_vk_off(self, chat: Chat, match):
        loop = asyncio.get_event_loop()
        cmd = f'systemctl --user stop {conf.vk_service}'.split()
        try:
            await loop.run_in_executor(None, lambda: subprocess_call(cmd))
        except Exception:
            msg = 'Error during stopping vk push service'
            logger.exception(msg)
            await chat.send_text(msg)
            return
        await chat.send_text('vk push service stopped')

    async def img_all_cams(self, chat: Chat, match):
        for cam in conf.cameras_list:
            await self.img_handler(chat, cam)

    async def img_handler(self, chat: Chat, cam):
        image = await CamHandler(cam, self._bot.session).get_img(regular=False)
        if not image:
            await chat.send_text(f'Error during image request for {cam.name}')
            return
        path = image.original_path if cam.resize else image.path
        markup = Markup([[
            InlineKeyboardButton(text='post',
                                 callback_data=f'post {cam.name} {path.name}')
        ]])
        with open(path, 'rb') as image:
            await chat.send_photo(image, reply_markup=markup.to_json())

    async def img_callback(self, chat, cq, match):
        await cq.answer()
        cam = await get_cam(match.group(1), chat)
        if not cam:
            return
        await self.img_handler(chat, cam)

    @ThreadSwitcherWithDB.optimized
    async def reg_channel(self, chat: Chat, match):
        async with db_in_thread():
            channel = db.query(Channel).filter(
                Channel.chat_id == chat.id).one_or_none()
        if channel:
            await self.notify_admins(f'Channel {chat.id} already registered!')
            return
        await chat.send_text('Choose cam for channel',
                             reply_markup=CamerasChannel().options.to_json())

    @ThreadSwitcherWithDB.optimized
    async def reg_photo_channel(self, chat: Chat, match):
        async with db_in_thread():
            channel = db.query(PhotoChannel).filter(
                PhotoChannel.chat_id == chat.id).one_or_none()
        if channel:
            await self.notify_admins(f'Channel {chat.id} already registered!')
            return
        await chat.send_text(
            'Choose cam for photo channel',
            reply_markup=CamerasChannel('choose_photo_cam').options.to_json())

    @ThreadSwitcherWithDB.optimized
    async def choose_cam_callback(self, chat, cq, match):
        cam = match.group(1)
        async with db_in_thread():
            channel = Channel(chat_id=chat.id, cam=cam)
            db.add(channel)
            db.commit()
        await cq.answer(text=f'Added channel for {cam}')
        await self.notify_admins(text=f'Added channel {chat.id} for {cam}')

    @ThreadSwitcherWithDB.optimized
    async def choose_photo_cam_callback(self, chat, cq, match):
        cam = match.group(1)
        async with db_in_thread():
            channel = PhotoChannel(chat_id=chat.id, cam=cam)
            db.add(channel)
            db.commit()
        await cq.answer(text=f'Added photo channel for {cam}')
        await self.notify_admins(
            text=f'Added photo channel {chat.id} for {cam}')

    async def post_photo(self, chat, cq, match):
        cam = match.group(1)
        photo = match.group(2)
        cam = conf.cameras[cam]
        path = Path(conf.root_dir) / 'data' / cam.name / 'imgs'
        if cam.resize:
            path /= 'original'
        path = path / '_'.join(photo.split('_')[:3]) / photo
        await self._post_photo(cam, path)
        await cq.answer()

    @ThreadSwitcherWithDB.optimized
    async def _post_photo(self, cam: Cam, photo: Path):
        async with db_in_thread():
            channels = db.query(PhotoChannel).filter(
                PhotoChannel.cam == cam.name).all()
        for channel in channels:
            chat = Chat(self._bot, channel.chat_id)
            with open(photo, 'rb') as ph:
                await chat.send_photo(ph)

    @ThreadSwitcherWithDB.optimized
    async def notify_admins(self, text, **options):
        async with db_in_thread():
            admins = db.query(Admin).all()
        for admin in admins:
            await self._bot.send_message(admin.chat_id, text, **options)

    @ThreadSwitcherWithDB.optimized
    async def admin_chats(self):
        async with db_in_thread():
            return db.query(Admin).all()

    async def menu(self, chat, match):
        await chat.send_text('Menu',
                             reply_markup=self.menu_markup.main_menu.to_json())

    async def select(self, chat: Chat, cq, match):
        await cq.answer()
        cam = match.group(1)
        await chat.edit_text(cq.src['message']['message_id'],
                             f'Camera: {cam}',
                             markup=dataclasses.asdict(
                                 self.menu_markup.cam_options[cam].markup))

    async def back(self, chat, cq, match):
        await cq.answer()
        await chat.edit_text(cq.src['message']['message_id'],
                             'Menu',
                             markup=dataclasses.asdict(
                                 self.menu_markup.main_menu))

    async def sync_gphotos(self, chat, cq, match):
        await cq.answer()
        cam = match.group(1)
        await chat.edit_text(cq.src['message']['message_id'],
                             f'Choose folder for {cam}',
                             markup=dataclasses.asdict(
                                 SyncFolders(cam).folders))

    async def run_sync_gphotos(self, chat, cq, match):
        _folder = match.group(1)
        folder = Path(conf.root_dir) / 'data' / _folder
        logger.debug(f'GOING TO SYNC FOLDER {folder}')
        await cq.answer(text=f'GOING TO SYNC FOLDER {folder}')
        await self.notify_admins(f'Started sync {folder}')
        try:
            await GooglePhotosManager().batch_upload(Path(folder))
        except Exception:
            logger.exception('Sync error!')
            await self.notify_admins(f'Error with {folder}!')
            return
        await self.notify_admins(f'{folder} successfully uploaded!')
        markup = Markup([[
            InlineKeyboardButton(text=f'{_folder}',
                                 callback_data=f'remove {_folder}')
        ]])
        await chat.send_text(f'Remove folder {folder.name}',
                             reply_markup=markup.to_json())

    async def remove_folder(self, chat, cq, match):
        await cq.answer(text='Removing folder..')
        folder = match.group(1)
        folder = Path(conf.root_dir) / 'data' / folder
        shutil.rmtree(folder)
        await chat.send_text('Successfully removed!')

    async def stats_command(self, chat: Chat, match):
        try:
            day = pendulum.from_format(match.group(1), 'DD_MM_YYYY')
        except IndexError:
            day = pendulum.today()
        await self.stats_request(day, chat.send_text)

    async def lstats_command(self, chat: Chat, match):
        try:
            day = pendulum.from_format(match.group(1), 'DD_MM_YYYY')
        except IndexError:
            day = pendulum.today()
        await self.stats_request(day, chat.send_text)

    @ThreadSwitcherWithDB.optimized
    async def db_data(self, chat: Chat, match):
        async with db_in_thread():
            md_data = db_data()
        await chat.send_text('\n'.join(md_data), parse_mode='Markdown')

    async def stats_request(self, day: pendulum.DateTime, send_command):
        logger.info(f'Getting stats info for {day}')
        try:
            markdown_result = await self.local_stats_handler(day)
        except Exception:
            logger.exception('Error during stats request')
            await send_command('Error during request stats')
            return
        day = day.format('DD_MM_YYYY')
        markup = Markup([[
            InlineKeyboardButton(text='check', callback_data=f'check_cb {day}')
        ], [
            InlineKeyboardButton(text='clear', callback_data=f'clear_cb {day}')
        ]])
        await send_command('\n'.join(markdown_result),
                           parse_mode='Markdown',
                           reply_markup=markup.to_json())

    async def stats_handler(self, day=None):
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(None, lambda: stats(day))
        album_stats = await self.agent.album_stats(day)
        markdown_result = [f'#stats *{day.format("DD/MM/YYYY")}*']
        for d in result['cameras']:
            stat = result['cameras'][d]
            count, size = stat['count'], convert_size(stat['size'])
            if count:
                avg = convert_size(stat['size'] / count)
            else:
                avg = 0
            media_count = album_stats[d]
            markdown_result.append(
                f'*{d}*: {count} - {media_count} - {size} - {avg} ')
        total = convert_size(result['total'])
        markdown_result.append(f'*total*: {total}')
        free = convert_size(result['free'])
        markdown_result.append(f'*free*: {free}')
        return markdown_result

    async def local_stats_handler(self, day=None):
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(None, lambda: stats(day))
        markdown_result = [f'#stats *{day.format("DD/MM/YYYY")}*']
        for d in result['cameras']:
            stat = result['cameras'][d]
            count, size = stat['count'], convert_size(stat['size'])
            if count:
                avg = convert_size(stat['size'] / count)
            else:
                avg = 0
            markdown_result.append(f'*{d}*: {count} - {size} - {avg} ')
        total = convert_size(result['total'])
        markdown_result.append(f'*total*: {total}')
        free = convert_size(result['free'])
        markdown_result.append(f'*free*: {free}')
        return markdown_result

    async def check_album(self, chat, match):
        cam = await get_cam(match.group(1), chat)
        if not cam:
            return
        day = match.group(2)
        await self.agent.check_album(cam, day)

    async def full_check_handler(self, chat, day):
        logger.info(f'Going to full check for {day}')
        for cam in conf.cameras_list:
            try:
                await self.agent.check_album(cam, day)
            except Exception:
                logger.exception(
                    f'Error during check and sync {cam.name} -- {day}')
                await chat.send_text(f'Error {cam.name} — {day}')
                continue
            await chat.send_text(f'Finished with {cam.name} — {day}')
        msg = f'Finished full check for {day}'
        logger.info(msg)
        await chat.send_text(msg)

    async def full_check(self, chat, match):
        day = match.group(1)
        await self.full_check_handler(chat, day)

    async def full_check_callback(self, chat, cq, match):
        day = match.group(1)
        await cq.answer(text=f'Running full check for {day}')
        await self.full_check_handler(chat, day)

    async def clear_handler(self, chat, day):
        logger.info(f'Going to clear for {day}')
        loop = asyncio.get_event_loop()
        for cam in conf.cameras_list:
            try:
                await loop.run_in_executor(None,
                                           lambda: clear_cam_storage(day, cam))
            except Exception:
                logger.exception(f'Error during clear {cam.name} -- {day}')
                await chat.send_text(f'Error {cam.name} — {day}')
                continue
            await chat.send_text(f'Finished with {cam.name} — {day}')
        logger.info(f'Finished clear for {day}')

    async def clear_command(self, chat, match):
        day = match.group(1)
        await self.clear_handler(chat, day)

    async def clear_callback(self, chat, cq, match):
        day = match.group(1)
        await cq.answer(text=f'Cleaning for {day}')
        await self.clear_handler(chat, day)