Image info, endpoint encryption

bot/handlers/admin/aliases.py

@@ -1,5 +1,5 @@
 from aiogram import types
-from bot.db import db, DBTables
+from bot.db import db, DBTables, encrypt
 import validators
 from bot.config import ADMIN
 from bot.utils.cooldown import throttle
@@ -16,7 +16,7 @@ async def set_endpoint(message: types.Message):
         await message.reply("❌ Specify correct url for endpoint")
         return
 
-    db[DBTables.config]['endpoint'] = message.get_args()
+    db[DBTables.config]['endpoint'] = encrypt(message.get_args())
 
     await db[DBTables.config].write()
 

bot/handlers/image_info/image_info.py

@@ -2,6 +2,7 @@ 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.keyboards.image_info import get_img_info_keyboard
 from bot.utils.trace_exception import PrettyException
 
 
@@ -12,14 +13,15 @@ async def imginfo(message: types.Message):
             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)):
+        if not 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))
-        # TODO: Pretty print this
+        await message.reply("Image was generated using this bot", reply_markup=get_img_info_keyboard(
+            message.reply_to_message.photo[0].file_unique_id
+        ))
 
     except IndexError:
         await message.reply('❌ Reply with this command on PICTURE', parse_mode='html')

bot/handlers/txt2img/set_settings.py

@@ -0,0 +1,34 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.utils.cooldown import throttle
+from bot.modules.api.objects.prompt_request import Prompt
+from bot.keyboards.exception import get_exception_keyboard
+from bot.utils.trace_exception import PrettyException
+
+
+@throttle(cooldown=5, admin_ids=db[DBTables.config].get('admins'))
+async def set_prompt_command(message: types.Message):
+    temp_message = await message.reply("⏳ Setting prompt...")
+    if not message.get_args():
+        await temp_message.edit_text("😶‍🌫️ Specify prompt for this command. Check /help setprompt")
+        return
+
+    try:
+        prompt: Prompt = db[DBTables.prompts].get(message.from_id, Prompt(message.get_args()))
+        prompt.prompt = message.get_args()
+        prompt.creator = message.from_id
+        db[DBTables.prompts][message.from_id] = prompt
+
+        await db[DBTables.config].write()
+
+        await message.reply('✅ Default prompt set')
+        await temp_message.delete()
+
+    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))
+        await temp_message.delete()
+        db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1
+        return

bot/handlers/txt2img/txt2img.py

@@ -12,9 +12,13 @@ from aiohttp import ClientConnectorError
 @throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
 async def txt2img_comand(message: types.Message):
     temp_message = await message.reply("⏳ Enqueued...")
-    if not message.get_args():
-        await temp_message.edit_text("😶‍🌫️ Specify prompt for this command. Check /help txt2img")
-        return
+
+    prompt: Prompt = db[DBTables.prompts].get(message.from_id)
+    if not prompt:
+        if message.get_args():
+            db[DBTables.prompts][message.from_id] = Prompt(message.get_args(), creator=message.from_id)
+
+    # TODO: Move it to other module
 
     try:
         db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 0) + 1