SeaCogs/aurora/models.py

129 lines
4.1 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)
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")