You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.3 KiB
108 lines
3.3 KiB
import asyncio
|
|
import logging
|
|
import os
|
|
|
|
import redis.asyncio as aioredis
|
|
from aiogram import Bot, Dispatcher
|
|
from aiogram.types import Message
|
|
from aiogram.utils import executor
|
|
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')]
|
|
tags = ' '.join(f'~{tag}' for tag in tag_list)
|
|
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.ext == 'webm':
|
|
await bot.send_video(int(os.environ['SEND_CHAT']),
|
|
post.file.url,
|
|
width=post.file.width,
|
|
height=post.file.height,
|
|
thumb=post.preview.url,
|
|
caption=caption)
|
|
else:
|
|
await bot.send_photo(int(os.environ['SEND_CHAT']),
|
|
post.file.url,
|
|
caption=caption)
|
|
await redis.sadd('e621:sent', post.id)
|
|
|
|
|
|
@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:
|
|
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)
|