54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
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")
|
|
)
|