import asyncio import logging import os from io import BytesIO from PIL import Image import httpx import redis.asyncio as aioredis from aiogram import Bot, Dispatcher from aiogram.types import Message from aiogram.utils import executor, exceptions import dotenv from e621 import E621 dotenv.load_dotenv('.env') redis = aioredis.from_url('redis://localhost') e621 = E621() logging.basicConfig(level=logging.INFO) bot = Bot(token=os.environ['BOT_TOKEN']) dp = Dispatcher(bot) async def check_updates(): async with redis.lock('e621:update'): tag_list = [t.decode() for t in await redis.smembers('e621:subs')] tag_list.sort() for tl_idx in range(0, len(tag_list), 40): tags = ' '.join(f'~{tag}' for tag in tag_list[tl_idx: tl_idx + 40]) posts = await e621.get_posts(tags) if not posts: return already_sent = await redis.smismember('e621:sent', [p.id for p in posts]) for i in range(len(posts)): if already_sent[i]: continue post = posts[i] monitored_tags = set(post.tags.flatten()) & set(tag_list) caption = f'Monitored tags: {" ".join(monitored_tags)}\n' \ f'Artist: {" ".join(post.tags.artist)}\n' \ f'Character: {" ".join(post.tags.character)}\n\n' \ f'https://e621.net/posts/{post.id}' if post.file.url: try: logging.warning(post.file.url) async with httpx.AsyncClient() as client: file = BytesIO() file.write((await client.get(post.file.url)).content) file.name = f'file.{post.file.ext}' file.seek(0) if post.file.ext == 'webm': await bot.send_video(int(os.environ['SEND_CHAT']), file, width=post.file.width, height=post.file.height, thumb=post.preview.url, caption=caption) elif post.file.ext == 'gif': await bot.send_animation(int(os.environ['SEND_CHAT']), file, width=post.file.width, height=post.file.height, thumb=post.preview.url, caption=caption) elif post.file.ext in ('png', 'jpg'): if post.file.size > 10000000: logging.warning('compressing') dl_im = Image.open(file).convert('RGBA') if dl_im.size[0] > 2000 or dl_im.size[1] > 2000: larger_dimension = max(dl_im.size) ratio = 2000 / larger_dimension dl_im = dl_im.resize((int(dl_im.size[0] * ratio), int(dl_im.size[1] * ratio)), Image.LANCZOS) im = Image.new('RGBA', dl_im.size, (255, 255, 255)) composite = Image.alpha_composite(im, dl_im).convert('RGB') file = BytesIO() composite.save(file, format='JPEG') file.seek(0) file.name = 'file.jpg' await bot.send_photo(int(os.environ['SEND_CHAT']), file, caption=caption) await redis.sadd('e621:sent', post.id) except exceptions.TelegramAPIError as e: logging.exception(e) @dp.message_handler(commands=['add']) async def add_tag(msg: Message): args = msg.get_args() if not args: await msg.reply('Please provide tag to subscribe to') return for tag in args.split(): await redis.sadd('e621:subs', tag) await msg.reply(f'Tags {args} added') @dp.message_handler(regexp=r'^\/del_\S+$') async def del_tag(msg: Message): args = msg.text[5:] if not args: await msg.reply('Please provide tag to subscribe to') return if ' ' in args: await msg.reply('Tag should not contain spaces') return if not await redis.sismember('e621:subs', args): await msg.reply('Tag not found') return await redis.srem('e621:subs', args) await msg.reply(f'Tag {msg} removed') @dp.message_handler(commands=['list']) async def list_tags(msg: Message): tags = [t.decode() for t in await redis.smembers('e621:subs')] lines = [] for tag in tags: lines.append(f'- {tag} [/del_{tag}]') lines = "\n".join(lines) await msg.reply(f'Monitored tags:\n\n{lines}') @dp.message_handler(commands=['update']) async def update(msg: Message): await check_updates() async def background_on_start(): while True: logging.warning('Checking updates...') await check_updates() await asyncio.sleep(60) async def on_bot_startup(dp: Dispatcher): asyncio.create_task(background_on_start()) if __name__ == '__main__': executor.start_polling(dp, on_startup=on_bot_startup)