From 8e3e6929d806c5620ccf3848873a967c2f49775c Mon Sep 17 00:00:00 2001
From: BarsTiger <kit.barsic@gmail.com>
Date: Thu, 2 Mar 2023 20:23:01 +0200
Subject: [PATCH] Adding and removing admins

---
 bot/handlers/admin/__init__.py            |  2 +
 bot/handlers/admin/aliases.py             | 60 +++++++++++++++++++++++
 bot/handlers/help_command/help_strings.py |  4 +-
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/bot/handlers/admin/__init__.py b/bot/handlers/admin/__init__.py
index 42fb38a..9291db7 100644
--- a/bot/handlers/admin/__init__.py
+++ b/bot/handlers/admin/__init__.py
@@ -7,4 +7,6 @@ from .tools import *
 def register():
     dp.register_message_handler(set_endpoint, commands='setendpoint')
     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(tools.hash_command, commands='hash')
diff --git a/bot/handlers/admin/aliases.py b/bot/handlers/admin/aliases.py
index aee96ed..c44180c 100644
--- a/bot/handlers/admin/aliases.py
+++ b/bot/handlers/admin/aliases.py
@@ -21,3 +21,63 @@ async def set_endpoint(message: types.Message):
     await db[DBTables.config].write()
 
     await message.reply("✅ New url set")
+
+
+@throttle(5)
+async def add_admin(message: types.Message):
+    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().isdecimal() and not hasattr(message.reply_to_message, 'text'):
+        await message.reply('❌ Put new admin ID to command arguments or answer to users message')
+        return
+    elif not message.get_args().isdecimal():
+        ID = message.reply_to_message.from_id
+    elif not hasattr(message.reply_to_message, 'text'):
+        ID = int(message.get_args())
+
+    if not isinstance(db[DBTables.config].get('admins'), list):
+        db[DBTables.config]['admins'] = list()
+
+    if ID not in db[DBTables.config].get('admins'):
+        admins_ = db[DBTables.config].get('admins')
+        admins_.append(ID)
+        db[DBTables.config]['admins'] = admins_
+    else:
+        await message.reply('❌ This admin is added already')
+        return
+
+    await db[DBTables.config].write()
+
+    await message.reply("✅ Added admin")
+
+
+@throttle(5)
+async def remove_admin(message: types.Message):
+    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().isdecimal() and not hasattr(message.reply_to_message, 'text'):
+        await message.reply('❌ Put admin ID to command arguments or answer to users message')
+        return
+    elif not message.get_args().isdecimal():
+        ID = message.reply_to_message.from_id
+    elif not hasattr(message.reply_to_message, 'text'):
+        ID = int(message.get_args())
+
+    if not isinstance(db[DBTables.config].get('admins'), list):
+        db[DBTables.config]['admins'] = list()
+
+    if ID not in db[DBTables.config].get('admins'):
+        await message.reply('❌ This admin is not added')
+        return
+    else:
+        admins_ = db[DBTables.config].get('admins')
+        admins_.remove(ID)
+        db[DBTables.config]['admins'] = admins_
+
+    await db[DBTables.config].write()
+
+    await message.reply("✅ Removed admin")
diff --git a/bot/handlers/help_command/help_strings.py b/bot/handlers/help_command/help_strings.py
index d8b8962..fb7f743 100644
--- a/bot/handlers/help_command/help_strings.py
+++ b/bot/handlers/help_command/help_strings.py
@@ -11,5 +11,7 @@ help_data = {
     'setsampler': 'Set StableDiffusion sampler',
     'setscale': 'Set CFG Scale (prompt stringency)',
     'setfaces': 'Set restore faces mode',
-    'setendpoint': '(admin) Set StableDiffusion API endpoint'
+    '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'
 }