Browse Source

Рефакторинг. Разношу логику.

master
Дмитрий 3 years ago
parent
commit
614e41e158
  1. 72
      main.py
  2. 11
      services.py

72
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)

11
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)

Loading…
Cancel
Save