import json import sqlite3 from datetime import datetime, timedelta from time import time from typing import Any, Dict, Iterable, List, Literal, Optional, Union import discord 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 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}" async def resolve(self, resolved_by: int, reason: str): if self.resolved: raise ValueError("Case is already resolved!") self.resolved = True self.resolved_by = resolved_by self.resolve_reason = reason if self.type in ["UNMUTE", "UNBAN"]: raise TypeError("Cannot resolve an unmute or unban case!") if self.type == "MUTE": try: guild: discord.Guild = await self.bot.fetch_guild(self.guild_id) member = await guild.fetch_member(self.target_id) await member.timeout( None, reason=f"Case {self.moderation_id} resolved by {resolved_by}{' for '+ reason if reason else ''}" ) except NotFound: pass if self.type in ["BAN", "TEMPBAN"]: try: guild: discord.Guild = await self.bot.fetch_guild(self.guild_id) await guild.unban(await self.get_target(), reason=f"Case {self.moderation_id} resolved by {resolved_by}{' for '+ reason if reason else ''}") except NotFound: pass if not self.changes: self.changes.append(Change.from_dict(self.bot, { "type": "ORIGINAL", "timestamp": self.timestamp, "reason": self.reason, "user_id": self.moderator_id, "duration": self.duration, "end_timestamp": self.end_timestamp, })) self.changes.append(Change.from_dict(self.bot, { "type": "RESOLVE", "timestamp": datetime.now(), "reason": reason, "user_id": resolved_by, })) self.update() def update(self): from aurora.utilities.database import connect from aurora.utilities.json import dumps 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.timestamp(), self.moderation_type, self.target_type, self.moderator_id, self.role_id, str(self.duration) if self.duration else None, self.end_timestamp.timestamp() if self.end_timestamp else None, self.reason, self.resolved, self.resolved_by, self.resolve_reason, self.expired, dumps(self.changes), dumps(self.metadata), self.moderation_id, )) cursor.close() logger.debug("Row updated in moderation_%s!\n%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s", self.moderation_id, self.guild_id, self.timestamp.timestamp(), self.moderation_type, self.target_type, self.moderator_id, self.role_id, str(self.duration) if self.duration else None, self.end_timestamp.timestamp() if self.end_timestamp else None, self.reason, self.resolved, self.resolved_by, self.resolve_reason, self.expired, dumps(self.changes), dumps(self.metadata), ) @classmethod def from_dict(cls, bot: Red, data: dict) -> "Moderation": return cls(bot=bot, **data) @classmethod def from_result(cls, bot: Red, result: Iterable, guild_id: int) -> "Moderation": if result[7] is not None: hours, minutes, seconds = map(int, result[7].split(':')) duration = timedelta(hours=hours, minutes=minutes, seconds=seconds) else: duration = None if result[14] is not None: changes = json.loads(result[14]) change_obj_list = [] for change in changes: change_obj_list.append(Change.from_dict(bot=bot, data=change)) case = { "moderation_id": int(result[0]), "guild_id": int(guild_id), "timestamp": datetime.fromtimestamp(result[1]), "moderation_type": str(result[2]), "target_type": str(result[3]), "target_id": int(result[4]), "moderator_id": int(result[5]), "role_id": int(result[6]) if result[6] is not None else None, "duration": duration, "end_timestamp": datetime.fromtimestamp(result[8]) if result[8] is not None else None, "reason": result[9], "resolved": bool(result[10]), "resolved_by": result[11], "resolve_reason": result[12], "expired": bool(result[13]), "changes": change_obj_list if result[14] else [], "metadata": json.loads(result[15].replace('\\"', '"').replace('["{', '[{').replace('}"]', '}]')) if result[15] else {}, } return cls.from_dict(bot=bot, data=case) @classmethod def from_sql(cls, bot: Red, moderation_id: int, guild_id: int) -> "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() cursor.close() if result and not moderation_id == 0: return cls.from_result(bot, result, guild_id) raise ValueError(f"Case {moderation_id} not found in moderation_{guild_id}!") @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 = { "moderation_id": moderation_id, "timestamp": timestamp.timestamp(), "moderation_type": moderation_type, "target_type": target_type, "target_id": target_id, "moderator_id": moderator_id, "role_id": role_id, "duration": str(duration) if duration else None, "end_timestamp": end_timestamp.timestamp() if end_timestamp else None, "reason": reason, "resolved": resolved, "resolved_by": resolved_by, "resolve_reason": resolved_reason, "expired": expired, "changes": dumps(changes), "metadata": dumps(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.values())) cursor.close() database.commit() if close_db: database.close() 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["moderation_id"], case["timestamp"], case["moderation_type"], case["target_type"], case["target_id"], case["moderator_id"], case["role_id"], case["duration"], case["end_timestamp"], case["reason"], case["resolved"], case["resolved_by"], case["resolve_reason"], case["expired"], case["changes"], case["metadata"], ) return cls.from_sql(bot=bot, moderation_id=moderation_id, guild_id=guild_id) 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": logger.debug("Creating Change from dict (%s): %s", type(data), data) if isinstance(data, str): data = json.loads(data) logger.debug("Change data was a string, converted to dict: %s", data) 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, "user_id": int(data["user_id"]) }) 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): if self.name == "Deleted Role" or self.name == "Forbidden Role": return self.name 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)