add admin logic

2025-02-23 00:37:04 +03:00
parent 8c191ca139
commit be24e124bb
9 changed files with 389 additions and 26 deletions
@@ -1,4 +1,5 @@
 import os
+from datetime import datetime
 from sshtunnel import SSHTunnelForwarder
 from dotenv import load_dotenv

@@ -6,8 +7,8 @@ load_dotenv()


 class config:
-    BOT_TOKEN = os.environ.get("TOKEN")
-    API_KEY = os.environ.get("API_KEY")
+    DATE_START = datetime.strptime(os.environ.get("DATE_START"), "%Y-%m-%d").date()
+    BOT_TOKEN = os.environ.get("API_TOKEN")
    ADMIN_ID = os.environ.get("ADMIN_ID")
    ENV = os.getenv("ENV", "PROD")

@@ -25,16 +26,12 @@ class config:
        DB_PORT = server.local_bind_port


-class BotCommands:
-    Start = ["start", "s"]
-    Help = ["help", "h"]
-
-
 class Messages:
    WELCOME_LIST = [
        "Привет! Это бот  Quick dates (Random Coffee). Чтобы зарегистрироваться, пожалуйста, ответь на следующие вопросы."  # несколько сообщений подряд
    ]
-    ERROR = "An error occurred"
+    HELP = "Помощь"
+    ERROR = "Произошла ошибка"
    ENTER_NAME = "Как тебя зовут?"
    ENTER_AGE = "Сколько тебе лет?"
    ENTER_GENDER = "Выбери пол (женский, мужской)"
@@ -20,3 +20,6 @@ class DB:

    def exist(self, query):
        return self.collection.find_one(query) is not None
+
+    def delete(self, query):
+        self.collection.delete_one(query)
@@ -2,11 +2,11 @@ import os
 import sys
 import telebot
 import threading
-from config import config, BotCommands, Messages
+from config import config, Messages
 from telebot import types
 from plugin.user import User
-from plugin.admin import Admin
-from plugin.bot_instance import bot
+from plugin.admin import Admin, AdminMessages
+from plugin.bot_instance import bot, BotCommands
 from plugin.register import start_reg_name
 import time

@@ -14,8 +14,13 @@ import time
 admin_id = config.ADMIN_ID


-@bot.message_handler(commands=["start"])
+@bot.message_handler(commands=BotCommands.Start)
 def start_message(message: types.Message):
+    if Admin.is_admin(message.chat.id):
+        help_message(message)
+
+        return
+
    if User(message.chat.id).name is not None:
        user = User(message.chat.id)
        markup = types.InlineKeyboardMarkup()
@@ -54,6 +59,20 @@ def start_message(message: types.Message):
    return


+@bot.message_handler(commands=BotCommands.Help)
+def help_message(message: types.Message):
+    bot.send_message(
+        message.chat.id,
+        (
+            Messages.HELP
+            if not Admin.is_admin(message.chat.id)
+            else AdminMessages.ADMIN_HELP
+        ),
+    )
+
+    return
+
+
 if __name__ == "__main__":
    print("/t--- Bot started ---")

@@ -1,21 +1,282 @@
 from db import DB
+import config
+from plugin.user import User
+from plugin.shuffle import Shuffle
+
+from plugin.bot_instance import bot
+
+"""
+логика: есть первый админ, который получается из config.ADMIN_ID все остальные админы добавляются только если есть свободная запись, сгенерированная другим админом - уникальный ключ. У админов есть свои команды для telebot, которые могут запускать только они
+"""


 class Admin:
    _db = DB("admins")

-    def __init__(self, user_id, name):
-        self.user_id = user_id
-        self.name = name
-        self._save()
-
    def __init__(self, user_id):
        self.user_id = user_id
        self._load()

-    def _save(self):
-        self._db.insert({"user_id": self.user_id, "name": self.name})
+        return

    def _load(self):
        data = self._db.find_one({"user_id": self.user_id})
-        self.name = data["name"]
+        if data:
+            self.name = data["name"]
+            self.added_by = data["added_by"]
+        else:
+            raise ValueError("Admin not found")
+
+        return
+
+    @staticmethod
+    def is_admin(user_id):
+        return Admin._db.find_one({"user_id": user_id}) is not None
+
+    @staticmethod
+    def add_admin(new_user_id, name, unique_key) -> int:
+        key_data = Admin._db.find_one({"unique_key": unique_key})
+
+        if new_user_id == config.ADMIN_ID and not Admin.is_admin(new_user_id):
+            Admin._db.insert({"user_id": new_user_id, "name": name})
+
+            return new_user_id
+
+        if not key_data:
+            raise ValueError("Unique key not found")
+
+        if not Admin.is_admin(new_user_id):
+            Admin._db.insert(
+                {
+                    "user_id": new_user_id,
+                    "name": name,
+                    "added_by": key_data["generated_by"],
+                }
+            )
+
+            if key_data:
+                Admin._db.delete({"unique_key": unique_key})
+
+                return key_data["generated_by"]
+        else:
+            raise ValueError("Admin already exists")
+
+        return
+
+    @staticmethod
+    def remove_admin(current_user_id, target_user_id):
+        if target_user_id == config.ADMIN_ID:
+            raise PermissionError("Cannot remove the first admin")
+
+        if Admin.is_admin(current_user_id):
+            if Admin.is_admin(target_user_id):
+                Admin._db.delete({"user_id": target_user_id})
+            else:
+                raise ValueError("Admin not found")
+        else:
+            raise PermissionError("No permission to remove admin")
+
+        return
+
+    @staticmethod
+    def generate_unique_key(admin_user_id):
+        if Admin.is_admin(admin_user_id):
+            import uuid
+
+            unique_key = str(uuid.uuid4())
+            while Admin._db.find_one({"unique_key": unique_key}):
+                unique_key = str(uuid.uuid4())
+
+            Admin._db.insert({"unique_key": unique_key, "generated_by": admin_user_id})
+
+            return unique_key
+        else:
+            raise PermissionError("Only admins can generate unique keys")
+
+        return
+
+    # ----------------- TELEBOT ----------------- #
+
+    def admin_only(func):
+        def wrapper(message, *args, **kwargs):
+            if Admin.is_admin(message.chat.id):
+                return func(message, *args, **kwargs)
+            else:
+                bot.send_message(
+                    message.chat.id, "У вас нет прав для выполнения этой команды."
+                )
+
+        return wrapper
+
+    @bot.message_handler(commands=["admin"])
+    def add_admin_command(message):
+        bot.send_message(message.chat.id, "Введите unique_key")
+        bot.register_next_step_handler(message, Admin.add_admin_step)
+
+        return
+
+    def add_admin_step(message):
+        try:
+            admin_id = Admin.add_admin(
+                message.chat.id, message.from_user.username, message.text
+            )
+            bot.send_message(message.chat.id, "Вы добавлены в админы!")
+            bot.send_message(
+                admin_id,
+                "По вашему ключу {} был добавлен новый админ {}".format(
+                    message.text, message.from_user.username
+                ),
+            )
+        except Exception as e:
+            bot.send_message(message.chat.id, str(e))
+
+        return
+
+    @bot.message_handler(commands=["generate_key"])
+    @admin_only
+    def generate_key_command(message):
+        try:
+            key = Admin.generate_unique_key(message.chat.id)
+            bot.send_message(
+                message.chat.id,
+                "Ваш ключ: <code>{}</code>".format(key),
+                parse_mode="HTML",
+            )
+        except Exception as e:
+            bot.send_message(message.chat.id, str(e))
+
+        return
+
+    @bot.message_handler(commands=["remove_admin"])
+    @admin_only
+    def remove_admin_command(message):
+        bot.send_message(message.chat.id, "Введите user_id")
+        bot.register_next_step_handler(message, Admin.remove_admin_step)
+
+        return
+
+    def remove_admin_step(message):
+        try:
+            Admin.remove_admin(message.chat.id, int(message.text))
+            bot.send_message(message.chat.id, "Админ удален!")
+        except Exception as e:
+            bot.send_message(message.chat.id, str(e))
+
+        return
+
+    @bot.message_handler(commands=["delete_user"])
+    @admin_only
+    def delete_user_command(message):
+        bot.send_message(message.chat.id, "Введите user_id")
+        bot.register_next_step_handler(message, Admin.delete_user_step)
+
+        return
+
+    def delete_user_step(message):
+        try:
+            User(int(message.text)).delete()
+            bot.send_message(message.chat.id, "Пользователь удален!")
+        except Exception as e:
+            bot.send_message(message.chat.id, str(e))
+
+        return
+
+    @bot.message_handler(commands=["stats"])
+    @admin_only
+    def stats_command(message):
+        bot.send_message(message.chat.id, "Статистика")
+        users = User.get_all()
+
+        with open("stats.txt", "w") as f:
+            f.write("Количество пользователей: {}\n".format(len(users)))
+
+            for user in users:
+                f.write(str(user) + "\n")
+
+        with open("stats.txt", "rb") as f:
+            bot.send_document(message.chat.id, f)
+
+        with open("admin_stats.txt", "w") as f:
+            admins = Admin._db.find({})
+            f.write("Количество админов: {}\n".format(len(admins)))
+
+            for admin in admins:
+                f.write(str(admin) + "\n")
+
+        with open("admin_stats.txt", "rb") as f:
+            bot.send_document(message.chat.id, f)
+
+        import os
+
+        os.remove("stats.txt")
+        os.remove("admin_stats.txt")
+
+        return
+
+    @bot.message_handler(commands=["send_message"])
+    @admin_only
+    def send_message_command(message):
+        bot.send_message(
+            message.chat.id, "Введите сообщение. Отправьте /cancel чтобы отменить"
+        )
+        bot.register_next_step_handler(message, Admin.send_message_step)
+
+        return
+
+    def send_message_step(message):
+        if message.text == "/cancel":
+            bot.send_message(message.chat.id, "Отменено")
+
+            return
+
+        users = User.get_all()
+        for user in users:
+            bot.send_message(user.user_id, message.text)
+
+        bot.send_message(message.chat.id, "Сообщения отправлены!")
+
+        return
+
+    @bot.message_handler(commands=["clear_all"])
+    @admin_only
+    def clear_all_command(message):
+        bot.send_message(
+            message.chat.id,
+            "Отправьте <code>delete all</code> для подтверждения",
+            parse_mode="HTML",
+        )
+
+        return
+
+    def clear_all_command_step(message):
+        if message.text != "delete all":
+            bot.send_message(message.chat.id, "Отменено")
+
+        users = User.get_all()
+        for user in users:
+            user.delete()
+
+        bot.send_message(message.chat.id, "Все пользователи удалены!")
+
+        return
+
+    @bot.message_handler(commands=["random"])
+    @admin_only
+    def random_command(message):
+        users = User.get_all()
+        Shuffle(users)
+
+        return
+
+
+class AdminMessages:
+    ADMIN_HELP = (
+        "Вы администратор! Вот что вы можете сделать:\n"
+        "/remove_admin - удалить админа\n"
+        "/generate_key - сгенерировать уникальный ключ"
+        "/delete_user - удалить пользователя"
+        "/stats - статистика"
+        "/send_message - отправить всем сообщение"
+        "/random - зарандомить людей для 1 тура"
+        "clear_all - удаляет ВСЕХ пользователей"
+    )
@@ -2,6 +2,12 @@ import telebot
 import config
 import threading

+
+class BotCommands:
+    Start = ["start", "s"]
+    Help = ["help", "h"]
+
+
 bot = telebot.TeleBot(
    config.BOT_TOKEN,
    colorful_logs=True,
@@ -1,4 +1,3 @@
-import telebot
 from telebot import types
 from plugin.user import User
 from config import Messages
@@ -19,7 +18,7 @@ def start_reg_age(message: types.Message):
    user = User(message.chat.id)
    user.age = int(message.text)

-    markup = types.InlineKeyboardMarkup(resize_keyboard=True)
+    markup = types.InlineKeyboardMarkup()
    markup.add(
        types.InlineKeyboardButton(Messages.GENDER_WOMAN, callback_data="gender_woman"),
        types.InlineKeyboardButton(Messages.GENDER_MAN, callback_data="gender_man"),
@@ -59,7 +58,7 @@ def start_reg_group(message: types.Message):
    user = User(message.chat.id)
    user.group = message.text

-    markup = types.InlineKeyboardMarkup(resize_keyboard=True)
+    markup = types.InlineKeyboardMarkup()
    markup.add(
        types.InlineKeyboardButton("Всё верно!", callback_data="registration_confirm"),
        types.InlineKeyboardButton("Изменить", callback_data="registration_change"),
@@ -97,6 +96,9 @@ def start_reg_again(call: types.CallbackQuery):
    bot.register_next_step_handler(call.message, start_reg_name)


+# --------------------------------- Test ---------------------------------
+
+
@bot.callback_query_handler(func=lambda call: call.data == "registration_confirm")
 def start_test(call: types.CallbackQuery):
    user = User(call.message.chat.id)
@@ -104,6 +106,65 @@ def start_test(call: types.CallbackQuery):
        call.message.chat.id, call.message.message_id, reply_markup=None
    )

-    bot.send_message(call.message.chat.id, TestMessages.TEST_WELCOME)
+    markup = types.InlineKeyboardMarkup()
+    markup.add(
+        types.InlineKeyboardButton("Начнём!", callback_data="test_yes"),
+    )
+    bot.send_message(
+        call.message.chat.id, TestMessages.TEST_WELCOME, reply_markup=markup
+    )
+    user.type = ""

-    pass
+    return
+
+
+@bot.callback_query_handler(func=lambda call: "test" in call.data)
+def test_question(call: types.CallbackQuery):
+    user = User(call.message.chat.id)
+
+    bot.edit_message_reply_markup(
+        call.message.chat.id, call.message.message_id, reply_markup=None
+    )
+
+    if len(user.type) == 4:
+        show_result(call.message)
+        return
+
+    if call.data != "test_yes":
+        user.type += call.data.split("_")[-1]
+
+        bot.edit_message_text(
+            call.message.text
+            + "\n\n"
+            + TestMessages.TEST_QUESTIONS[len(user.type) - 1]["answers"][
+                int(call.data[-1])
+            ][0],
+            call.message.chat.id,
+            call.message.message_id,
+        )  # упростить
+
+    question = TestMessages.TEST_QUESTIONS[len(user.type)]
+    markup = types.InlineKeyboardMarkup()
+    for answer in question["answers"]:
+        markup.add(
+            types.InlineKeyboardButton(answer[0], callback_data=f"test_{answer[1]}")
+        )
+
+    return
+
+
+def show_result(message: types.Message):
+    user = User(message.chat.id)
+    text = TestMessages.TEST_RESULT.format(
+        TestMessages.TEST_RESULTS[user.type][0],
+        TestMessages.TEST_RESULTS[user.type][1],
+    )
+    bot.send_photo(
+        message.chat.id,
+        open(f"pics/{user.type}.jpg", "rb"),
+        caption=text,
+    )
+
+    bot.send_message(message.chat.id, TestMessages.TEST_FINISH)
+
+    return
@@ -0,0 +1,6 @@
+from plugin.user import User
+
+
+def Shuffle(users: list[User]):
+
+    return
@@ -1,5 +1,7 @@
 class TestMessages:
    TEST_WELCOME = "Пройди следующий короткий тест, чтобы было легче подобрать тебе пару. А еще ты узнаешь, кто ты из героев советских мультфильмов."
+    TEST_RESULT = "Ты - {}!\n\n{}"
+    TEST_FINISH = "Поздравляем! Теперь вы зарегистрированы в системе."
    TEST_QUESTIONS = [
        {
            "question": "Когда вы оказываетесь в новой компании людей, вы, как правило:",
@@ -99,6 +99,14 @@ class User:
    def update_data(self):
        self._db.update({"_id": self.user_id}, self._into_json())

+    def delete(self):
+        self._db.delete({"_id": self.user_id})
+
+    def get_all() -> list["User"]:
+        users_data = User._db.find({})
+
+        return [User(data["_id"]) for data in users_data]
+
    def _save(self):
        if not self._db.exist({"_id": self.user_id}):
            self._db.insert(self._into_json())
@@ -107,7 +115,7 @@ class User:

    def _load(self):
        if not self._db.exist({"_id": self.user_id}):
-            return
+            raise ValueError("User not found")

        data = self._db.find_one({"_id": self.user_id})
        self._name = data["name"]