From be24e124bb796640e5fedcbe7b6c965c8859370a Mon Sep 17 00:00:00 2001 From: krosh Date: Sun, 23 Feb 2025 00:37:04 +0300 Subject: [PATCH] add admin logic --- app/config.py | 13 +- app/db.py | 3 + app/main.py | 27 +++- app/plugin/admin.py | 277 +++++++++++++++++++++++++++++++++++-- app/plugin/bot_instance.py | 6 + app/plugin/register.py | 71 +++++++++- app/plugin/shuffle.py | 6 + app/plugin/test.py | 2 + app/plugin/user.py | 10 +- 9 files changed, 389 insertions(+), 26 deletions(-) create mode 100644 app/plugin/shuffle.py diff --git a/app/config.py b/app/config.py index 7f7a6de..4958e15 100644 --- a/app/config.py +++ b/app/config.py @@ -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 = "Выбери пол (женский, мужской)" diff --git a/app/db.py b/app/db.py index cf3d414..970effd 100644 --- a/app/db.py +++ b/app/db.py @@ -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) diff --git a/app/main.py b/app/main.py index 08ad55f..8181703 100644 --- a/app/main.py +++ b/app/main.py @@ -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 ---") diff --git a/app/plugin/admin.py b/app/plugin/admin.py index 9ceca41..819a429 100644 --- a/app/plugin/admin.py +++ b/app/plugin/admin.py @@ -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, + "Ваш ключ: {}".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, + "Отправьте delete all для подтверждения", + 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 - удаляет ВСЕХ пользователей" + ) diff --git a/app/plugin/bot_instance.py b/app/plugin/bot_instance.py index 3c94a83..6f77a77 100644 --- a/app/plugin/bot_instance.py +++ b/app/plugin/bot_instance.py @@ -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, diff --git a/app/plugin/register.py b/app/plugin/register.py index c5725e3..55d7426 100644 --- a/app/plugin/register.py +++ b/app/plugin/register.py @@ -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 diff --git a/app/plugin/shuffle.py b/app/plugin/shuffle.py new file mode 100644 index 0000000..514b7fc --- /dev/null +++ b/app/plugin/shuffle.py @@ -0,0 +1,6 @@ +from plugin.user import User + + +def Shuffle(users: list[User]): + + return diff --git a/app/plugin/test.py b/app/plugin/test.py index 1276312..476a6a4 100644 --- a/app/plugin/test.py +++ b/app/plugin/test.py @@ -1,5 +1,7 @@ class TestMessages: TEST_WELCOME = "Пройди следующий короткий тест, чтобы было легче подобрать тебе пару. А еще ты узнаешь, кто ты из героев советских мультфильмов." + TEST_RESULT = "Ты - {}!\n\n{}" + TEST_FINISH = "Поздравляем! Теперь вы зарегистрированы в системе." TEST_QUESTIONS = [ { "question": "Когда вы оказываетесь в новой компании людей, вы, как правило:", diff --git a/app/plugin/user.py b/app/plugin/user.py index a4c1ccb..110faae 100644 --- a/app/plugin/user.py +++ b/app/plugin/user.py @@ -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"]