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"]