From 921bbc1809b82af4f58686b213b723cc82d5cc15 Mon Sep 17 00:00:00 2001 From: BarsTiger Date: Mon, 20 Feb 2023 23:30:36 +0200 Subject: [PATCH] Added imginfo and hidden hash command --- bot/handlers/admin/__init__.py | 2 ++ bot/handlers/admin/tools.py | 32 +++++++++++++++++++++++ bot/handlers/help_command/help_strings.py | 3 ++- bot/handlers/image_info/__init__.py | 6 +++++ bot/handlers/image_info/image_info.py | 31 ++++++++++++++++++++++ bot/handlers/register.py | 3 ++- bot/modules/get_hash/__init__.py | 0 bot/modules/get_hash/get_hash.py | 23 ++++++++++++++++ 8 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 bot/handlers/admin/tools.py create mode 100644 bot/handlers/image_info/__init__.py create mode 100644 bot/handlers/image_info/image_info.py create mode 100644 bot/modules/get_hash/__init__.py create mode 100644 bot/modules/get_hash/get_hash.py diff --git a/bot/handlers/admin/__init__.py b/bot/handlers/admin/__init__.py index 1e759b7..42fb38a 100644 --- a/bot/handlers/admin/__init__.py +++ b/bot/handlers/admin/__init__.py @@ -1,8 +1,10 @@ from bot.common import dp from .aliases import * from .reset import * +from .tools import * def register(): dp.register_message_handler(set_endpoint, commands='setendpoint') dp.register_message_handler(reset.resetqueue, commands='resetqueue') + dp.register_message_handler(tools.hash_command, commands='hash') diff --git a/bot/handlers/admin/tools.py b/bot/handlers/admin/tools.py new file mode 100644 index 0000000..f2602c6 --- /dev/null +++ b/bot/handlers/admin/tools.py @@ -0,0 +1,32 @@ +from aiogram import types +from bot.db import db, DBTables +from bot.utils.cooldown import throttle +from bot.keyboards.exception import get_exception_keyboard +from bot.utils.trace_exception import PrettyException +from bot.modules.get_hash.get_hash import get_hash + + +@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins')) +async def hash_command(message: types.Message): + try: + if not hasattr(message.reply_to_message, 'photo') or not hasattr(message.reply_to_message, 'document'): + await message.reply('❌ REPLY with this command on picture or file', parse_mode='html') + return + + if len(message.reply_to_message.photo) < 1: + file_hash = await get_hash(message.reply_to_message.document.file_id) + + else: + file_hash = await get_hash(message.reply_to_message.photo[0].file_id) + + await message.reply((lambda x: x if x else "❌ Hash not returned")(file_hash)) + + except IndexError: + await message.reply('❌ Reply with this command on PICTURE OR FILE', parse_mode='html') + + except Exception as e: + exception_id = f'{message.message_thread_id}-{message.message_id}' + db[DBTables.exceptions][exception_id] = PrettyException(e) + await message.reply('❌ Error happened while processing your request', parse_mode='html', + reply_markup=get_exception_keyboard(exception_id)) + return diff --git a/bot/handlers/help_command/help_strings.py b/bot/handlers/help_command/help_strings.py index c34d2c9..a224509 100644 --- a/bot/handlers/help_command/help_strings.py +++ b/bot/handlers/help_command/help_strings.py @@ -1,3 +1,4 @@ help_data = { - 'setendpoint': '(admin) Set StableDiffusion API endpoint' + 'setendpoint': '(admin) Set StableDiffusion API endpoint', + 'imginfo': 'Get information about image, that was generated using this bot' } diff --git a/bot/handlers/image_info/__init__.py b/bot/handlers/image_info/__init__.py new file mode 100644 index 0000000..75c1387 --- /dev/null +++ b/bot/handlers/image_info/__init__.py @@ -0,0 +1,6 @@ +from bot.common import dp +from .image_info import * + + +def register(): + dp.register_message_handler(imginfo, commands='imginfo') diff --git a/bot/handlers/image_info/image_info.py b/bot/handlers/image_info/image_info.py new file mode 100644 index 0000000..c4f71a3 --- /dev/null +++ b/bot/handlers/image_info/image_info.py @@ -0,0 +1,31 @@ +from aiogram import types +from bot.db import db, DBTables +from bot.utils.cooldown import throttle +from bot.keyboards.exception import get_exception_keyboard +from bot.utils.trace_exception import PrettyException + + +@throttle(cooldown=10, admin_ids=db[DBTables.config].get('admins')) +async def imginfo(message: types.Message): + try: + if not hasattr(message.reply_to_message, 'photo'): + await message.reply('❌ Reply with this command on picture', parse_mode='html') + return + + if not (original_r := db[DBTables.generated].get(message.reply_to_message.photo[0].file_unique_id)): + await message.reply('❌ This picture wasn\'t generated using this bot ' + 'or doesn\'t exist in database. Note this only works on ' + 'files forwarded from bot.', parse_mode='html') + return + + await message.reply(str(original_r)) + + except IndexError: + await message.reply('❌ Reply with this command on PICTURE', parse_mode='html') + + except Exception as e: + exception_id = f'{message.message_thread_id}-{message.message_id}' + db[DBTables.exceptions][exception_id] = PrettyException(e) + await message.reply('❌ Error happened while processing your request', parse_mode='html', + reply_markup=get_exception_keyboard(exception_id)) + return diff --git a/bot/handlers/register.py b/bot/handlers/register.py index b050c6b..75ded6d 100644 --- a/bot/handlers/register.py +++ b/bot/handlers/register.py @@ -3,12 +3,13 @@ from rich import print def register_handlers(): from bot.handlers import ( - initialize, admin, help_command, txt2img + initialize, admin, help_command, txt2img, image_info ) initialize.register() admin.register() help_command.register() txt2img.register() + image_info.register() print('[gray]All handlers registered[/]') diff --git a/bot/modules/get_hash/__init__.py b/bot/modules/get_hash/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bot/modules/get_hash/get_hash.py b/bot/modules/get_hash/get_hash.py new file mode 100644 index 0000000..6ef898d --- /dev/null +++ b/bot/modules/get_hash/get_hash.py @@ -0,0 +1,23 @@ +import os +from bot.common import bot +import hashlib +import aiohttp + + +async def get_hash(file_id: str): + url = bot.get_file_url( + (await bot.get_file(file_id)).file_path + ) + + async with aiohttp.ClientSession() as session: + async with session.get(url) as resp: + assert resp.status == 200 + data = await resp.read() + + with open('file_to_get_hash', "wb") as f: + f.write(data) + + file_hash = hashlib.md5(open('file_to_get_hash', 'rb').read()).hexdigest() + + os.remove('file_to_get_hash') + return file_hash