Add whitelist to bot

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

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

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

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

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

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,