diff --git a/aurora/aurora.py b/aurora/aurora.py index 6e2e2fc..309bbb7 100644 --- a/aurora/aurora.py +++ b/aurora/aurora.py @@ -20,8 +20,7 @@ from redbot.core import app_commands, commands, data_manager from redbot.core.app_commands import Choice from redbot.core.bot import Red from redbot.core.commands.converter import parse_relativedelta, parse_timedelta -from redbot.core.utils.chat_formatting import (box, error, humanize_list, - humanize_timedelta, warning) +from redbot.core.utils.chat_formatting import box, error, humanize_list, humanize_timedelta, warning from aurora.importers.aurora import ImportAuroraView from aurora.importers.galacticbot import ImportGalacticBotView @@ -29,20 +28,14 @@ from aurora.menus.addrole import Addrole from aurora.menus.guild import Guild from aurora.menus.immune import Immune from aurora.menus.overrides import Overrides -from aurora.models import Change, Moderation +from aurora.models.change import Change +from aurora.models.moderation import Moderation from aurora.utilities.config import config, register_config from aurora.utilities.database import connect, create_guild_table -from aurora.utilities.factory import (addrole_embed, case_factory, - changes_factory, evidenceformat_factory, - guild_embed, immune_embed, - message_factory, overrides_embed) - +from aurora.utilities.factory import addrole_embed, case_factory, changes_factory, evidenceformat_factory, guild_embed, immune_embed, message_factory, overrides_embed from aurora.utilities.json import dump from aurora.utilities.logger import logger -from aurora.utilities.utils import (check_moddable, check_permissions, - get_footer_image, log, send_evidenceformat, - timedelta_from_relativedelta) - +from aurora.utilities.utils import check_moddable, check_permissions, get_footer_image, log, send_evidenceformat, timedelta_from_relativedelta class Aurora(commands.Cog): """Aurora is a fully-featured moderation system. diff --git a/aurora/importers/aurora.py b/aurora/importers/aurora.py index 0885ee9..7d1b273 100644 --- a/aurora/importers/aurora.py +++ b/aurora/importers/aurora.py @@ -8,7 +8,7 @@ from discord import ButtonStyle, Interaction, Message, ui from redbot.core import commands from redbot.core.utils.chat_formatting import box, warning -from aurora.models import Moderation +from aurora.models.moderation import Moderation from aurora.utilities.database import connect, create_guild_table diff --git a/aurora/importers/galacticbot.py b/aurora/importers/galacticbot.py index 4cdc42e..50d51b9 100644 --- a/aurora/importers/galacticbot.py +++ b/aurora/importers/galacticbot.py @@ -7,7 +7,7 @@ from discord import ButtonStyle, Interaction, Message, ui from redbot.core import commands from redbot.core.utils.chat_formatting import box, warning -from aurora.models import Change, Moderation +from aurora.models.moderation import Change, Moderation from aurora.utilities.database import connect, create_guild_table diff --git a/aurora/models/base.py b/aurora/models/base.py new file mode 100644 index 0000000..31d01ac --- /dev/null +++ b/aurora/models/base.py @@ -0,0 +1,24 @@ +from typing import Any + +from pydantic import BaseModel, ConfigDict +from redbot.core.bot import Red + + +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 ( # pylint: disable=cyclic-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 ( # pylint: disable=cyclic-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) diff --git a/aurora/models/change.py b/aurora/models/change.py new file mode 100644 index 0000000..ceb4641 --- /dev/null +++ b/aurora/models/change.py @@ -0,0 +1,62 @@ +import json +from datetime import datetime, timedelta +from typing import Literal, Optional + +from redbot.core.bot import Red + +from aurora.models.base import AuroraBaseModel +from aurora.models.partials import PartialUser +from aurora.utilities.logger import logger + + +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.trace("Creating Change from dict (%s): %s", type(data), data) + if isinstance(data, str): + data = json.loads(data) + logger.trace("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) diff --git a/aurora/models.py b/aurora/models/moderation.py similarity index 65% rename from aurora/models.py rename to aurora/models/moderation.py index ab63be1..7bea736 100644 --- a/aurora/models.py +++ b/aurora/models/moderation.py @@ -2,34 +2,19 @@ import json import sqlite3 from datetime import datetime, timedelta from time import time -from typing import Any, Dict, Iterable, List, Literal, Optional, Union +from typing import Dict, Iterable, List, Optional, Union import discord -from discord import Forbidden, HTTPException, InvalidData, NotFound -from pydantic import BaseModel, ConfigDict +from discord import NotFound from redbot.core.bot import Red +from aurora.models.base import AuroraGuildModel +from aurora.models.change import Change +from aurora.models.partials import PartialChannel, PartialRole, PartialUser 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 # pylint: disable=cyclic-import - 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 # pylint: disable=cyclic-import - 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 @@ -335,129 +320,3 @@ class Moderation(AuroraGuildModel): ) 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.trace("Creating Change from dict (%s): %s", type(data), data) - if isinstance(data, str): - data = json.loads(data) - logger.trace("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 in ["Deleted Channel", "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 in ["Deleted Role", "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) diff --git a/aurora/models/partials.py b/aurora/models/partials.py new file mode 100644 index 0000000..48ff2a0 --- /dev/null +++ b/aurora/models/partials.py @@ -0,0 +1,79 @@ +from discord import Forbidden, HTTPException, InvalidData, NotFound +from redbot.core.bot import Red + +from aurora.models.base import AuroraBaseModel, AuroraGuildModel + + +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 in ["Deleted Channel", "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 in ["Deleted Role", "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) diff --git a/aurora/utilities/database.py b/aurora/utilities/database.py index 588ad35..16f8a7e 100644 --- a/aurora/utilities/database.py +++ b/aurora/utilities/database.py @@ -5,7 +5,7 @@ import sqlite3 from discord import Guild from redbot.core import data_manager -from .logger import logger +from aurora.utilities.logger import logger def connect() -> sqlite3.Connection: diff --git a/aurora/utilities/factory.py b/aurora/utilities/factory.py index 11d00b9..1b508bc 100644 --- a/aurora/utilities/factory.py +++ b/aurora/utilities/factory.py @@ -6,7 +6,8 @@ from discord import Color, Embed, Guild, Interaction, InteractionMessage, Member from redbot.core import commands from redbot.core.utils.chat_formatting import bold, box, error, humanize_timedelta, warning -from aurora.models import Moderation, PartialUser +from aurora.models.moderation import Moderation +from aurora.models.partials import PartialUser from aurora.utilities.config import config from aurora.utilities.utils import get_bool_emoji, get_next_case_number, get_pagesize_str diff --git a/aurora/utilities/json.py b/aurora/utilities/json.py index 554679d..73595ee 100644 --- a/aurora/utilities/json.py +++ b/aurora/utilities/json.py @@ -3,7 +3,7 @@ from datetime import datetime, timedelta from redbot.core.bot import Red -from aurora.models import AuroraBaseModel +from aurora.models.base import AuroraBaseModel class JSONEncoder(json.JSONEncoder): diff --git a/aurora/utilities/utils.py b/aurora/utilities/utils.py index ca9cc85..e85fd58 100644 --- a/aurora/utilities/utils.py +++ b/aurora/utilities/utils.py @@ -125,7 +125,7 @@ def get_next_case_number(guild_id: str, cursor=None) -> int: async def log(interaction: Interaction, moderation_id: int, resolved: bool = False) -> None: """This function sends a message to the guild's configured logging channel when an infraction takes place.""" - from aurora.models import Moderation + from aurora.models.moderation import Moderation from aurora.utilities.factory import log_factory logging_channel_id = await config.guild(interaction.guild).log_channel() @@ -147,7 +147,7 @@ async def log(interaction: Interaction, moderation_id: int, resolved: bool = Fal async def send_evidenceformat(interaction: Interaction, moderation_id: int) -> None: """This function sends an ephemeral message to the moderator who took the moderation action, with a pre-made codeblock for use in the mod-evidence channel.""" - from aurora.models import Moderation + from aurora.models.moderation import Moderation from aurora.utilities.factory import evidenceformat_factory send_evidence_bool = (