from dataclasses import dataclass from datetime import datetime from typing import Union from bson import ObjectId from libbot import config_get from libbot.pyrogram.classes import PyroClient from modules.database import col_users @dataclass class PyroUser: """Dataclass of DB entry of a user""" _id: ObjectId id: int locale: Union[str, None] banned: bool cooldown: datetime subscription: dict async def update_locale(self, locale: str) -> None: await col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}}) async def update_cooldown(self, time: datetime = datetime.now()) -> None: await col_users.update_one({"_id": self._id}, {"$set": {"cooldown": time}}) async def block(self) -> None: """Ban user from using command and submitting content.""" await col_users.update_one({"_id": self._id}, {"$set": {"banned": True}}) async def unblock(self) -> None: """Allow user to use command and submit posts again.""" await col_users.update_one({"_id": self._id}, {"$set": {"banned": False}}) async def is_limited(self, app: Union[PyroClient, None] = None) -> bool: """Check if user is on a cooldown after submitting something. ### Returns: `bool`: Must be `True` if on the cooldown and `False` if not """ admins = ( await config_get("admins", "bot") + [await config_get("owner", "bot")] if app is None else app.admins ) return (datetime.now() - self.cooldown).total_seconds() < ( app.config["submission"]["timeout"] if app is not None else await config_get("timeout", "submission") )