diff --git a/main.py b/main.py index 8f039d4..cd1c78d 100644 --- a/main.py +++ b/main.py @@ -1,16 +1,16 @@ +"""Главный модуль программы.""" import logging from aiogram import Bot, Dispatcher, executor, types from aiogram.utils.exceptions import MessageToDeleteNotFound import asyncio from datetime import timedelta, datetime import secrets -from sqlalchemy import create_engine -from sqlalchemy.orm import sessionmaker + from sqlalchemy.exc import IntegrityError from models import Allowed_user, Token, Messages_to_delete, Base from aiogram.utils.markdown import escape_md -from services import auth, only_admins -from . import settings +import settings +import services logging.basicConfig(level=logging.INFO) @@ -18,30 +18,26 @@ logging.basicConfig(level=logging.INFO) bot = Bot(token=settings.API_TOKEN) dp = Dispatcher(bot) -engine = create_engine("sqlite:///" + str(settings.DB_PATH)) -Session = sessionmaker(bind=engine) -session = Session() - 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 ) - session.add(msg) - session.commit() + 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) - session.delete(msg) - session.commit() + services.session.delete(msg) + services.session.commit() except MessageToDeleteNotFound: print("Ошибка удаления сообщения") @dp.message_handler(commands=["start", "help"]) -@auth +@services.auth async def send_welcome(message: types.Message): """Справка по командам.""" await message.reply( @@ -51,7 +47,7 @@ async def send_welcome(message: types.Message): @dp.message_handler(commands=["channel_id"]) -@only_admins +@services.only_admins async def get_channel_id(message: types.Message) -> None: """Возвращает id канала по его имени.""" @@ -69,7 +65,7 @@ async def get_channel_id(message: types.Message) -> None: @dp.message_handler(commands=["mes"]) -@auth +@services.auth async def send_to_chanel(message: types.Message) -> None: deletion_time = message.date + timedelta(minutes=settings.DELAY_TIME) @@ -85,7 +81,7 @@ async def send_to_chanel(message: types.Message) -> None: @dp.message_handler(commands=["q"]) -@auth +@services.auth async def send_question(message: types.Message) -> None: """Отправить вопрос через бота. @@ -116,7 +112,7 @@ async def send_question(message: types.Message) -> None: @dp.message_handler(commands=["get_invite"]) -@only_admins +@services.only_admins async def get_invite(message: types.Message) -> None: token = secrets.token_urlsafe(16) @@ -126,8 +122,8 @@ async def get_invite(message: types.Message) -> None: exp = datetime.now() + timedelta(minutes=5) test = Token(token=token, expire=exp) - session.add(test) - session.commit() + services.session.add(test) + services.session.commit() msg = f"Новый токен: *{token}*\nБудет действителен 5 минут\n\nЧтобы вступить в чат напиши боту {escape_md(bot_link)} \n`/add_me {token}`" await message.reply( @@ -147,7 +143,7 @@ async def add_me(message: types.Message) -> None: """ user_token = message.text[8:].strip() - tokens_list = session.query(Token).order_by(Token.id).all() + 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) @@ -157,12 +153,12 @@ async def add_me(message: types.Message) -> None: user_id=message["from"]["id"], date_add=datetime.now() ) try: - session.add(new_user) - session.delete(finded_token) - session.commit() + services.session.add(new_user) + services.session.delete(finded_token) + services.session.commit() except IntegrityError: print(f"Юзер {new_user.user_id} уже в базе") - session.rollback() + services.session.rollback() # генерируем ссылку, истечет через 5 минут, может вступить 1 человек link = await bot.create_chat_invite_link( @@ -184,7 +180,7 @@ async def add_me(message: types.Message) -> None: @dp.message_handler(commands=["stop"]) -@auth +@services.auth async def stop_and_panic(message: types.Message) -> None: """Panic-button для бота. @@ -194,13 +190,15 @@ async def stop_and_panic(message: types.Message) -> None: """ u_id = message["from"]["id"] finded_user: Allowed_user = ( - session.query(Allowed_user).filter(Allowed_user.user_id == u_id).one_or_none() + services.session.query(Allowed_user) + .filter(Allowed_user.user_id == u_id) + .one_or_none() ) await bot.kick_chat_member( chat_id=settings.CHAT_ID, user_id=finded_user.user_id, revoke_messages=True ) - session.delete(finded_user) - session.commit() + services.session.delete(finded_user) + services.session.commit() # рассылка сообщения админам что юзер нажал паническую кнопку. Нужны ид чатов. # for admin_id in ADMINS: @@ -208,10 +206,10 @@ async def stop_and_panic(message: types.Message) -> None: @dp.message_handler(commands=["wlc"]) -@only_admins +@services.only_admins async def wlc(message: types.Message) -> None: """Тестирование приветственного сообщения.""" - with open("welcome_msg.txt", "r") as f: + with open("static/welcome_msg.txt", "r") as f: msg = f.read() await message.reply( msg, reply=False, parse_mode="Markdown", disable_web_page_preview=True @@ -220,7 +218,9 @@ async def wlc(message: types.Message) -> None: # https://stackoverflow.com/questions/67637631/create-a-background-process-using-aiogram async def messages_cleanup() -> None: - msg_list = session.query(Messages_to_delete).order_by(Messages_to_delete.id).all() + msg_list = ( + services.session.query(Messages_to_delete).order_by(Messages_to_delete.id).all() + ) while True: for m in msg_list: if m.deletion_needed: @@ -228,12 +228,12 @@ async def messages_cleanup() -> None: await bot.delete_message( chat_id=settings.CHAT_ID, message_id=m.message_id ) - session.delete(m) - session.commit() + services.session.delete(m) + services.session.commit() except MessageToDeleteNotFound: print("Сообщение не найдено!") - session.delete(m) - session.commit() + services.session.delete(m) + services.session.commit() await asyncio.sleep(600) @@ -242,5 +242,5 @@ async def on_startup(dispatcher: Dispatcher) -> None: if __name__ == "__main__": - Base.metadata.create_all(engine) + Base.metadata.create_all(services.engine) executor.start_polling(dp, skip_updates=True, on_startup=on_startup) diff --git a/services.py b/services.py index 589d611..f815a78 100644 --- a/services.py +++ b/services.py @@ -1,6 +1,15 @@ """Славный модуль с бизнес-логикой.""" from models import Allowed_user +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +import settings + +engine = create_engine("sqlite:///" + str(settings.DB_PATH)) +Session = sessionmaker(bind=engine) +session = Session() + def auth(func): """Аутентификация. @@ -33,7 +42,7 @@ def only_admins(func): async def wrapper(message): uid = message["from"]["id"] - if uid not in ADMINS: + if uid not in settings.ADMINS: return return await func(message)