Added imginfo and hidden hash command

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')

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

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'
 }

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')

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

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[/]')

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