126 lines
3.2 KiB
Python
126 lines
3.2 KiB
Python
|
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)
|