From fb57b11feb90338c96c0aebe79f2098d5f4f7aad Mon Sep 17 00:00:00 2001
From: BarsTiger <kit.barsic@gmail.com>
Date: Sat, 11 Mar 2023 23:59:05 +0200
Subject: [PATCH] Disabling generation function

---
 bot/callbacks/config/admin_settings.py |  1 +
 bot/handlers/admin/__init__.py         |  1 +
 bot/handlers/admin/on_off.py           | 17 +++++++++++++++++
 bot/handlers/initialize/start.py       |  3 +++
 bot/handlers/txt2img/txt2img.py        |  4 ++++
 bot/keyboards/config.py                |  2 ++
 6 files changed, 28 insertions(+)
 create mode 100644 bot/handlers/admin/on_off.py

diff --git a/bot/callbacks/config/admin_settings.py b/bot/callbacks/config/admin_settings.py
index c8f3b53..dca05ed 100644
--- a/bot/callbacks/config/admin_settings.py
+++ b/bot/callbacks/config/admin_settings.py
@@ -31,6 +31,7 @@ async def on_admin_settings_set(call: types.CallbackQuery, callback_data: dict):
         if 'aliases' in overload and 'admin' in overload and ('add' in overload or 'remove' in overload)
         else f"⚒️ Type new endpoint address: " if "aliases.set_endpoint" in overload
         else f"⚒️ Type \"reset\" if you REALLY want to reset queue: " if "reset.resetqueue" in overload
+        else f"⚒️ Type \"on\" or \"off\" to change generation mode: " if "on_off" in overload
         else f"❌ Not found...",
         reply_markup=types.InlineKeyboardMarkup().add(types.InlineKeyboardButton(
             "👈 Back",
diff --git a/bot/handlers/admin/__init__.py b/bot/handlers/admin/__init__.py
index 9291db7..c840d8a 100644
--- a/bot/handlers/admin/__init__.py
+++ b/bot/handlers/admin/__init__.py
@@ -2,6 +2,7 @@ from bot.common import dp
 from .aliases import *
 from .reset import *
 from .tools import *
+from .on_off import *
 
 
 def register():
diff --git a/bot/handlers/admin/on_off.py b/bot/handlers/admin/on_off.py
new file mode 100644
index 0000000..714e058
--- /dev/null
+++ b/bot/handlers/admin/on_off.py
@@ -0,0 +1,17 @@
+from aiogram import types
+from bot.db import db, DBTables
+from bot.config import ADMIN
+from bot.utils.cooldown import throttle
+
+
+@throttle(5)
+async def on_off_call(message: types.Message, is_command=None):
+    if message.from_id not in db[DBTables.config].get('admins') and message.from_id != ADMIN:
+        await message.reply('❌ You are not permitted to do that. '
+                            'It is only for this bot instance maintainers and admins')
+        return
+
+    db[DBTables.config]['enabled'] = False if message.text.lower() == 'off' else True
+    await db[DBTables.config].write()
+
+    await message.reply(f"Generation enabled: {'💚 (yes)' if db[DBTables.config]['enabled'] else '💔 (no)'}")
diff --git a/bot/handlers/initialize/start.py b/bot/handlers/initialize/start.py
index 37c0b16..6780159 100644
--- a/bot/handlers/initialize/start.py
+++ b/bot/handlers/initialize/start.py
@@ -18,6 +18,9 @@ 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 db[DBTables.config].get('enabled') is None:
+            db[DBTables.config]['enabled'] = True
+            await message.reply(f'✅ Generation is enabled now')
         return
 
     await message.reply(f'👋 Hello, {message.from_user.username}. Use /help to see available commands.')
diff --git a/bot/handlers/txt2img/txt2img.py b/bot/handlers/txt2img/txt2img.py
index e44a6c0..c164474 100644
--- a/bot/handlers/txt2img/txt2img.py
+++ b/bot/handlers/txt2img/txt2img.py
@@ -14,6 +14,10 @@ from bot.utils.errorable_command import wrap_exception
 @throttle(cooldown=30, admin_ids=db[DBTables.config].get('admins'))
 async def generate_command(message: types.Message):
     temp_message = await message.reply("⏳ Enqueued...")
+    if not db[DBTables.config]['enabled']:
+        await message.reply('💔 Generation is disabled by admins now. Try again later')
+        await temp_message.delete()
+        return
 
     try:
         prompt = get_prompt(user_id=message.from_id,
diff --git a/bot/keyboards/config.py b/bot/keyboards/config.py
index 2504d82..e8fc549 100644
--- a/bot/keyboards/config.py
+++ b/bot/keyboards/config.py
@@ -57,6 +57,8 @@ def get_admin_settings_keyboard() -> types.InlineKeyboardMarkup:
         types.InlineKeyboardButton("Remove admin", callback_data=admin_settings_data.new("aliases.remove_admin")),
         types.InlineKeyboardButton("Set API endpoint", callback_data=admin_settings_data.new("aliases.set_endpoint")),
         types.InlineKeyboardButton("Reset generation queue", callback_data=admin_settings_data.new("reset.resetqueue")),
+        types.InlineKeyboardButton("Turn on/off generation",
+                                   callback_data=admin_settings_data.new("on_off.on_off_call")),
         types.InlineKeyboardButton("👈 Back", callback_data="config_back"),
         types.InlineKeyboardButton("🔻 Close", callback_data="close_keyboard")
     ]