update aiogram; add channel sending

main.py

@@ -1,11 +1,14 @@
 import asyncio
+import base64
 import logging
 import os
 import random
+import re
+import traceback
 from asyncio import sleep
 from io import BytesIO
 from pathlib import Path
-from tempfile import TemporaryFile, TemporaryDirectory
+from tempfile import TemporaryDirectory
 from time import time
 from typing import List
 
@@ -15,10 +18,10 @@ from PIL import Image
 
 import httpx
 import redis.asyncio as aioredis
-from aiogram import Bot, Dispatcher
-from aiogram.dispatcher import filters
-from aiogram.types import Message, ParseMode, ChatActions, InputFile
-from aiogram.utils import executor, exceptions
+from aiogram import Bot, Dispatcher, filters, exceptions, F
+from aiogram.enums import ChatAction, ParseMode
+from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, BufferedInputFile, \
+    CallbackQuery
 import dotenv
 
 from e621 import E621, E621Post
@@ -30,13 +33,13 @@ e621 = E621()
 
 logging.basicConfig(level=logging.INFO)
 bot = Bot(token=os.environ['BOT_TOKEN'])
-dp = Dispatcher(bot)
+dp = Dispatcher()
 
-user_ids = list(map(int, os.environ['USERS'].split(',')))
+ChatFilter = F.chat.id.in_(set(map(int, os.environ['USERS'].split(','))))
 
 
 async def send_post(post: E621Post, tag_list: List[str]):
-    await bot.send_chat_action(int(os.environ['SEND_CHAT']), action=ChatActions.TYPING)
+    await bot.send_chat_action(int(os.environ['SEND_CHAT']), action=ChatAction.TYPING)
     monitored_tags = set(post.tags.flatten()) & set(tag_list)
     artist_tags = post.tags.artist
     character_tags = post.tags.character
@@ -83,6 +86,10 @@ async def send_post(post: E621Post, tag_list: List[str]):
                         src_path.unlink()
                         mp4_path.unlink()
                 elif post.file.ext in ('png', 'jpg'):
+                    markup = InlineKeyboardMarkup(inline_keyboard=[[
+                        InlineKeyboardButton(text='Safe', callback_data='send pics'),
+                        InlineKeyboardButton(text='NSFW', callback_data='send nsfw'),
+                    ]])
                     # if post.file.size > 10000000:
                     logging.warning('compressing')
                     dl_im = Image.open(file).convert('RGBA')
@@ -98,11 +105,11 @@ async def send_post(post: E621Post, tag_list: List[str]):
                     file = BytesIO()
                     composite.save(file, format='JPEG')
                     file.seek(0)
-                    file.name = 'file.jpg'
                     await bot.send_photo(int(os.environ['SEND_CHAT']),
-                                         file,
+                                         BufferedInputFile(file.read(), 'file.jpg'),
                                          caption=caption,
-                                         parse_mode=ParseMode.HTML)
+                                         parse_mode=ParseMode.HTML,
+                                         reply_markup=markup)
             await redis.sadd('e621:sent', post.id)
         except exceptions.TelegramAPIError as e:
             logging.exception(e)
@@ -132,9 +139,9 @@ async def check_updates():
                 await sleep(1)
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['add'])
+@dp.message(filters.Command('add'), ChatFilter)
 async def add_tag(msg: Message):
-    args = msg.get_args()
+    args = ' '.join(msg.text.split()[1:])
     if not args:
         await msg.reply('Please provide tag to subscribe to')
         return
@@ -145,23 +152,25 @@ async def add_tag(msg: Message):
     await msg.reply(f'Tags {args} added')
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['mark_old_as_sent'])
+@dp.message(filters.Command('mark_old_as_sent'), ChatFilter)
 async def mark_old_as_sent(msg: Message):
-    tag_list = [t.decode() for t in await redis.smembers('e621:subs')]
-    m = await msg.reply(f'0/{len(tag_list)} tags have old posts marked as sent')
-    for i, tag in enumerate(tag_list, 1):
-        posts = await e621.get_posts(tag)
-        await redis.sadd('e621:sent', *[post.id for post in posts])
-        await m.edit_text(f'{i}/{len(tag_list)} tags have old posts marked as sent')
-        await sleep(1)
-    await m.edit_text(f'Done marking old posts as sent for {len(tag_list)} tags')
+    logging.warning('Waiting for lock...')
+    async with redis.lock('e621:update'):
+        tag_list = [t.decode() for t in await redis.smembers('e621:subs')]
+        m = await msg.reply(f'0/{len(tag_list)} tags have old posts marked as sent')
+        for i, tag in enumerate(tag_list, 1):
+            posts = await e621.get_posts(tag)
+            await redis.sadd('e621:sent', *[post.id for post in posts])
+            await m.edit_text(f'{i}/{len(tag_list)} tags have old posts marked as sent')
+            await sleep(1)
+        await m.edit_text(f'Done marking old posts as sent for {len(tag_list)} tags')
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), regexp=r'^\/del_\S+$')
+@dp.message(filters.Command(re.compile(r'del_\S+')), ChatFilter)
 async def del_tag(msg: Message):
     args = msg.text[5:]
     if not args:
-        await msg.reply('Please provide tag to subscribe to')
+        await msg.reply('Please provide tag to unsubscribe from')
         return
     if ' ' in args:
         await msg.reply('Tag should not contain spaces')
@@ -173,9 +182,9 @@ async def del_tag(msg: Message):
     await msg.reply(f'Tag {args} removed')
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['del'])
+@dp.message(filters.Command('del'), ChatFilter)
 async def del_command(msg: Message):
-    args = msg.get_args()
+    args = ' '.join(msg.text.split()[1:])
     if not args:
         await msg.reply('Please provide tag to subscribe to')
         return
@@ -184,7 +193,7 @@ async def del_command(msg: Message):
     await msg.reply(f'Tags {args} removed')
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['list'])
+@dp.message(filters.Command('list'), ChatFilter)
 async def list_tags(msg: Message):
     tags = [t.decode() for t in await redis.smembers('e621:subs')]
     tags.sort()
@@ -202,14 +211,14 @@ async def list_tags(msg: Message):
     await msg.reply(f'Monitored tags:\n\n{lines}')
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['update'])
+@dp.message(filters.Command('update'), ChatFilter)
 async def update(msg: Message):
     await check_updates()
 
 
-@dp.message_handler(filters.IDFilter(chat_id=user_ids), commands=['test'])
+@dp.message(filters.Command('test'), ChatFilter)
 async def test(msg: Message):
-    args = msg.get_args()
+    args = ' '.join(msg.text.split()[1:])
     if not args:
         await msg.reply('Please provide post id')
         return
@@ -221,6 +230,33 @@ async def test(msg: Message):
     await send_post(post[0], [])
 
 
+@dp.callback_query(F.data.startswith('send '))
+async def send_callback(cq: CallbackQuery):
+    _, destination = cq.data.split()
+    img_bytes = BytesIO()
+    await bot.download(cq.message.photo[-1], img_bytes)
+    img_bytes.seek(0)
+    data = base64.b64encode(img_bytes.read()).decode()
+    async with httpx.AsyncClient() as client:
+        try:
+            r = await client.post(f'https://bots.bakatrouble.me/bots_rpc/{destination}/', json={
+                "method": "post_photo",
+                "params": [data, True],
+                "jsonrpc": "2.0",
+                "id": 0,
+            })
+            resp = r.json()
+            if 'result' in resp and resp['result'] == True:
+                await cq.answer('Sent')
+            elif 'result' in resp and resp['result'] == 'duplicate':
+                await cq.answer('Duplicate')
+            else:
+                raise Exception(resp)
+        except:
+            traceback.print_exc()
+            await cq.answer('An error has occurred, check logs')
+
+
 async def background_on_start():
     await redis.delete('e621:update')
     while True:
@@ -233,9 +269,14 @@ async def background_on_start():
         await asyncio.sleep(600)
 
 
-async def on_bot_startup(dp: Dispatcher):
+async def on_bot_startup():
     asyncio.create_task(background_on_start())
 
 
+async def main():
+    dp.startup.register(on_bot_startup)
+    await dp.start_polling(bot)
+
+
 if __name__ == '__main__':
-    executor.start_polling(dp, on_startup=on_bot_startup)
+    asyncio.run(main())