129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
from datetime import datetime, timedelta
|
|
from typing import Any, ClassVar, Dict, List, Optional, Union
|
|
|
|
from discord import Forbidden, HTTPException, InvalidData, NotFound
|
|
from pydantic import BaseModel
|
|
from redbot.core.bot import Red
|
|
|
|
from aurora.utilities.utils import generate_dict
|
|
|
|
|
|
class AuroraBaseModel(BaseModel):
|
|
"""Base class for all models in Aurora."""
|
|
class Moderation(AuroraBaseModel):
|
|
bot: ClassVar[Red]
|
|
moderation_id: int
|
|
guild_id: int
|
|
timestamp: datetime
|
|
moderation_type: str
|
|
target_type: str
|
|
target_id: int
|
|
moderator_id: int
|
|
role_id: Optional[int]
|
|
duration: Optional[timedelta]
|
|
end_timestamp: Optional[datetime]
|
|
reason: Optional[str]
|
|
resolved: bool
|
|
resolved_by: Optional[int]
|
|
resolve_reason: Optional[str]
|
|
expired: bool
|
|
changes: List[Dict]
|
|
metadata: Dict
|
|
|
|
@property
|
|
def id(self) -> int:
|
|
return self.moderation_id
|
|
|
|
@property
|
|
def type(self) -> str:
|
|
return self.moderation_type
|
|
|
|
async def get_moderator(self) -> "PartialUser":
|
|
return await PartialUser.from_id(self.bot, self.moderator_id)
|
|
|
|
async def get_target(self) -> Union["PartialUser", "PartialChannel"]:
|
|
if self.target_type == "user":
|
|
return await PartialUser.from_id(self.bot, self.target_id)
|
|
else:
|
|
return await PartialChannel.from_id(self.bot, self.target_id)
|
|
|
|
async def get_resolved_by(self) -> Optional["PartialUser"]:
|
|
if self.resolved_by:
|
|
return await PartialUser.from_id(self.bot, self.resolved_by)
|
|
return None
|
|
|
|
def __str__(self):
|
|
return f"{self.moderation_type} {self.target_type} {self.target_id} {self.reason}"
|
|
|
|
@classmethod
|
|
def from_sql(cls, bot: Red, moderation_id: int, guild_id: int) -> Optional["Moderation"]:
|
|
from aurora.utilities.database import connect
|
|
query = f"SELECT * FROM moderation_{guild_id} WHERE moderation_id = ?;"
|
|
|
|
with connect() as database:
|
|
cursor = database.cursor()
|
|
cursor.execute(query, (moderation_id,))
|
|
result = cursor.fetchone()
|
|
|
|
if result:
|
|
case = generate_dict(result, guild_id)
|
|
cursor.close()
|
|
return cls.from_dict(bot, case)
|
|
|
|
return None
|
|
|
|
@classmethod
|
|
def from_dict(cls, bot: Red, data: dict) -> "Moderation":
|
|
return cls(bot=bot, **data)
|
|
|
|
def to_json(self, indent: int = None, file: Any = None):
|
|
from aurora.utilities.json import dump, dumps
|
|
return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent)
|
|
|
|
|
|
class PartialUser(AuroraBaseModel):
|
|
id: int
|
|
username: str
|
|
discriminator: int
|
|
|
|
@property
|
|
def name(self):
|
|
return f"{self.username}#{self.discriminator}" if self.discriminator == 0 else self.username
|
|
|
|
def __str__(self):
|
|
return self.name
|
|
|
|
@classmethod
|
|
async def from_id(cls, bot: Red, user_id: int) -> "PartialUser":
|
|
user = bot.get_user(user_id)
|
|
if not user:
|
|
try:
|
|
user = await bot.fetch_user(user_id)
|
|
return cls(id=user.id, username=user.name, discriminator=user.discriminator)
|
|
except NotFound:
|
|
return cls(id=user_id, username="Deleted User", discriminator=0)
|
|
|
|
class PartialChannel(AuroraBaseModel):
|
|
id: int
|
|
name: str
|
|
|
|
@property
|
|
def mention(self):
|
|
if self.name == "Deleted Channel" or self.name == "Forbidden Channel":
|
|
return self.name
|
|
return f"<#{self.id}>"
|
|
|
|
def __str__(self):
|
|
return self.mention
|
|
|
|
@classmethod
|
|
async def from_id(cls, bot: Red, channel_id: int) -> "PartialChannel":
|
|
user = bot.get_channel(channel_id)
|
|
if not user:
|
|
try:
|
|
user = await bot.fetch_channel(channel_id)
|
|
return cls(id=user.id, username=user.name, discriminator=user.discriminator)
|
|
except (NotFound, InvalidData, HTTPException, Forbidden) as e:
|
|
if e == Forbidden:
|
|
return cls(id=channel_id, name="Forbidden Channel")
|
|
return cls(id=channel_id, name="Deleted Channel")
|