diff --git a/bot/handlers/admin/__init__.py b/bot/handlers/admin/__init__.py index fb1be1f..a8e9ba9 100644 --- a/bot/handlers/admin/__init__.py +++ b/bot/handlers/admin/__init__.py @@ -12,4 +12,5 @@ def register(): 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(aliases.get_whitelist, commands='getwhitelist') dp.register_message_handler(tools.hash_command, commands='hash') diff --git a/bot/handlers/admin/aliases.py b/bot/handlers/admin/aliases.py index 5521731..7a6d9bc 100644 --- a/bot/handlers/admin/aliases.py +++ b/bot/handlers/admin/aliases.py @@ -23,6 +23,7 @@ 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: @@ -57,6 +58,7 @@ async def add_whitelist(message: types.Message, is_command: bool = True): await message.reply("✅ Added whitelist") + @throttle(5) async def remove_whitelist(message: types.Message, is_command: bool = True): if message.from_id != ADMIN: @@ -64,7 +66,7 @@ async def remove_whitelist(message: types.Message, is_command: bool = True): 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: + 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 \ @@ -92,6 +94,16 @@ async def remove_whitelist(message: types.Message, is_command: bool = True): await message.reply("✅ Removed whitelist") +@throttle(5) +async def get_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 + + await message.reply(f"✅ Whitelisted ids: {db[DBTables.config].get('whitelist')}" + if db[DBTables.config].get('whitelist') else + '❌ Whitelist is disabled. Everyone can use the bot. Add someone to whitelist to enable it') + @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 bb795da..6169eb4 100644 --- a/bot/handlers/help_command/help_strings.py +++ b/bot/handlers/help_command/help_strings.py @@ -18,6 +18,7 @@ help_data = { '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', - 'addwhitelist': '(admin) Add new whitelist - reply to message or type user ID', - 'rmwhitelist': '(admin) Remove whitelist - reply to message or type user ID' + 'addwhitelist': '(admin) Add to whitelist - reply to message or type user ID', + 'rmwhitelist': '(admin) Remove from whitelist - reply to message or type user ID', + 'getwhitelist': '(admin) Get current whitelist and check if it is enabled' }