diff --git a/daemons/bot.py b/daemons/bot.py index feb46c0..b8ec4b3 100644 --- a/daemons/bot.py +++ b/daemons/bot.py @@ -1,40 +1,15 @@ -import logging import os import telebot -from cachetools import TTLCache -from requests import get from telebot.types import Message import settings from helpers.mongo import mongo +from helpers.platform import platform bot = telebot.TeleBot(os.getenv("TELEGRAM_TOKEN")) -STAFF_CACHE = TTLCache(100, 60) - - -def is_staff(telegram_id): - logging.info("is staff check") - staff_data = STAFF_CACHE.get(telegram_id) - if staff_data is None: - staff_data = get( - 'https://platform.sprinthub.ru/is_staff', - headers={'X-Security-Token': settings.PLATFORM_SECURITY_TOKEN}, - params={ - 'telegram_id': telegram_id - } - ) - if staff_data.status_code != 200: - logging.info('return status ' + str(staff_data.status_code)) - return False - staff_data = staff_data.json()['is_staff'] - logging.info('got info ' + str(staff_data)) - STAFF_CACHE[telegram_id] = staff_data - return staff_data - - @bot.message_handler(commands=['start']) def on_start(message: Message): mongo.users_collection.delete_many({"chat_id": message.chat.id}) @@ -43,7 +18,7 @@ def on_start(message: Message): @bot.message_handler() def do_action(message: Message): - if settings.STAGE == 'development' and not is_staff(message.chat.id): + if settings.STAGE == 'development' and not platform.get_staff(message.chat.id): return from helpers.answer import Answer Answer(message).process() diff --git a/helpers/answer.py b/helpers/answer.py index 47ab1db..6be6dbb 100644 --- a/helpers/answer.py +++ b/helpers/answer.py @@ -94,12 +94,12 @@ class Answer: if len(lessons) == 0: self.send_message( "Сегодня у тебя нет пар, отдыхай", - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) else: self.send_message( "Твои пары сегодня:\n" + ruz.schedule_builder(lessons), - reply_markup=main_keyboard(), + reply_markup=main_keyboard(self.user['chat_id']), parse_mode='Markdown', ) self.set_state("ready") @@ -150,7 +150,7 @@ class Answer: text = "Я не понимаю такой команды, используй кнопки." self.send_message( text, - reply_markup=main_keyboard(), + reply_markup=main_keyboard(self.user['chat_id']), parse_mode='Markdown' ) @@ -171,7 +171,7 @@ class Answer: self.set_state("ready") self.send_message( self.message_text, - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) else: self.send_message("Используй кнопки!", reply_markup=notify_type()) @@ -192,7 +192,7 @@ class Answer: return mongo.users_collection.update_one({"chat_id": self.user['chat_id']}, {"$set": {"first_lesson_notify": time_notify}}) self.set_state("ready") - self.send_message("Запомнил!", reply_markup=main_keyboard()) + self.send_message("Запомнил!", reply_markup=main_keyboard(self.user['chat_id'])) def handle_state_wait_for_daily_notify_type(self): if self.message_text == "Текущий день": @@ -216,7 +216,7 @@ class Answer: elif self.message_text == "Не уведомлять": self.send_message( "Принято! Я не буду уведомлять тебя.", - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) mongo.users_collection.update_one( {"chat_id": self.user['chat_id']}, @@ -225,7 +225,7 @@ class Answer: elif self.message_text == "Назад": self.send_message( "Возвращаюсь!", - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) mongo.users_collection.update_one( {"chat_id": self.user['chat_id']}, @@ -259,7 +259,7 @@ class Answer: text = f"Принято! Буду уведомлять тебя за {self.message_text}." else: text = f"Принято! Я не буду уведомлять тебя." - self.send_message(text, reply_markup=main_keyboard()) + self.send_message(text, reply_markup=main_keyboard(self.user['chat_id'])) self.set_state("ready") def _validate_time(self, line: str) -> bool: @@ -281,7 +281,7 @@ class Answer: if self.message_text == "Не уведомлять": self.send_message( "Принято! Я не буду присылать тебе ежедневные уведомления.", - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) mongo.users_collection.update_one( {"chat_id": self.user['chat_id']}, @@ -298,7 +298,7 @@ class Answer: next_time = get_next_daily_notify_time(self.user) self.send_message( f"Принято! Буду уведомлять тебя каждый день в {self.message_text}.", - reply_markup=main_keyboard() + reply_markup=main_keyboard(self.user['chat_id']) ) mongo.users_collection.update_one( {"chat_id": self.user['chat_id']}, @@ -314,7 +314,7 @@ class Answer: mongo.users_collection.delete_one({"email": self.user['email']}) self.send_message("Настройки сброшены, ждем твоего возвращения", reply_markup=again_keyboard()) elif self.message_text == "Нет": - self.send_message("Возращаюсь к прежнему режиму", reply_markup=main_keyboard()) + self.send_message("Возращаюсь к прежнему режиму", reply_markup=main_keyboard(self.user['chat_id'])) self.set_state("ready") else: self.send_message("Я не понимаю, используй кнопки", reply_markup=yes_no_keyboard()) diff --git a/helpers/keyboards.py b/helpers/keyboards.py index b748157..c1c9bb2 100644 --- a/helpers/keyboards.py +++ b/helpers/keyboards.py @@ -1,15 +1,15 @@ import telebot -import settings +from helpers.platform import platform -def main_keyboard(): +def main_keyboard(telegram_id): kb = telebot.types.ReplyKeyboardMarkup(True, False) kb.row("Пары сегодня", "Пары завтра") kb.row("Расписание на неделю") kb.row("Напоминания о парах") kb.row("Ежедневные уведомления") - if settings.DEBUG: + if platform.experiment_enabled_for_user('alice', telegram_id): kb.row("Подключение Алисы") kb.row("Сброс настроек") return kb diff --git a/helpers/platform.py b/helpers/platform.py new file mode 100644 index 0000000..7d2b36b --- /dev/null +++ b/helpers/platform.py @@ -0,0 +1,97 @@ +import os + +from cachetools import TTLCache +from requests import get + + +class Platform: + def __init__(self): + self.staff_cache = TTLCache(1000, 60) + self.exp_cache = TTLCache(1000, 60) + self.configs = TTLCache(1000, 60) + self.stage = os.getenv("STAGE", "local") + self.token = os.getenv("PLATFORM_SECURITY_TOKEN") + self.project = "РУЗ Бот" + + def get_config(self, config_name): + config = self.configs.get(config_name) + if config is None: + config = get( + 'https://platform.sprinthub.ru/configs/get', + headers={'X-Security-Token': self.token}, + params={ + 'project': self.project, + 'stage': self.stage, + 'name': config_name + } + ) + if config.status_code != 200: + return {} + config = config.json() + self.configs[config_name] = config + return config + + def get_experiment(self, experiment_name): + exp = self.exp_cache.get(experiment_name) + if exp is None: + exp = get( + 'https://platform.sprinthub.ru/experiments/get', + headers={'X-Security-Token': self.token}, + params={ + 'project': self.project, + 'stage': self.stage, + 'name': experiment_name + } + ) + if exp.status_code != 200: + return { + 'enabled': False, + 'condition': False + } + experiment_name = exp.json() + self.exp_cache[experiment_name] = exp + return exp + + def get_staff(self, telegram_id): + is_staff = self.staff_cache.get(telegram_id) + if is_staff is None: + exp = get( + 'https://platform.sprinthub.ru/is_staff', + headers={'X-Security-Token': self.token}, + params={ + 'telegram_id': telegram_id, + } + ) + if exp.status_code != 200: + return False + data = exp.json() + is_staff = data['is_staff'] + self.staff_cache[telegram_id] = is_staff + return is_staff + + def experiment_enabled_for_user(self, experiment_name, telegram_id): + exp_data = self.get_experiment(experiment_name) + if not exp_data['enabled']: + return False + + class User: + def __init__(self, platform, telegram_id): + self.platform = platform + self.telegram_id = telegram_id + + @property + def is_staff(self): + return self.platform.get_staff(self.telegram_id) + + @property + def is_superuser(self): + return False + + user = User(self, telegram_id) + try: + return bool(eval(exp_data['condition'])) + except: + return False + + +platform = Platform()