diff --git a/bot/commands/subscription.py b/bot/commands/subscription.py index 8851c50..429608f 100644 --- a/bot/commands/subscription.py +++ b/bot/commands/subscription.py @@ -6,6 +6,7 @@ from telegram.ext import CommandHandler, CallbackContext from bot.core import BotTelegramCore +from bot.messages import ADMIN_RESTRICTED from utils.utils import build_menu from db.subject import Subject from db.observer import UserObserver @@ -19,9 +20,14 @@ def subscribe(update: Update, context: CallbackContext): + user = update.effective_user chat = update.effective_chat message = update.effective_message + if not BotTelegramCore.instance().is_admin(user.id, chat.id): + message.reply_text(ADMIN_RESTRICTED) + return + try: observer = UserObserver.objects.get(chat_id=f"{chat.id}") except DoesNotExist: @@ -49,9 +55,14 @@ def subscribe(update: Update, context: CallbackContext): def unsubscribe(update: Update, context: CallbackContext): + user = update.effective_user chat = update.effective_chat message = update.effective_message + if not BotTelegramCore.instance().is_admin(user.id, chat.id): + message.reply_text(ADMIN_RESTRICTED) + return + try: observer = UserObserver.objects.get(chat_id=f"{chat.id}") except DoesNotExist: @@ -79,9 +90,14 @@ def unsubscribe(update: Update, context: CallbackContext): def subscriptions(update: Update, context: CallbackContext): + user = update.effective_user chat = update.effective_chat message = update.effective_message + if not BotTelegramCore.instance().is_admin(user.id, chat.id): + message.reply_text(ADMIN_RESTRICTED) + return + try: observer = UserObserver.objects.get(chat_id=f"{chat.id}") except DoesNotExist: diff --git a/bot/core.py b/bot/core.py index 7f146da..63f0a21 100644 --- a/bot/core.py +++ b/bot/core.py @@ -3,7 +3,7 @@ from decouple import config from telegram.ext import Updater, Handler -from telegram import Update, ParseMode +from telegram import Update, ParseMode, TelegramError logging.basicConfig( @@ -49,20 +49,28 @@ def chat_id(self): @property def chat(self): - return self._updater.bot.get_chat(self.chat_id) + return self.get_chat(self.chat_id) + + def get_chat(self, chat_id): + return self._updater.bot.get_chat(chat_id) def is_from_official_chat(self, update: Update): return self.chat_id == update.message.chat.id - @property - def administrators(self): + @staticmethod + def get_administrators_ids(chat): return [chat_member.user.id for - chat_member in self.chat.get_administrators()] + chat_member in chat.get_administrators()] @classmethod - def is_admin(cls, user_id): + def is_admin(cls, user_id, chat_id=None): instance = cls.instance() - return user_id in instance.administrators + try: + chat = instance.get_chat(chat_id) + except TelegramError: + chat = instance.chat + return (chat.type == chat.PRIVATE or + user_id in instance.get_administrators_ids(chat)) @classmethod def send_message(cls, text, chat_id, parse_mode=None): diff --git a/bot/messages.py b/bot/messages.py index e3d4bd2..5d4a6b8 100644 --- a/bot/messages.py +++ b/bot/messages.py @@ -4,3 +4,4 @@ "Read my documentation or create a issue on " \ "https://github.com/dcr-guys/JackBot/" +ADMIN_RESTRICTED = "This command is restricted to admins only!"