from datetime import datetime, timedelta from typing import Any, Dict, List, Literal, Optional, Union from discord import Forbidden, HTTPException, InvalidData, NotFound from pydantic import BaseModel, ConfigDict from redbot.core.bot import Red from aurora.utilities.utils import generate_dict 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): moderation_id: int timestamp: datetime 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 resolved: bool resolved_by: Optional[int] = None resolve_reason: Optional[str] = None expired: bool changes: List["Change"] metadata: Dict @property def id(self) -> int: return self.moderation_id @property def type(self) -> str: return self.moderation_type @property def unix_timestamp(self) -> int: return int(self.timestamp.timestamp()) 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 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(bot, 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) class Change(AuroraBaseModel): type: Literal["ORIGINAL", "RESOLVE", "EDIT"] timestamp: datetime reason: str user_id: int duration: Optional[timedelta] = None end_timestamp: Optional[datetime] = None @property def unix_timestamp(self) -> int: return int(self.timestamp.timestamp()) 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": if data["duration"] is not None: hours, minutes, seconds = map(int, data["duration"].split(':')) duration = timedelta(hours=hours, minutes=minutes, seconds=seconds) else: duration = None data.update({ "timestamp": datetime.fromtimestamp(data["timestamp"]), "end_timestamp": datetime.fromtimestamp(data["end_timestamp"]) if data["end_timestamp"] else None, "duration": duration }) 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) return cls(bot=bot, id=user.id, username=user.name, discriminator=user.discriminator) 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") return cls(bot=bot, guild_id=channel.guild.id, id=channel.id, username=channel.name) 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)