56 lines
1.7 KiB
Python
56 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):
|
||
|
|
col_users.update_one({"_id": self._id}, {"$set": {"locale": locale}})
|
||
|
|
|
||
|
|
async def update_cooldown(self, time: datetime = datetime.now()):
|
||
|
|
col_users.update_one({"_id": self._id}, {"$set": {"cooldown": time}})
|
||
|
|
|
||
|
|
async def block(self) -> None:
|
||
|
|
"""Ban user from using command and submitting content."""
|
||
|
|
col_users.update_one({"_id": self._id}, {"$set": {"banned": True}})
|
||
|
|
|
||
|
|
async def unblock(self) -> None:
|
||
|
|
"""Allow user to use command and submit posts again."""
|
||
|
|
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 = (
|
||
|
|
app.admins
|
||
|
|
if app is not None
|
||
|
|
else (
|
||
|
|
await config_get("admins", "bot") + [await config_get("owner", "bot")]
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
return (datetime.now() - self.cooldown).total_seconds() < (
|
||
|
|
app.config["submission"]["timeout"]
|
||
|
|
if app is not None
|
||
|
|
else await config_get("timeout", "submission")
|
||
|
|
)
|