SeaCogs/aurora/models.py

184 lines
6.2 KiB
Python
Raw Normal View History

from datetime import datetime, timedelta
from typing import Any, Dict, List, Literal, Optional, Union
2024-05-04 13:41:11 -04:00
from discord import Forbidden, HTTPException, InvalidData, NotFound
from pydantic import BaseModel, ConfigDict
from redbot.core.bot import Red
2024-05-04 13:42:58 -04:00
from aurora.utilities.utils import generate_dict
2024-05-04 13:41:11 -04:00
class AuroraBaseModel(BaseModel):
"""Base class for all models in Aurora."""
model_config = ConfigDict(ignored_types=(Red,), arbitrary_types_allowed=True)
bot: Red
def to_json(self, indent: int = None, file: Any = None, **kwargs):
from aurora.utilities.json import dump, dumps
return dump(self.model_dump(exclude={"bot"}), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot"}), indent=indent, **kwargs)
class AuroraGuildModel(AuroraBaseModel):
"""Subclass of AuroraBaseModel that includes a guild_id attribute, and a modified to_json() method to match."""
guild_id: int
def to_json(self, indent: int = None, file: Any = None, **kwargs):
from aurora.utilities.json import dump, dumps
return dump(self.model_dump(exclude={"bot", "guild_id"}), file, indent=indent, **kwargs) if file else dumps(self.model_dump(exclude={"bot", "guild_id"}), indent=indent, **kwargs)
class Moderation(AuroraGuildModel):
2024-05-04 13:41:11 -04:00
moderation_id: int
timestamp: datetime
2024-05-04 13:41:11 -04:00
moderation_type: str
target_type: str
target_id: int
moderator_id: int
role_id: Optional[int] = None
duration: Optional[timedelta] = None
end_timestamp: Optional[datetime] = None
reason: Optional[str] = None
2024-05-04 13:41:11 -04:00
resolved: bool
resolved_by: Optional[int] = None
resolve_reason: Optional[str] = None
2024-05-04 13:41:11 -04:00
expired: bool
changes: List["Change"]
2024-05-04 13:41:11 -04:00
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
async def get_role(self) -> Optional["PartialRole"]:
if self.role_id:
return await PartialRole.from_id(self.bot, self.guild_id, self.role_id)
return None
2024-05-04 13:41:11 -04:00
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 = ?;"
2024-05-04 13:48:57 -04:00
with connect() as database:
cursor = database.cursor()
cursor.execute(query, (moderation_id,))
result = cursor.fetchone()
if result:
case = generate_dict(bot, result, guild_id)
2024-05-04 13:48:57 -04:00
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)
class Change(AuroraBaseModel):
type: Literal["ORIGINAL", "RESOLVE", "EDIT"]
timestamp: datetime
reason: str
user_id: int
duration: Optional[timedelta] = None
end_timestamp: Optional[datetime] = None
def __str__(self):
return f"{self.type} {self.user_id} {self.reason}"
async def get_user(self) -> "PartialUser":
return await PartialUser.from_id(self.bot, self.user_id)
@classmethod
def from_dict(cls, bot: Red, data: dict) -> "Change":
2024-05-04 18:18:57 -04:00
return cls(bot=bot, **data)
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(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator)
except NotFound:
return cls(bot=bot, id=user_id, username="Deleted User", discriminator=0)
class PartialChannel(AuroraGuildModel):
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":
channel = bot.get_channel(channel_id)
if not channel:
try:
channel = await bot.fetch_channel(channel_id)
return cls(bot=bot, guild_id=channel.guild.id, id=channel.id, username=channel.name)
except (NotFound, InvalidData, HTTPException, Forbidden) as e:
if e == Forbidden:
return cls(bot=bot, guild_id=0, id=channel_id, name="Forbidden Channel")
return cls(bot=bot, guild_id=0, id=channel_id, name="Deleted Channel")
class PartialRole(AuroraGuildModel):
id: int
name: str
@property
def mention(self):
return f"<@&{self.id}>"
def __str__(self):
return self.mention
@classmethod
async def from_id(cls, bot: Red, guild_id: int, role_id: int) -> "PartialRole":
try:
guild = await bot.fetch_guild(guild_id, with_counts=False)
except (Forbidden, HTTPException):
return cls(bot=bot, guild_id=guild_id, id=role_id, name="Forbidden Role")
role = guild.get_role(role_id)
if not role:
return cls(bot=bot, guild_id=guild_id, id=role_id, name="Deleted Role")
return cls(bot=bot, guild_id=guild_id, id=role.id, name=role.name)