372 lines
13 KiB
Python
372 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
|
|
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_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
|
|
case_safe["duration"] = str(case_safe["duration"])
|
|
case_safe["resolved"] = int(case_safe["resolved"])
|
|
case_safe["expired"] = int(case_safe["expired"])
|
|
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, 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)
|