From e9b6fb5c3fa364a919e62bfb66f1dff9e358ae2c 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: Thu, 17 Mar 2022 17:57:40 +0300 Subject: [PATCH] =?UTF-8?q?=D0=92=D1=8B=D0=BD=D0=B5=D1=81=20=D0=BD=D0=B0?= =?UTF-8?q?=D1=81=D1=82=D1=80=D0=BE=D0=B9=D0=BA=D0=B8=20=D0=B2=20=D0=BE?= =?UTF-8?q?=D0=BA=D1=80=D1=83=D0=B6=D0=B5=D0=BD=D0=B8=D0=B5,=20=D0=B2?= =?UTF-8?q?=D1=8B=D0=BD=D0=B5=D1=81=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB=D0=B8?= =?UTF-8?q?=20=D0=B4=D0=B0=D0=BD=D0=BD=D1=8B=D1=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 84 ++++++++++++++----------------------------------------- models.py | 66 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 63 deletions(-) create mode 100644 models.py diff --git a/main.py b/main.py index 08828e6..7c27dd4 100644 --- a/main.py +++ b/main.py @@ -6,78 +6,28 @@ import asyncio from datetime import timedelta, datetime import secrets from sqlalchemy import create_engine -from sqlalchemy import Column, Integer, String, DateTime from sqlalchemy.orm import declarative_base, sessionmaker +from sqlalchemy.exc import IntegrityError + +from models import Allowed_user, Token, Messages_to_delete logging.basicConfig(level=logging.INFO) -DELAY_TIME = 1 -DELAY_TIME_Q = 3 -CHAT_ID = -1001612776177 +DELAY_TIME = int(os.getenv("GBOT_DELAY_TIME")) +DELAY_TIME_Q = int(os.getenv("GBOT_DELAY_TIME_Q")) +CHAT_ID = os.getenv("GBOT_CHAT_ID") ALLOWED_USERS = [106693654, 7063133] ADMINS = [106693654, 7063133, 4978608] API_TOKEN = os.getenv("TELEGRAM_API_TOKEN") +DB_PATH = os.getenv("GBOT_DB_PATH") bot = Bot(token=API_TOKEN) dp = Dispatcher(bot) -engine = create_engine("sqlite:///support.db") +engine = create_engine("sqlite:///" + str(DB_PATH)) Base = declarative_base() Session = sessionmaker(bind=engine) session = Session() -class Allowed_user(Base): - __tablename__ = "allowed_users" - id = Column(Integer, primary_key=True) - user_id = Column(Integer, nullable=False) - date_add = Column(DateTime, nullable=False) - - -class Token(Base): - """Класс для токенов. - - Хранит сам ключ, и дату истечения. - - """ - - __tablename__ = "tokens" - id = Column(Integer, primary_key=True) - token = Column(String, nullable=False) - expire = Column(DateTime, nullable=False) - - def expired(self) -> bool: - """Проверка срока жизни токена. - - Вовзращает истину, если срок действия токена истек. - """ - if datetime.now() > self.expire: - return True - - return False - - def __repr__(self): - """Базовое представление класса.""" - return f"{self.id}: {self.token} | {self.expire}" - - -class Messages_to_delete(Base): - """Сообщения, которые нужно будет удалить.""" - - __tablename__ = "messages_to_delete" - id = Column(Integer, primary_key=True) - message_id = Column(Integer, nullable=False) - deletion_date = Column(DateTime, nullable=False) - - def deletion_needed(self) -> bool: - """Нужно ли удалить это сообщение. - - Мы могли пропустить удаление из-за перезапуска. - """ - if datetime.now() > self.deletion_date: - return True - - return False - - def auth(func): async def wrapper(message): users_list: Allowed_user = ( @@ -89,7 +39,7 @@ def auth(func): ) if not finded_user: - return # await message.reply("Access Denied", reply=False) + return return await func(message) return wrapper @@ -210,8 +160,13 @@ async def add_me(message: types.Message): new_user: Allowed_user = Allowed_user( user_id=message["from"]["id"], date_add=datetime.now() ) - session.add(new_user) - session.commit() + try: + session.add(new_user) + session.delete(finded_token) + session.commit() + except IntegrityError: + print(f"Юзер {new_user.user_id} уже в базе") + session.rollback() # генерируем ссылку, истечет через 5 минут, может вступить 1 человек link = await bot.create_chat_invite_link( @@ -250,9 +205,12 @@ async def messages_cleanup() -> None: if m.deletion_needed: try: await bot.delete_message(chat_id=CHAT_ID, message_id=m.message_id) - + session.delete(m) + session.commit() except MessageToDeleteNotFound: - print("Ошибка удаления сообщения") + print("Сообщение не найдено!") + session.delete(m) + session.commit() await asyncio.sleep(600) diff --git a/models.py b/models.py new file mode 100644 index 0000000..096d7a0 --- /dev/null +++ b/models.py @@ -0,0 +1,66 @@ +"""Модель для данных.""" +from sqlalchemy import Column, Integer, String, DateTime, UniqueConstraint +from sqlalchemy.orm import declarative_base +from datetime import datetime + +Base = declarative_base() + + +class Allowed_user(Base): + """Белый список пользователей. + + При добавлении в базу user_id проверяется на уникальность (в + базе), чтобы избежать повторов. + """ + + __tablename__ = "allowed_users" + id = Column(Integer, primary_key=True) + user_id = Column(Integer, nullable=False) + date_add = Column(DateTime, nullable=False) + __table_args__ = (UniqueConstraint("user_id", name="unique_uid"),) + + +class Token(Base): + """Класс для токенов. + + Хранит сам ключ, и дату истечения. + + """ + + __tablename__ = "tokens" + id = Column(Integer, primary_key=True) + token = Column(String, nullable=False) + expire = Column(DateTime, nullable=False) + + def expired(self) -> bool: + """Проверка срока жизни токена. + + Вовзращает истину, если срок действия токена истек. + """ + if datetime.now() > self.expire: + return True + + return False + + def __repr__(self): + """Базовое представление класса.""" + return f"{self.id}: {self.token} | {self.expire}" + + +class Messages_to_delete(Base): + """Сообщения, которые нужно будет удалить.""" + + __tablename__ = "messages_to_delete" + id = Column(Integer, primary_key=True) + message_id = Column(Integer, nullable=False) + deletion_date = Column(DateTime, nullable=False) + + def deletion_needed(self) -> bool: + """Нужно ли удалить это сообщение. + + Мы могли пропустить удаление из-за перезапуска. + """ + if datetime.now() > self.deletion_date: + return True + + return False