GloryBot/classes/pycord_check.py

126 lines
3.2 KiB
Python
Raw Normal View History

2024-06-02 22:51:07 +03:00
import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from random import choice
from typing import Optional
from bson import ObjectId
from classes.errors import (
CheckAlreadyAssignedError,
CheckChallengeNotFoundError,
CheckNotFoundError,
GuildChallengesEmptyError,
)
from classes.pycord_challenge import PycordChallenge
from modules.database import col_challenges, col_checks
logger = logging.getLogger(__name__)
@dataclass
class PycordCheck:
"""Dataclass of DB entry of a security check"""
__slots__ = (
"_id",
"guild",
"active",
"thread_id",
"member",
"date_created",
"date_modified",
"challenge",
)
_id: ObjectId
guild: ObjectId
active: bool
thread_id: int
member: ObjectId
date_created: datetime
date_modified: datetime
challenge: ObjectId
@classmethod
async def find(
cls,
guild: ObjectId,
member: ObjectId,
):
db_entry = await col_checks.find_one(
{"active": True, "guild": guild, "member": member}
)
if not db_entry:
raise CheckNotFoundError(guild, member)
return cls(**db_entry)
@classmethod
async def create(
cls,
guild: ObjectId,
thread_id: int,
member: ObjectId,
challenge: Optional[ObjectId] = None,
):
# Check whether active check already exists
if await col_checks.find_one(
{"active": True, "guild": guild, "member": member}
):
raise CheckAlreadyAssignedError(guild, member)
# Get all enabled guild challenges
guild_challenges = [
challenge
async for challenge in col_challenges.find(
{"guild": guild, "enabled": True, "archived": False}
)
]
# Check whether guild has challenges at all
if not challenge and not guild_challenges:
raise GuildChallengesEmptyError(guild)
# Create a check dict
check = {
"guild": guild,
"active": True,
"thread_id": thread_id,
"member": member,
"date_created": datetime.now(UTC),
"date_modified": datetime.now(UTC),
"challenge": challenge or choice(guild_challenges)["_id"],
}
# Insert the check into the database
inserted = await col_checks.insert_one(check)
check["_id"] = inserted.inserted_id
return cls(**check)
async def deactivate(self) -> None:
"""Deactivate the check"""
await col_checks.update_one(
{"_id": self._id},
{"$set": {"active": False, "date_modified": datetime.now(UTC)}},
)
async def get_challenge(self) -> PycordChallenge:
"""Get check's bound challenge
### Raises:
* `CheckChallengeNotFoundError`: Challenge could not be found
### Returns:
* `PycordChallenge`: Captcha challenge
"""
challenge = await col_challenges.find_one({"_id": self.challenge})
if not challenge:
raise CheckChallengeNotFoundError(self._id, self.challenge)
return PycordChallenge(**challenge)