diff --git a/docker-compose.yml b/docker-compose.yml index 3d7b5b8..ab3dc17 100755 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -3,7 +3,7 @@ version: "3.6" services: database: - image: mongo:4.0.8 + image: mongo:5.0.2 restart: unless-stopped env_file: - ./docker/config/database.env diff --git a/src/app.py b/src/app.py index 79ec234..7c24a4f 100755 --- a/src/app.py +++ b/src/app.py @@ -3,7 +3,7 @@ from telegram.ext.dispatcher import Dispatcher from config.credentials import BOT_TOKEN from handler.abstractHandler import AbstractHandler -from handler import (inHandler, mentionHandler, outHandler) +from handler import (inHandler, mentionHandler, outHandler, silentMentionHandler, groupsHandler) class App: diff --git a/src/config/contents.py b/src/config/contents.py index a3345c1..f1a0abf 100755 --- a/src/config/contents.py +++ b/src/config/contents.py @@ -6,3 +6,4 @@ opted_in_failed = re.escape('You already opted-in for everyone-mentions.') opted_off = re.escape('You have opted-off for everyone-mentions.') opted_off_failed = re.escape('You need to opt-in first before processing this command.') mention_failed = re.escape('There are no users to mention.') +no_groups = re.escape('There are no groups for this chat.') \ No newline at end of file diff --git a/src/database/client.py b/src/database/client.py index 4f4af04..781c24a 100755 --- a/src/database/client.py +++ b/src/database/client.py @@ -1,3 +1,4 @@ +from typing import Optional from urllib.parse import quote_plus from config.credentials import (MONGODB_DATABASE, MONGODB_HOSTNAME, @@ -34,3 +35,6 @@ class Client(): filter, { "$set" : data } ) + + def aggregate(self, collection, pipeline: list): + return self.database.get_collection(collection).aggregate(pipeline) \ No newline at end of file diff --git a/src/entity/group.py b/src/entity/group.py new file mode 100644 index 0000000..ff526f9 --- /dev/null +++ b/src/entity/group.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class Group(): + chat_id: str + group_name: str + users_count: int + + default_name: str = 'default' diff --git a/src/exception/alreadyExistsException.py b/src/exception/alreadyExistsException.py deleted file mode 100644 index d3498b1..0000000 --- a/src/exception/alreadyExistsException.py +++ /dev/null @@ -1,2 +0,0 @@ -class AlreadyExistsException(Exception): - pass \ No newline at end of file diff --git a/src/exception/invalidArgumentException.py b/src/exception/invalidArgumentException.py new file mode 100644 index 0000000..53acd51 --- /dev/null +++ b/src/exception/invalidArgumentException.py @@ -0,0 +1,2 @@ +class InvalidArgumentException(Exception): + pass \ No newline at end of file diff --git a/src/handler/abstractHandler.py b/src/handler/abstractHandler.py index 2b800ba..5455009 100755 --- a/src/handler/abstractHandler.py +++ b/src/handler/abstractHandler.py @@ -14,8 +14,11 @@ class AbstractHandler: @abstractmethod def handle(self, update: Update, context: CallbackContext) -> None: raise Exception('handle method is not implemented') - def get_update_data(self, update: Update) -> UpdateData: - return UpdateData.create_from_update(update) + def get_update_data(self, update: Update, context: CallbackContext) -> UpdateData: + return UpdateData.create_from_arguments(update, context) - def reply(self, update: Update, message: str) -> None: + def reply_markdown(self, update: Update, message: str) -> None: update.effective_message.reply_markdown_v2(text=message) + + def reply_html(self, update: Update, html: str) -> None: + update.effective_message.reply_html(text=html) diff --git a/src/handler/groupsHandler.py b/src/handler/groupsHandler.py new file mode 100644 index 0000000..c6bd580 --- /dev/null +++ b/src/handler/groupsHandler.py @@ -0,0 +1,40 @@ +from typing import Iterable + +import prettytable as pt +from config.contents import no_groups +from entity.group import Group +from repository.groupRepository import GroupRepository +from telegram.ext.callbackcontext import CallbackContext +from telegram.ext.commandhandler import CommandHandler +from telegram.update import Update + +from handler.abstractHandler import AbstractHandler + + +class GroupsHandler(AbstractHandler): + bot_handler: CommandHandler + group_repository: GroupRepository + + def __init__(self) -> None: + self.bot_handler = CommandHandler('groups', self.handle) + self.group_repository = GroupRepository() + + def handle(self, update: Update, context: CallbackContext) -> None: + real_chat_id = str(update.effective_chat.id) + + groups = self.group_repository.get_by_chat_id(real_chat_id) + + if groups: + return self.reply_html(update, self.build_groups_message(groups)) + + self.reply_markdown(update, no_groups) + + def get_bot_handler(self) -> CommandHandler: + return self.bot_handler + + def build_groups_message(self, groups: Iterable[Group]) -> str: + resultTable = pt.PrettyTable(['Name', 'Members']) + + resultTable.add_rows([[record.group_name, record.users_count] for record in groups]) + + return f'
{str(resultTable)}
' diff --git a/src/handler/inHandler.py b/src/handler/inHandler.py index a7d6731..8fff3ef 100755 --- a/src/handler/inHandler.py +++ b/src/handler/inHandler.py @@ -1,4 +1,5 @@ from config.contents import opted_in, opted_in_failed +from exception.invalidArgumentException import InvalidArgumentException from exception.notFoundException import NotFoundException from repository.userRepository import UserRepository from telegram.ext.callbackcontext import CallbackContext @@ -17,14 +18,16 @@ class InHandler(AbstractHandler): self.user_repository = UserRepository() def handle(self, update: Update, context: CallbackContext) -> None: - update_data = self.get_update_data(update) + try: + update_data = self.get_update_data(update, context) + except InvalidArgumentException as e: + return self.reply_markdown(update, str(e)) try: user = self.user_repository.get_by_id(update_data.user_id) if user.is_in_chat(update_data.chat_id): - self.reply(update, opted_in_failed) - return + return self.reply_markdown(update, opted_in_failed) user.add_to_chat(update_data.chat_id) self.user_repository.save(user) @@ -32,7 +35,7 @@ class InHandler(AbstractHandler): except NotFoundException: self.user_repository.save_by_update_data(update_data) - self.reply(update, opted_in) + self.reply_markdown(update, opted_in) def get_bot_handler(self) -> CommandHandler: return self.bot_handler diff --git a/src/handler/mentionHandler.py b/src/handler/mentionHandler.py index 4a23865..9735b9e 100755 --- a/src/handler/mentionHandler.py +++ b/src/handler/mentionHandler.py @@ -2,6 +2,7 @@ from typing import Iterable from config.contents import mention_failed from entity.user import User +from exception.invalidArgumentException import InvalidArgumentException from repository.userRepository import UserRepository from telegram.ext.callbackcontext import CallbackContext from telegram.ext.commandhandler import CommandHandler @@ -13,35 +14,31 @@ from handler.abstractHandler import AbstractHandler class MentionHandler(AbstractHandler): bot_handler: CommandHandler user_repository: UserRepository - silent: str = 'silent' def __init__(self) -> None: self.bot_handler = CommandHandler('everyone', self.handle) self.user_repository = UserRepository() def handle(self, update: Update, context: CallbackContext) -> None: - updateData = self.get_update_data(update) + try: + updateData = self.get_update_data(update, context) + except InvalidArgumentException as e: + return self.reply_markdown(update, str(e)) + users = self.user_repository.get_all_for_chat(updateData.chat_id) if users: - self.reply(update, self.build_mention_message(users, self.isSilent(context))) - return + return self.reply_markdown(update, self.build_mention_message(users)) - self.reply(update, mention_failed) + self.reply_markdown(update, mention_failed) def get_bot_handler(self) -> CommandHandler: return self.bot_handler - def build_mention_message(self, users: Iterable[User], silent: bool = False) -> str: + def build_mention_message(self, users: Iterable[User]) -> str: result = '' for user in users: - if not silent: - result += f'*[{user.username}](tg://user?id={user.user_id})* ' - else: - result += f'*{user.username}\({user.user_id}\)*\n' + result += f'*[{user.username}](tg://user?id={user.user_id})* ' return result - - def isSilent(self, context: CallbackContext) -> bool: - return self.silent in context.args \ No newline at end of file diff --git a/src/handler/outHandler.py b/src/handler/outHandler.py index bb010c7..dc200a4 100755 --- a/src/handler/outHandler.py +++ b/src/handler/outHandler.py @@ -1,4 +1,5 @@ from config.contents import opted_off, opted_off_failed +from exception.invalidArgumentException import InvalidArgumentException from exception.notFoundException import NotFoundException from repository.userRepository import UserRepository from telegram.ext.callbackcontext import CallbackContext @@ -17,7 +18,10 @@ class OutHandler(AbstractHandler): self.user_repository = UserRepository() def handle(self, update: Update, context: CallbackContext) -> None: - updateData = self.get_update_data(update) + try: + updateData = self.get_update_data(update, context) + except InvalidArgumentException as e: + return self.reply_markdown(update, str(e)) try: user = self.user_repository.get_by_id(updateData.user_id) @@ -25,13 +29,12 @@ class OutHandler(AbstractHandler): if not user.is_in_chat(updateData.chat_id): raise NotFoundException() except NotFoundException: - self.reply(update, opted_off_failed) - return + return self.reply_markdown(update, opted_off_failed) user.remove_from_chat(updateData.chat_id) self.user_repository.save(user) - self.reply(update, opted_off) + self.reply_markdown(update, opted_off) def get_bot_handler(self) -> CommandHandler: return self.bot_handler diff --git a/src/handler/silentMentionHandler.py b/src/handler/silentMentionHandler.py new file mode 100644 index 0000000..e9ff22a --- /dev/null +++ b/src/handler/silentMentionHandler.py @@ -0,0 +1,21 @@ +from typing import Iterable + +from entity.user import User +from telegram.ext.commandhandler import CommandHandler + +from handler.abstractHandler import AbstractHandler +from handler.mentionHandler import MentionHandler + + +class MentionHandler(MentionHandler, AbstractHandler): + def __init__(self) -> None: + super().__init__() + self.bot_handler = CommandHandler('silent', self.handle) + + def build_mention_message(self, users: Iterable[User]) -> str: + result = '' + + for user in users: + result += f'*{user.username}\({user.user_id}\)*\n' + + return result diff --git a/src/handler/vo/updateData.py b/src/handler/vo/updateData.py index eb1a3bf..8739480 100644 --- a/src/handler/vo/updateData.py +++ b/src/handler/vo/updateData.py @@ -1,9 +1,14 @@ from __future__ import annotations from dataclasses import dataclass +import re import names +from telegram.ext.callbackcontext import CallbackContext from telegram.update import Update +from entity.group import Group + +from exception.invalidArgumentException import InvalidArgumentException @dataclass @@ -13,9 +18,17 @@ class UpdateData(): username: str @staticmethod - def create_from_update(update: Update) -> UpdateData: - user_id = str(update.effective_user.id) + def create_from_arguments(update: Update, context: CallbackContext) -> UpdateData: chat_id = str(update.effective_chat.id) + + if context.args and context.args[0]: + if not context.args[0].isalpha() or context.args[0] == Group.default_name: + raise InvalidArgumentException(re.escape(f'Group name must contain only letters and can not be `{Group.default_name}`.')) + else: + chat_id += f'~{context.args[0]}'.lower() + + + user_id = str(update.effective_user.id) username = update.effective_user.username or update.effective_user.first_name if not username: diff --git a/src/repository/groupRepository.py b/src/repository/groupRepository.py new file mode 100644 index 0000000..82363df --- /dev/null +++ b/src/repository/groupRepository.py @@ -0,0 +1,53 @@ +import itertools +import re +from typing import Iterable + +from database.client import Client +from entity.group import Group +from entity.user import User + + +class GroupRepository(): + client: Client + + count: str = 'count' + + def __init__(self) -> None: + self.client = Client() + + def get_by_chat_id(self, chat_id: str) -> Iterable[Group]: + groups = self.client.aggregate( + User.collection, + [ + { "$unwind": f'${User.chats_index}' }, + { + "$match": { + User.chats_index: { "$regex": re.compile(f'^{chat_id}.*$') }, + }, + }, + { + "$group": { + "_id": { + "$last": { "$split": [f'${User.chats_index}', "~"] }, + }, + self.count: { "$count": {} }, + }, + }, + { + "$sort": { '_id': 1 } + } + ] + ) + + result = [] + for group in groups: + group_name = group['_id'] + + if group_name == chat_id: + group_name = Group.default_name + + result.append( + Group(chat_id, group_name, group[self.count]) + ) + + return result diff --git a/src/requirements.txt b/src/requirements.txt index 81d707c..8db8dcc 100755 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,4 +1,5 @@ python-dotenv==0.19.0 python-telegram-bot==13.7 pymongo==3.12.0 -names==0.3.0 \ No newline at end of file +names==0.3.0 +prettytable==2.2.1 \ No newline at end of file