diff --git a/classes/user.py b/classes/user.py index d2e7c77..7e94896 100644 --- a/classes/user.py +++ b/classes/user.py @@ -1,5 +1,7 @@ +from modules.app import app from datetime import datetime -from modules.database import col_banned +from modules.database import col_banned, col_users +from modules.utils import configGet class PosterUser(): @@ -14,4 +16,17 @@ class PosterUser(): col_banned.insert_one({"user": self.id, "date": datetime.now()}) def unblock(self) -> None: - col_banned.find_one_and_delete({"user": self.id}) \ No newline at end of file + col_banned.find_one_and_delete({"user": self.id}) + + def is_limited(self) -> bool: + if self.id in app.admins: + return False + else: + db_record = col_users.find_one({"user": self.id}) + if db_record is None: + return False + return True if (datetime.now() - db_record["cooldown"]).total_seconds() < configGet("timeout", "submission") else False + + def limit(self) -> None: + if col_users.find_one_and_update({"user": self.id}, {"$set": {"cooldown": datetime.now()}}) is None: + col_users.insert_one({"user": self.id, "cooldown": datetime.now()}) \ No newline at end of file diff --git a/modules/submissions.py b/modules/submissions.py deleted file mode 100644 index e529835..0000000 --- a/modules/submissions.py +++ /dev/null @@ -1,18 +0,0 @@ -from modules.app import app -from datetime import datetime -from modules.utils import configGet -from modules.database import col_users -from pyrogram.types.user_and_chats import User - -def subLimit(user: User) -> None: - if col_users.find_one_and_update({"user": user.id}, {"$set": {"cooldown": datetime.now()}}) is None: - col_users.insert_one({"user": user.id, "cooldown": datetime.now()}) - -def subLimited(user: User) -> bool: - if user.id in app.admins: - return False - else: - db_record = col_users.find_one({"user": user.id}) - if db_record is None: - return False - return True if (datetime.now() - db_record["cooldown"]).total_seconds() < configGet("timeout", "submission") else False \ No newline at end of file diff --git a/plugins/handlers/submission.py b/plugins/handlers/submission.py index 6b683af..e09937e 100644 --- a/plugins/handlers/submission.py +++ b/plugins/handlers/submission.py @@ -1,4 +1,4 @@ -from datetime import datetime, timezone +from datetime import datetime from os import makedirs, path, sep from traceback import format_exc from uuid import uuid4 @@ -8,11 +8,11 @@ from pyrogram.types import InlineKeyboardButton, InlineKeyboardMarkup, Message from pyrogram.enums.chat_action import ChatAction from classes.exceptions import SubmissionDuplicatesError from classes.poster_client import PosterClient +from classes.user import PosterUser from modules.app import app from modules.database import col_banned, col_submitted from modules.logger import logWrite -from modules.submissions import subLimit, subLimited from modules.utils import configGet, locale from classes.enums.submission_types import SubmissionType @@ -31,7 +31,7 @@ async def get_submission(app: PosterClient, msg: Message): save_tmp = True contents = None - if subLimited(msg.from_user): + if PosterUser(msg.from_user.id).is_limited(): await msg.reply_text(locale("sub_cooldown", "message", locale=user_locale).format(str(configGet("timeout", "submission")))) return @@ -180,7 +180,7 @@ async def get_submission(app: PosterClient, msg: Message): # ] ] - subLimit(msg.from_user) + PosterUser(msg.from_user.id).limit() if msg.from_user.id != app.owner: await msg.reply_text(locale("sub_sent", "message", locale=user_locale), quote=True)