From c129c7bc60103b742846f4a649e389256552dd31 Mon Sep 17 00:00:00 2001 From: seel2304 <37144504+seel2304@users.noreply.github.com> Date: Wed, 3 May 2023 10:36:12 +0000 Subject: [PATCH 1/7] Disable cooldown for admins --- bot/utils/cooldown.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bot/utils/cooldown.py b/bot/utils/cooldown.py index 47bb29b..d997cf5 100644 --- a/bot/utils/cooldown.py +++ b/bot/utils/cooldown.py @@ -30,7 +30,7 @@ def throttle(cooldown: int = 5, by_id: bool = True, admin_ids: list = None): if not last_time: last_time = delta - if last_time <= delta: + if last_time <= delta or user_id in admin_ids: try: f_name_dict = db[DBTables.cooldown][func.__name__] f_name_dict[user_id] = now @@ -42,7 +42,7 @@ def throttle(cooldown: int = 5, by_id: bool = True, admin_ids: list = None): try: return asyncio.create_task(func(*args, **kwargs)) except Exception as e: - assert e + assert e else: return not_allowed(*args, cooldown, by_id) From 2dda780d3f558e46540b4d2ed2c772f9a8988065 Mon Sep 17 00:00:00 2001 From: seel2304 <37144504+seel2304@users.noreply.github.com> Date: Wed, 3 May 2023 15:50:15 +0000 Subject: [PATCH 2/7] Add whitelist to bot --- bot/handlers/admin/__init__.py | 2 + bot/handlers/admin/aliases.py | 69 +++++++++++++++++++++++ bot/handlers/help_command/help_strings.py | 4 +- bot/handlers/initialize/start.py | 7 +++ bot/handlers/register.py | 2 +- bot/handlers/txt2img/txt2img.py | 4 ++ 6 files changed, 86 insertions(+), 2 deletions(-) diff --git a/bot/handlers/admin/__init__.py b/bot/handlers/admin/__init__.py index c840d8a..fb1be1f 100644 --- a/bot/handlers/admin/__init__.py +++ b/bot/handlers/admin/__init__.py @@ -10,4 +10,6 @@ def register(): dp.register_message_handler(reset.resetqueue, commands='resetqueue') dp.register_message_handler(aliases.add_admin, commands='addadmin') dp.register_message_handler(aliases.remove_admin, commands='rmadmin') + dp.register_message_handler(aliases.add_whitelist, commands='addwhitelist') + dp.register_message_handler(aliases.remove_whitelist, commands='rmwhitelist') dp.register_message_handler(tools.hash_command, commands='hash') diff --git a/bot/handlers/admin/aliases.py b/bot/handlers/admin/aliases.py index b6d47aa..5521731 100644 --- a/bot/handlers/admin/aliases.py +++ b/bot/handlers/admin/aliases.py @@ -23,6 +23,75 @@ async def set_endpoint(message: types.Message, is_command: bool = True): await message.reply("✅ New url set") +@throttle(5) +async def add_whitelist(message: types.Message, is_command: bool = True): + if message.from_id != ADMIN: + await message.reply('❌ You are not permitted to do that. It is only for main admin') + return + + if not (message.get_args() if is_command else message.text).isdecimal() and not \ + hasattr(message.reply_to_message, 'text') and (message.chat.id >= 0): + await message.reply('❌ Put new whitelist chat ID to command arguments') + return + elif not (message.get_args() if is_command else message.text).isdecimal() and not \ + hasattr(message.reply_to_message, 'text') and (message.chat.id < 0): + ID = message.chat.id + await message.reply(f'Chat ID: {message.chat.id} Chat title: {message.chat.title}') + elif not (message.get_args() if is_command else message.text).isdecimal(): + ID = message.reply_to_message.from_id + elif not hasattr(message.reply_to_message, 'text'): + ID = int((message.get_args() if is_command else message.text)) + + if not isinstance(db[DBTables.config].get('whitelist'), list): + db[DBTables.config]['whitelist'] = list() + + if ID not in db[DBTables.config].get('whitelist'): + whitelist_ = db[DBTables.config].get('whitelist') + whitelist_.append(ID) + db[DBTables.config]['whitelist'] = whitelist_ + else: + await message.reply('❌ This whitelist is added already') + return + + await db[DBTables.config].write() + + await message.reply("✅ Added whitelist") + +@throttle(5) +async def remove_whitelist(message: types.Message, is_command: bool = True): + if message.from_id != ADMIN: + await message.reply('❌ You are not permitted to do that. It is only for main admin') + return + + if not (message.get_args() if is_command else message.text).isdecimal() and not \ + hasattr(message.reply_to_message, 'text') and (message.chat.id >= 0: + await message.reply('❌ Put whitelist ID to command arguments or answer to users message') + return + elif not (message.get_args() if is_command else message.text).isdecimal() and not \ + hasattr(message.reply_to_message, 'text') and (message.chat.id < 0): + ID = message.chat.id + await message.reply(f'Chat ID: {message.chat.id} Chat title: {message.chat.title}') + elif not (message.get_args() if is_command else message.text).isdecimal(): + ID = message.reply_to_message.from_id + elif not hasattr(message.reply_to_message, 'text'): + ID = int((message.get_args() if is_command else message.text)) + + if not isinstance(db[DBTables.config].get('whitelist'), list): + db[DBTables.config]['whitelist'] = list() + + if ID not in db[DBTables.config].get('whitelist'): + await message.reply('❌ This whitelist is not added') + return + else: + whitelist_ = db[DBTables.config].get('whitelist') + whitelist_.remove(ID) + db[DBTables.config]['whitelist'] = whitelist_ + + await db[DBTables.config].write() + + await message.reply("✅ Removed whitelist") + + @throttle(5) async def add_admin(message: types.Message, is_command: bool = True): diff --git a/bot/handlers/help_command/help_strings.py b/bot/handlers/help_command/help_strings.py index 686cf72..bb795da 100644 --- a/bot/handlers/help_command/help_strings.py +++ b/bot/handlers/help_command/help_strings.py @@ -17,5 +17,7 @@ help_data = { 'setmodel': '(global) Sets StableDiffusion model for all users. Can be used only once an hour', 'setendpoint': '(admin) Set StableDiffusion API endpoint', 'addadmin': '(admin) Add new admin - reply to message or type user ID', - 'rmadmin': '(admin) Remove admin - reply to message or type user ID' + 'rmadmin': '(admin) Remove admin - reply to message or type user ID', + 'addwhitelist': '(admin) Add new whitelist - reply to message or type user ID', + 'rmwhitelist': '(admin) Remove whitelist - reply to message or type user ID' } diff --git a/bot/handlers/initialize/start.py b/bot/handlers/initialize/start.py index 6780159..e832a48 100644 --- a/bot/handlers/initialize/start.py +++ b/bot/handlers/initialize/start.py @@ -18,6 +18,13 @@ async def start_command(message: types.Message): await db[DBTables.config].write() await message.reply(f'✅ Added {message.from_user.username} to admins. You can add other admins, ' f'check bot settings menu') + if ADMIN not in db[DBTables.config].get('whitelist'): + whitelist_ = db[DBTables.config].get('whitelist') + whitelist_.append(ADMIN) + db[DBTables.config]['whitelist'] = whitelist_ + await db[DBTables.config].write() + await message.reply(f'✅ Added {message.from_user.username} to whitelist. You can add other users to whitelist, ' + f'check bot settings menu') if db[DBTables.config].get('enabled') is None: db[DBTables.config]['enabled'] = True await message.reply(f'✅ Generation is enabled now') diff --git a/bot/handlers/register.py b/bot/handlers/register.py index de3a159..fb82a95 100644 --- a/bot/handlers/register.py +++ b/bot/handlers/register.py @@ -3,7 +3,7 @@ from rich import print def register_handlers(): from bot.handlers import ( - initialize, admin, help_command, txt2img, image_info, config + initialize, admin, help_command, txt2img, image_info, config, whitelist ) initialize.register() diff --git a/bot/handlers/txt2img/txt2img.py b/bot/handlers/txt2img/txt2img.py index c164474..42ebb15 100644 --- a/bot/handlers/txt2img/txt2img.py +++ b/bot/handlers/txt2img/txt2img.py @@ -18,6 +18,10 @@ async def generate_command(message: types.Message): await message.reply('💔 Generation is disabled by admins now. Try again later') await temp_message.delete() return + elif (message.chat.id not in db[DBTables.config]['whitelist'] and message.from_id not in db[DBTables.config]['whitelist']): + await message.reply('❌You are not on the white list, access denied. Contact admin @kilisauros for details') + await temp_message.delete() + return try: prompt = get_prompt(user_id=message.from_id, From 75f6404e94b75c35fc461d622ac0c64bb2ce361d Mon Sep 17 00:00:00 2001 From: seel2304 <37144504+seel2304@users.noreply.github.com> Date: Wed, 3 May 2023 17:09:25 +0000 Subject: [PATCH 3/7] Fixed changing models. Changed cooldown for changing models to 30m --- bot/handlers/txt2img/set_model.py | 4 ++-- bot/modules/api/models.py | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/bot/handlers/txt2img/set_model.py b/bot/handlers/txt2img/set_model.py index ddb31e3..75248fe 100644 --- a/bot/handlers/txt2img/set_model.py +++ b/bot/handlers/txt2img/set_model.py @@ -1,5 +1,5 @@ from aiogram import types -from bot.db import db, DBTables +from bot.db import db, DBTables, decrypt from bot.utils.cooldown import throttle from bot.keyboards.set_model import get_set_model_keyboard from bot.modules.api.models import get_models @@ -7,7 +7,7 @@ from bot.utils.errorable_command import wrap_exception @wrap_exception() -@throttle(cooldown=60*60, admin_ids=db[DBTables.config].get('admins'), by_id=False) +@throttle(cooldown=30*60, admin_ids=db[DBTables.config].get('admins'), by_id=False) async def set_model_command(message: types.Message): models = await get_models() if models is not None and len(models) > 0: diff --git a/bot/modules/api/models.py b/bot/modules/api/models.py index 77fba92..9000950 100644 --- a/bot/modules/api/models.py +++ b/bot/modules/api/models.py @@ -1,14 +1,16 @@ import aiohttp +import json from bot.db import db, DBTables, decrypt async def get_models() -> list | None: endpoint = decrypt(db[DBTables.config].get('endpoint')) async with aiohttp.ClientSession() as session: - r = await session.get(endpoint + "/sdapi/v1/sd-models") + r = await session.get(endpoint + '/sdapi/v1/sd-models') if r.status != 200: return None - return [x["title"] for x in await r.json()] + json_data = await r.json() + return [x["title"] for x in json_data] async def set_model(model_name: str): From 8ec4ddcc0aab815711d600d9e99f99f20820be40 Mon Sep 17 00:00:00 2001 From: vasmarfas <37144504+vasmarfas@users.noreply.github.com> Date: Wed, 3 May 2023 21:35:23 +0300 Subject: [PATCH 4/7] Update set_model.py --- bot/handlers/txt2img/set_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bot/handlers/txt2img/set_model.py b/bot/handlers/txt2img/set_model.py index 75248fe..7d55562 100644 --- a/bot/handlers/txt2img/set_model.py +++ b/bot/handlers/txt2img/set_model.py @@ -15,5 +15,5 @@ async def set_model_command(message: types.Message): else: await message.reply('❌ No models available') return - + await message.reply("Examples of models (with additional info): https://telegra.ph/Opisanie-raboty-modelej-05-03") await message.reply("🪄 You can choose model from available:", reply_markup=get_set_model_keyboard(0)) From 1653e184a81ec8ea99aa49344baef8ada00bfc45 Mon Sep 17 00:00:00 2001 From: vasmarfas <37144504+vasmarfas@users.noreply.github.com> Date: Wed, 3 May 2023 21:39:16 +0300 Subject: [PATCH 5/7] Update set_model.py Whitelisted set_model Chanched cooldown to 5 minutes --- bot/handlers/txt2img/set_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/bot/handlers/txt2img/set_model.py b/bot/handlers/txt2img/set_model.py index 7d55562..dcc5634 100644 --- a/bot/handlers/txt2img/set_model.py +++ b/bot/handlers/txt2img/set_model.py @@ -7,8 +7,11 @@ from bot.utils.errorable_command import wrap_exception @wrap_exception() -@throttle(cooldown=30*60, admin_ids=db[DBTables.config].get('admins'), by_id=False) +@throttle(cooldown=5*60, admin_ids=db[DBTables.config].get('admins'), by_id=False) async def set_model_command(message: types.Message): + if (message.chat.id not in db[DBTables.config]['whitelist'] and message.from_id not in db[DBTables.config]['whitelist']): + await message.reply('❌You are not on the white list, access denied. Contact admin @kilisauros for details') + return models = await get_models() if models is not None and len(models) > 0: db[DBTables.config]['models'] = models From 0bc37fc87f28a5fe8e15dbd25320b8f69ce6fce7 Mon Sep 17 00:00:00 2001 From: seel2304 Date: Fri, 5 May 2023 02:10:45 +0300 Subject: [PATCH 6/7] Add TG Chat DB (KB not working currently) --- bot/handlers/txt2img/txt2img.py | 13 +++++++++++++ bot/utils/private_keyboard.py | 7 +++++++ 2 files changed, 20 insertions(+) diff --git a/bot/handlers/txt2img/txt2img.py b/bot/handlers/txt2img/txt2img.py index 42ebb15..458e9b2 100644 --- a/bot/handlers/txt2img/txt2img.py +++ b/bot/handlers/txt2img/txt2img.py @@ -1,4 +1,5 @@ import re +from bot.common import bot from aiogram import types from bot.db import db, DBTables from bot.utils.cooldown import throttle @@ -8,6 +9,8 @@ from bot.modules.api.objects.prompt_request import Generated from bot.modules.api.status import wait_for_status from bot.keyboards.image_info import get_img_info_keyboard from bot.utils.errorable_command import wrap_exception +from bot.callbacks.factories.image_info import (prompt_only, full_prompt, import_prompt, back) + @wrap_exception([ValueError], custom_loading=True) @@ -43,6 +46,16 @@ async def generate_command(message: types.Message): db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1 image = await txt2img(prompt) image_message = await message.reply_photo(photo=image[0]) + + #Send photo to SD Image Archive + # message_data = + + archive_message = f'User ID: {message.from_id} \n \ + User nickname: {message.from_user.full_name} \n \ + User username: @{message.from_user.username}' + # archive_message = f'{callback_data=full_prompt.new(p_id=f"{image_message.photo[0].file_unique_id}")} + "User ID: " + message.from_id + " User nickname: " + message.from_user.full_name + " User username: " + message.from_user.username' + await bot.send_photo(-929754401, photo=image[0], caption=archive_message, reply_markup=get_img_info_keyboard(image_message.photo[0].file_unique_id)) + db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated( prompt=prompt, diff --git a/bot/utils/private_keyboard.py b/bot/utils/private_keyboard.py index 4a69b35..ce0cf62 100644 --- a/bot/utils/private_keyboard.py +++ b/bot/utils/private_keyboard.py @@ -1,10 +1,17 @@ from aiogram import types +from bot.config import ADMIN async def other_user(call: types.CallbackQuery) -> bool: if not hasattr(call.message.reply_to_message, 'from_id'): await call.answer('Error, original call was removed', show_alert=True) return True + elif call.from_user.id == ADMIN: + print(f"call.from_user.id: {call.from_user.id}") + print(f"ADMIN: {ADMIN}") + print("User is Admin") + return False + elif call.message.reply_to_message.from_id != call.from_user.id: await call.answer('It is not your menu!', show_alert=True) return True From 18678b8d03b14bc13064c133383034394dbbb8fd Mon Sep 17 00:00:00 2001 From: seel2304 Date: Fri, 5 May 2023 03:15:21 +0300 Subject: [PATCH 7/7] Fixed TG Image DB --- bot/handlers/txt2img/txt2img.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/bot/handlers/txt2img/txt2img.py b/bot/handlers/txt2img/txt2img.py index 458e9b2..1da8ee2 100644 --- a/bot/handlers/txt2img/txt2img.py +++ b/bot/handlers/txt2img/txt2img.py @@ -47,15 +47,23 @@ async def generate_command(message: types.Message): image = await txt2img(prompt) image_message = await message.reply_photo(photo=image[0]) - #Send photo to SD Image Archive - # message_data = + #Send photo to SD Image Archive archive_message = f'User ID: {message.from_id} \n \ User nickname: {message.from_user.full_name} \n \ - User username: @{message.from_user.username}' - # archive_message = f'{callback_data=full_prompt.new(p_id=f"{image_message.photo[0].file_unique_id}")} + "User ID: " + message.from_id + " User nickname: " + message.from_user.full_name + " User username: " + message.from_user.username' - await bot.send_photo(-929754401, photo=image[0], caption=archive_message, reply_markup=get_img_info_keyboard(image_message.photo[0].file_unique_id)) - + User username: @{message.from_user.username} \n \ + Chat ID: {message.chat.id} \n \ + Chat title: {message.chat.title} \n \ + Info: \n \ + 🖤 Prompt: {prompt.prompt} \n \ + 🐊 Negative: {prompt.negative_prompt} \n \ + 💫 Model: In development \n \ + 🪜 Steps: {prompt.steps} \n \ + 🧑‍🎨 CFG Scale: {prompt.cfg_scale} \n \ + 🖥️ Size: {prompt.width}x{prompt.height} \n \ + 😀 Restore faces: {prompt.restore_faces} \n \ + ⚒️ Sampler: {prompt.sampler} \n ' + await bot.send_photo(-929754401, photo=image[0], caption=archive_message) db[DBTables.generated][image_message.photo[0].file_unique_id] = Generated( prompt=prompt, @@ -74,4 +82,4 @@ async def generate_command(message: types.Message): await message.reply(f'❌ Error! {e.args[0]}') await temp_message.delete() db[DBTables.queue]['n'] = db[DBTables.queue].get('n', 1) - 1 - return + return \ No newline at end of file