import sqlite3 from datetime import datetime, timedelta from time import time 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.logger import logger from aurora.utilities.utils import generate_dict, get_next_case_number 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}" def update(self): from aurora.utilities.database import connect query = f"UPDATE moderation_{self.guild_id} SET timestamp = ?, moderation_type = ?, target_type = ?, moderator_id = ?, role_id = ?, duration = ?, end_timestamp = ?, reason = ?, resolved = ?, resolved_by = ?, resolve_reason = ?, expired = ?, changes = ?, metadata = ? WHERE moderation_id = ?;" with connect() as database: cursor = database.cursor() cursor.execute(query, ( self.timestamp, self.moderation_type, self.target_type, self.moderator_id, self.role_id, self.duration, self.end_timestamp, self.reason, self.resolved, self.resolved_by, self.resolve_reason, self.expired, self.changes, self.metadata, self.moderation_id )) cursor.close() logger.info("Updated moderation case %s in guild %s with the following data:\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s\n%s", self.moderation_id, self.guild_id, self.timestamp, self.moderation_type, self.target_type, self.moderator_id, self.role_id, self.duration, self.end_timestamp, self.reason, self.resolved, self.resolved_by, self.resolve_reason, self.expired, self.changes, self.metadata, ) @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) @classmethod def log( cls, bot: Red, guild_id: int, moderator_id: int, moderation_type: str, target_type: str, target_id: int, role_id: int, duration: timedelta = None, reason: str = None, database: sqlite3.Connection = None, timestamp: datetime = None, resolved: bool = False, resolved_by: int = None, resolved_reason: str = None, expired: bool = None, changes: list = None, metadata: dict = None, ) -> "Moderation": from aurora.utilities.database import connect from aurora.utilities.json import dumps if not timestamp: timestamp = datetime.fromtimestamp(time()) elif not isinstance(timestamp, datetime): timestamp = datetime.fromtimestamp(timestamp) if duration != "NULL" and duration is not None: end_timestamp = timestamp + duration else: duration = None end_timestamp = None if not expired: if timestamp > end_timestamp: expired = True else: expired = False if reason == "NULL": reason = None if resolved_by == "NULL": resolved_by = None if resolved_reason == "NULL": resolved_reason = None if role_id == 0: role_id = None if not database: database = connect() close_db = True else: close_db = False cursor = database.cursor() moderation_id = get_next_case_number(guild_id=guild_id, cursor=cursor) case = { "guild_id": guild_id, "moderation_id": moderation_id, "timestamp": timestamp, "moderation_type": moderation_type, "target_type": target_type, "target_id": target_id, "moderator_id": moderator_id, "role_id": role_id, "duration": duration, "end_timestamp": end_timestamp, "reason": reason, "resolved": resolved, "resolved_by": resolved_by, "resolve_reason": resolved_reason, "expired": expired, "changes": changes, "metadata": metadata } case_safe = case.copy() case_safe.pop("guild_id") case_safe["duration"] = str(case_safe["duration"]) if case_safe["duration"] else None case_safe["timestamp"] = case_safe["timestamp"].timestamp() case_safe["end_timestamp"] = case_safe["end_timestamp"].timestamp() if case_safe["end_timestamp"] else None case_safe["changes"] = dumps(case_safe["changes"]) case_safe["metadata"] = dumps(case_safe["metadata"]) sql = f"INSERT INTO `moderation_{guild_id}` (moderation_id, timestamp, moderation_type, target_type, target_id, moderator_id, role_id, duration, end_timestamp, reason, resolved, resolved_by, resolve_reason, expired, changes, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)" cursor.execute(sql, tuple(case_safe.values())) cursor.close() database.commit() if close_db: database.close() case_safe.update({"guild_id": guild_id}) logger.debug( "Row inserted into moderation_%s!\n%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s", guild_id, case_safe["moderation_id"], case_safe["timestamp"], case_safe["moderation_type"], case_safe["target_type"], case_safe["target_id"], case_safe["moderator_id"], case_safe["role_id"], case_safe["duration"], case_safe["end_timestamp"], case_safe["reason"], case_safe["resolved"], case_safe["resolved_by"], case_safe["resolve_reason"], case_safe["expired"], case_safe["changes"], case_safe["metadata"], ) for change in case["changes"]: change.update({"bot": bot}) return cls.from_dict(bot=bot, data=case) 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 "duration" in data and data["duration"] and not isinstance(data["duration"], timedelta): hours, minutes, seconds = map(int, data["duration"].split(':')) duration = timedelta(hours=hours, minutes=minutes, seconds=seconds) elif "duration" in data and isinstance(data["duration"], timedelta): duration = data["duration"] else: duration = None if "end_timestamp" in data and data["end_timestamp"] and not isinstance(data["end_timestamp"], datetime): end_timestamp = datetime.fromtimestamp(data["end_timestamp"]) elif "end_timestamp" in data and isinstance(data["end_timestamp"], datetime): end_timestamp = data["end_timestamp"] else: end_timestamp = None if not isinstance(data["timestamp"], datetime): timestamp = datetime.fromtimestamp(data["timestamp"]) else: timestamp = data["timestamp"] data.update({ "timestamp": timestamp, "end_timestamp": end_timestamp, "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)