from datetime import datetime, timedelta from typing import Any, ClassVar, Dict, List, Literal, 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.""" def to_json(self, indent: int = None, file: Any = None, **kwargs): from aurora.utilities.json import dump, dumps return dump(self.model_dump(), file, indent=indent, **kwargs) if file else dumps(self.model_dump(), indent=indent, **kwargs) 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 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(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, **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 Change(AuroraBaseModel): type: Literal["ORIGINAL", "RESOLVE", "EDIT"] timestamp: datetime reason: str user_id: int duration: Optional[timedelta] end_timestamp: Optional[datetime] def __str__(self): return f"{self.type} {self.user_id} {self.reason}" @classmethod def from_dict(cls, data: dict) -> "Change": return cls(**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(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") class PartialRole(AuroraBaseModel): id: int guild_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(id=role_id, guild_id=guild_id, name="Forbidden Role") role = guild.get_role(role_id) if not role: return cls(id=role_id, guild_id=guild_id, name="Deleted Role") return cls(id=role.id, guild_id=guild_id, name=role.name)