From 36fc400cd37a86b02c791a137570fa3422f725f6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=BC=D0=B8=D1=82=D1=80=D0=B8=D0=B9?= Date: Sat, 26 Mar 2022 17:07:22 +0300 Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D1=84=D0=B0=D0=BA=D1=82=D0=BE?= =?UTF-8?q?=D1=80=D0=B8=D0=BD=D0=B3.=20=D0=92=D1=8B=D0=BD=D0=BE=D1=88?= =?UTF-8?q?=D1=83=20=D0=BB=D0=BE=D0=B3=D0=B8=D0=BA=D1=83=20=D0=B8=D0=B7=20?= =?UTF-8?q?main=20=D0=B2=20services?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 81 +++++++++++++++++------------------------------------ services.py | 68 +++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 57 deletions(-) diff --git a/main.py b/main.py index cd1c78d..2688159 100644 --- a/main.py +++ b/main.py @@ -19,23 +19,6 @@ bot = Bot(token=settings.API_TOKEN) dp = Dispatcher(bot) -async def clean_up(message_id: int, delay_min: int): - exp = datetime.now() + timedelta(minutes=delay_min) - msg: Messages_to_delete = Messages_to_delete( - message_id=message_id, deletion_date=exp - ) - services.session.add(msg) - services.session.commit() - - await asyncio.sleep(delay_min * 60) - try: - await bot.delete_message(chat_id=settings.CHAT_ID, message_id=message_id) - services.session.delete(msg) - services.session.commit() - except MessageToDeleteNotFound: - print("Ошибка удаления сообщения") - - @dp.message_handler(commands=["start", "help"]) @services.auth async def send_welcome(message: types.Message): @@ -50,7 +33,6 @@ async def send_welcome(message: types.Message): @services.only_admins async def get_channel_id(message: types.Message) -> None: """Возвращает id канала по его имени.""" - c_id = message.text[12:] result = await bot.send_message( @@ -77,7 +59,9 @@ async def send_to_chanel(message: types.Message) -> None: parse_mode="Markdown", ) print(result.sender_chat.id) - await clean_up(message_id=result.message_id, delay_min=settings.DELAY_TIME) + await services._clean_up( + message_id=result.message_id, delay_min=settings.DELAY_TIME, bot=bot + ) @dp.message_handler(commands=["q"]) @@ -108,22 +92,20 @@ async def send_question(message: types.Message) -> None: parse_mode="Markdown", ) - await clean_up(message_id=result.message_id, delay_min=settings.DELAY_TIME_Q) + await services._clean_up( + message_id=result.message_id, delay_min=settings.DELAY_TIME_Q, bot=bot + ) @dp.message_handler(commands=["get_invite"]) @services.only_admins async def get_invite(message: types.Message) -> None: - - token = secrets.token_urlsafe(16) - + """Получить токен для нового пользователя.""" bot_info = await bot.get_me() bot_link = "@" + str(bot_info.username) - exp = datetime.now() + timedelta(minutes=5) - test = Token(token=token, expire=exp) - services.session.add(test) - services.session.commit() + token = services.get_new_token() + msg = f"Новый токен: *{token}*\nБудет действителен 5 минут\n\nЧтобы вступить в чат напиши боту {escape_md(bot_link)} \n`/add_me {token}`" await message.reply( @@ -143,40 +125,27 @@ async def add_me(message: types.Message) -> None: """ user_token = message.text[8:].strip() - tokens_list = services.session.query(Token).order_by(Token.id).all() - finded_token: Token = next((t for t in tokens_list if t.token == user_token), None) + if not services.check_token(user_token): + return - # если с токеном все впорядке, то добавляем юзера в базу - if finded_token and not finded_token.expired(): - new_user: Allowed_user = Allowed_user( - user_id=message["from"]["id"], date_add=datetime.now() - ) - try: - services.session.add(new_user) - services.session.delete(finded_token) - services.session.commit() - except IntegrityError: - print(f"Юзер {new_user.user_id} уже в базе") - services.session.rollback() - - # генерируем ссылку, истечет через 5 минут, может вступить 1 человек - link = await bot.create_chat_invite_link( - chat_id=settings.CHAT_ID, expire_date=timedelta(minutes=5), member_limit=1 - ) + services.add_new_user_to_allowed_list(message["from"]["id"]) + + # генерируем ссылку, истечет через 5 минут, может вступить 1 человек + link = await bot.create_chat_invite_link( + chat_id=settings.CHAT_ID, expire_date=timedelta(minutes=5), member_limit=1 + ) + with open("static/welcome_msg_short.txt", "r") as file: + msg = file.read() await message.reply( - "Welcome to the internet" - "\nHave a look around" - "\nAnything that brain of yours can think of can be found" - f"\n\nТвоя ссылка для вступления: {link['invite_link']}", - reply=False, - parse_mode="Markdown", + msg.format(invite=link["invite_link"]), reply=False, parse_mode="Markdown" ) - await asyncio.sleep(10) - with open("static/welcome_msg.txt", "r") as f: - msg = f.read() - await message.reply(msg, reply=False, parse_mode="Markdown") + + await asyncio.sleep(10) + with open("static/welcome_msg.txt", "r") as f: + msg = f.read() + await message.reply(msg, reply=False, parse_mode="Markdown") @dp.message_handler(commands=["stop"]) diff --git a/services.py b/services.py index f815a78..ab4c1bd 100644 --- a/services.py +++ b/services.py @@ -1,10 +1,18 @@ """Славный модуль с бизнес-логикой.""" -from models import Allowed_user +from models import Allowed_user, Messages_to_delete, Token from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker +from sqlalchemy.exc import IntegrityError import settings +import asyncio +from secrets import token_urlsafe + +from datetime import datetime, timedelta + +from aiogram.utils.exceptions import MessageToDeleteNotFound +from aiogram import Bot engine = create_engine("sqlite:///" + str(settings.DB_PATH)) Session = sessionmaker(bind=engine) @@ -48,3 +56,61 @@ def only_admins(func): return await func(message) return wrapper + + +async def _clean_up(message_id: int, delay_min: int, bot: Bot) -> None: + """Удаление сообщенния message_id через delay_min минут.""" + exp = datetime.now() + timedelta(minutes=delay_min) + msg: Messages_to_delete = Messages_to_delete( + message_id=message_id, deletion_date=exp + ) + session.add(msg) + session.commit() + + await asyncio.sleep(delay_min * 60) + + try: + await bot.delete_message(chat_id=settings.CHAT_ID, message_id=message_id) + session.delete(msg) + session.commit() + except MessageToDeleteNotFound: + print("Ошибка удаления сообщения") + + +def get_new_token() -> str: + """Возвращает токен для приглашения пользователя. + + Длина токена 16 символов, срок действия 5 минут. Токен записывается в базу. + + """ + token = token_urlsafe(16) + exp = datetime.now() + timedelta(minutes=5) + test = Token(token=token, expire=exp) + session.add(test) + session.commit() + + return token + + +def check_token(token: str) -> bool: + """Проверяет, находится ли токен в базе, возвращает True, если токен есть и не истек. затем удаляет токен из базы.""" + finded_token = session.query(Token).filter(Token.token == token).one_or_none() + + if finded_token and not finded_token.expired(): + session.delete(finded_token) + session.commit() + return True + + return False + + +def add_new_user_to_allowed_list(user_id: int) -> None: + """Добавляет пользователя в белый список. Если пользователь в нем уже есть, то игнорирует ошибку и ведет себя так, словно его добавили впервые.""" + new_user: Allowed_user = Allowed_user(user_id=user_id, date_add=datetime.now()) + + try: + session.add(new_user) + session.commit() + except IntegrityError: + print(f"Юзер {new_user.user_id} уже в базе") + session.rollback()