SeaCogs/aurora/models.py

366 lines
13 KiB
Python

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,
author_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
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_timedelta = timestamp + duration
end_timestamp = int(end_timedelta.timestamp())
else:
duration = None
end_timestamp = None
if not expired:
if timestamp.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": author_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["timestamp"] = case_safe["timestamp"].timestamp()
case_safe["end_timestamp"] = case_safe["end_timestamp"].timestamp() if case_safe["end_timestamp"] else None
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, case_safe.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_safe.values()
)
return cls.from_dict(bot=bot, **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)