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,