diff --git a/jarvis/__init__.py b/jarvis/__init__.py index 9129aa5..5c568dc 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -7,14 +7,15 @@ import pymongo from discord import DMChannel, Intents, Member, Message from discord.ext import commands from discord.ext.tasks import loop -from discord.utils import find, get +from discord.utils import find from discord_slash import SlashCommand from psutil import Process from jarvis import logo, utils from jarvis.config import get_config +from jarvis.db import DBManager +from jarvis.db.types import Autoreact, Ban, Lock, Mute, Setting, Warning from jarvis.utils import build_embed -from jarvis.utils.db import DBManager from jarvis.utils.field import Field if asyncio.get_event_loop().is_closed(): @@ -35,10 +36,11 @@ jarvis = commands.Bot( ) slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True) jarvis_self = Process() -__version__ = "1.0.1" +__version__ = "1.1.0" db = DBManager(get_config().mongo).mongo +jarvis_db = db.jarvis @jarvis.event @@ -67,25 +69,16 @@ async def on_ready(): @jarvis.event async def on_member_join(user: Member): guild = user.guild - db = DBManager(get_config().mongo).mongo - mutes = list( - db.jarvis.mutes.find( - {"active": True, "user": user.id, "guild": guild.id} - ) - ) + mutes = Mute.get_active(guild=guild.id) if mutes and len(mutes) >= 1: - mute_role = db.jarvis.settings.find_one( - {"guild": guild.id, "setting": "mute"} - ) - role = guild.get_role(mute_role["value"]) + mute_role = Setting.get(guild=guild.id, setting="mute") + role = guild.get_role(mute_role.value) await user.add_roles( role, reason="User is muted still muted from prior mute" ) - unverified = db.jarvis.settings.find_one( - {"guild": guild.id, "setting": "unverified"} - ) + unverified = Setting.get(guild=guild.id, setting="unverified") if unverified: - role = guild.get_role(unverified["value"]) + role = guild.get_role(unverified.value) await user.add_roles(role, reason="User just joined and is unverified") @@ -95,32 +88,32 @@ async def on_message(message: Message): not isinstance(message.channel, DMChannel) and message.author.id != jarvis.user.id ): - autoreact = db.jarvis.autoreact.find_one( - {"guild": message.guild.id, "channel": message.channel.id} + autoreact = Autoreact.get( + guild=message.guild.id, + channel=message.channel.id, ) if autoreact: - for reaction in autoreact["reactions"]: + for reaction in autoreact.reactions: await message.add_reaction(reaction) - massmention = db.jarvis.settings.find_one( - {"guild": message.guild.id, "setting": "massmention"} + massmention = Setting.get( + guild=message.guild.id, + setting="massmention", ) if ( - massmention["value"] > 0 + massmention.value > 0 and len(message.mentions) - (1 if message.author in message.mentions else 0) - > massmention["value"] + > massmention.value ): - db.jarvis.warns.insert_one( - { - "user": message.author.id, - "reason": "Mass Mention", - "admin": get_config().client_id, - "time": datetime.now(), - "guild": message.guild.id, - "duration": 24, - "active": True, - } + warning = Warning( + active=True, + admin=get_config().client_id, + duration=24, + guild=message.guild.id, + reason="Mass Mention", + user=message.author.id, ) + warning.insert() fields = [Field("Reason", "Mass Mention", False)] embed = build_embed( title="Warning", @@ -138,9 +131,7 @@ async def on_message(message: Message): + f"| {message.author.id}" ) await message.channel.send(embed=embed) - roleping = db.jarvis.settings.find_one( - {"guild": message.guild.id, "setting": "roleping"} - ) + roleping = Setting.get(guild=message.guild.id, setting="roleping") roles = [] for mention in message.role_mentions: roles.append(mention.id) @@ -149,22 +140,18 @@ async def on_message(message: Message): roles.append(role.id) if ( roleping - and any(x in roleping["value"] for x in roles) - and not any( - x.id in roleping["value"] for x in message.author.roles - ) + and any(x in roleping.value for x in roles) + and not any(x.id in roleping.value for x in message.author.roles) ): - db.jarvis.warns.insert_one( - { - "user": message.author.id, - "reason": "Pinged a blocked role/user with a blocked role", - "admin": get_config().client_id, - "time": datetime.now(), - "guild": message.guild.id, - "duration": 24, - "active": True, - } + warning = Warning( + active=True, + admin=get_config().client_id, + duration=24, + guild=message.guild.id, + reason="Pinged a blocked role/user with a blocked role", + user=message.author.id, ) + warning.insert() fields = [ Field( "Reason", @@ -188,11 +175,9 @@ async def on_message(message: Message): + f"| {message.author.id}" ) await message.channel.send(embed=embed) - autopurge = db.jarvis.autopurge.find_one( - {"guild": message.guild.id, "channel": message.channel.id} - ) + autopurge = Setting.get(guild=message.guild.id, setting="autopurge") if autopurge: - await message.delete(delay=autopurge["delay"]) + await message.delete(delay=autopurge.delay) content = re.sub(r"\s+", "", message.content) match = invites.search(content) if match: @@ -203,17 +188,15 @@ async def on_message(message: Message): ] if match.group(1) not in allowed: await message.delete() - db.jarvis.warns.insert_one( - { - "user": message.author.id, - "reason": "Sent an invite link", - "admin": get_config().client_id, - "time": datetime.now(), - "guild": message.guild.id, - "duration": 24, - "active": True, - } + warning = Warning( + active=True, + admin=get_config().client_id, + duration=24, + guild=message.guild.id, + reason="Sent an invite link", + user=message.author.id, ) + warning.insert() fields = [ Field( "Reason", @@ -255,56 +238,57 @@ async def on_guild_join(guild): await general.send("Importing all preferences from home interface...") # Set some default settings - db = DBManager(get_config().mongo).mongo - db.jarvis.settings.insert_one( - {"guild": guild.id, "setting": "massmention", "value": 5} - ) + setting = Setting(guild=guild.id, setting="massmention", value=5) + setting.insert() await general.send("Systems are now fully operational") @loop(minutes=1) async def unmute(): - db = DBManager(get_config().mongo).mongo - mutes = list(db.jarvis.mutes.find({"active": True, "length": {"$gt": 0}})) - mute_roles = list( - db.jarvis.settings.find({"setting": "mute"}).sort( - [("guild", pymongo.ASCENDING)] - ) - ) + mutes = Mute.get_active(duration={"$gt": 0}) + mute_roles = Setting.get(setting="mute") updates = [] for mute in mutes: - if mute["time"] + timedelta(minutes=mute["length"]) < datetime.now(): - mute_role = [ - x["value"] for x in mute_roles if x["guild"] == mute["guild"] - ][0] - guild = await jarvis.fetch_guild(mute["guild"]) + if ( + mute.created_at + timedelta(minutes=mute.duration) + < datetime.utcnow() + ): + mute_role = [x.value for x in mute_roles if x.guild == mute.guild][ + 0 + ] + guild = await jarvis.fetch_guild(mute.guild) role = guild.get_role(mute_role) - user = await guild.fetch_member(mute["user"]) + user = await guild.fetch_member(mute.user) if user: if role in user.roles: await user.remove_roles(role, reason="Mute expired") + + # Objects can't handle bulk_write, so handle it via raw methods updates.append( pymongo.UpdateOne( - {"user": user.id, "guild": guild.id, "time": mute["time"]}, + { + "user": user.id, + "guild": guild.id, + "created_at": mute.created_at, + }, {"$set": {"active": False}}, ) ) if updates: - db.jarvis.mutes.bulk_write(updates) + jarvis_db.mutes.bulk_write(updates) @loop(minutes=10) async def unban(): - db = DBManager(get_config().mongo).mongo - bans = list(db.jarvis.bans.find({"active": True, "type": "temp"})) + bans = Ban.get_active(type="temp") updates = [] for ban in bans: - if ban["time"] + timedelta( - hours=ban["length"] - ) < datetime.now() + timedelta(minutes=10): - guild = await jarvis.fetch_guild(ban["guild"]) - user = await jarvis.fetch_user(ban["user"]) + if ban.created_at + timedelta( + hours=ban.duration + ) < datetime.utcnow() + timedelta(minutes=10): + guild = await jarvis.fetch_guild(ban.guild) + user = await jarvis.fetch_user(ban.user) if user: guild.unban(user) updates.append( @@ -312,25 +296,27 @@ async def unban(): { "user": user.id, "guild": guild.id, - "time": ban["time"], + "created_at": ban.created_at, "type": "temp", }, {"$set": {"active": False}}, ) ) if updates: - db.jarvis.bans.bulk_write(updates) + jarvis_db.bans.bulk_write(updates) @loop(minutes=1) async def unlock(): - db = DBManager(get_config().mongo).mongo - locks = list(db.jarvis.locks.find({"active": True})) + locks = Lock.get_active() updates = [] for lock in locks: - if lock["time"] + timedelta(minutes=lock["duration"]) < datetime.now(): - guild = await jarvis.fetch_guild(lock["guild"]) - channel = await jarvis.fetch_channel(lock["channel"]) + if ( + lock.created_at + timedelta(minutes=lock.duration) + < datetime.utcnow() + ): + guild = await jarvis.fetch_guild(lock.guild) + channel = await jarvis.fetch_channel(lock.channel) if channel: roles = await guild.fetch_roles() for role in roles: @@ -344,29 +330,31 @@ async def unlock(): { "channel": channel.id, "guild": guild.id, - "time": lock["time"], + "created_at": lock.created_at, }, {"$set": {"active": False}}, ) ) if updates: - db.jarvis.locks.bulk_write(updates) + jarvis_db.locks.bulk_write(updates) @loop(hours=1) async def unwarn(): - db = DBManager(get_config().mongo).mongo - warns = list(db.jarvis.warns.find({"active": True})) + warns = Warning.get_active() updates = [] for warn in warns: - if warn["time"] + timedelta(hours=warn["duration"]) < datetime.now(): + if ( + warn.created_at + timedelta(hours=warn.duration) + < datetime.utcnow() + ): updates.append( pymongo.UpdateOne( - {"_id": warn["_id"]}, {"$set": {"active": False}} + {"_id": warn._id}, {"$set": {"active": False}} ) ) if updates: - db.jarvis.warns.bulk_write(updates) + jarvis_db.warns.bulk_write(updates) def run(ctx=None): diff --git a/jarvis/cogs/admin.py b/jarvis/cogs/admin.py index 5b1ed64..84451c8 100644 --- a/jarvis/cogs/admin.py +++ b/jarvis/cogs/admin.py @@ -1,5 +1,5 @@ import re -from datetime import datetime, timedelta +from datetime import datetime from typing import Union import pymongo @@ -10,8 +10,19 @@ from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_choice, create_option import jarvis +from jarvis.db import DBManager +from jarvis.db.types import ( + Autopurge, + Ban, + Kick, + Lock, + Mute, + Purge, + Setting, + Unban, + Warning, +) from jarvis.utils import build_embed -from jarvis.utils.db import DBManager from jarvis.utils.field import Field from jarvis.utils.permissions import admin_or_permissions @@ -26,7 +37,7 @@ class AdminCog(commands.Cog): def __init__(self, bot: commands.Bot): self.bot = bot config = jarvis.config.get_config() - self.db = DBManager(config.mongo).mongo + self.db = DBManager(config.mongo).mongo.jarvis @cog_ext.cog_slash( name="ban", @@ -83,6 +94,11 @@ class AdminCog(commands.Cog): "You cannot set a temp ban to < 0 hours.", hidden=True ) return + elif type == "temp" and duration > 744: + await ctx.send( + "You cannot set a temp ban to > 1 month", hidden=True + ) + return if len(reason) > 100: await ctx.send("Reason must be < 100 characters", hidden=True) return @@ -100,11 +116,8 @@ class AdminCog(commands.Cog): f"You have been {mtype}banned from {guild_name}." + f" Reason:\n{reason}" ) - time = datetime.now() - expiry = None if mtype == "temp": user_message += f"\nDuration: {duration} hours" - expiry = time + timedelta(hours=duration) try: await ctx.guild.ban(user, reason=reason) @@ -124,41 +137,30 @@ class AdminCog(commands.Cog): if type == "soft": active = False - self.db.jarvis.bans.insert_one( - { - "user": user.id, - "username": user.name, - "discrim": user.discriminator, - "reason": reason, - "admin": ctx.author.id, - "time": datetime.now(), - "guild": ctx.guild.id, - "type": type, - "duration": duration, - "expiry": expiry, - "active": active, - } - ) + _ = Ban( + user=user.id, + username=user.name, + discrim=user.discriminator, + reason=reason, + admin=ctx.author.id, + guild=ctx.guild.id, + type=type, + duration=duration, + active=active, + ).insert() async def discord_apply_unban( self, ctx: SlashContext, user: User, reason: str ): await ctx.guild.unban(user, reason=reason) - self.db.jarvis.unbans.insert_one( - { - "user": user.id, - "username": user.name, - "discrim": user.discriminator, - "guild": ctx.guild.id, - "admin": ctx.author.id, - "reason": reason, - "time": datetime.now(), - } - ) - _ = self.db.jarvis.bans.update( - {"user": user.id, "guild": ctx.guild.id}, - {"$set": {"active": False}}, - ) + _ = Unban( + user=user.id, + username=user.name, + discrim=user.discriminator, + guild=ctx.guild.id, + admin=ctx.author.id, + reason=reason, + ).insert() await ctx.send("User successfully unbanned.\nReason: " + reason) @cog_ext.cog_slash( @@ -239,8 +241,8 @@ class AdminCog(commands.Cog): # We take advantage of the previous checks to save CPU cycles if not discord_ban_info: if isinstance(user, int): - database_ban_info = self.db.jarvis.bans.find_one( - {"guild": ctx.guild.id, "user": user, "active": True} + database_ban_info = Ban.get( + guild=ctx.guild.id, user=user, active=True ) else: search = { @@ -250,7 +252,7 @@ class AdminCog(commands.Cog): } if discrim: search["discrim"] = discrim - database_ban_info = self.db.jarvis.bans.find_one(search) + database_ban_info = Ban.get(**search) if not discord_ban_info and not database_ban_info: await ctx.send(f"Unable to find user {orig_user}", hidden=True) @@ -266,21 +268,16 @@ class AdminCog(commands.Cog): ctx, discord_ban_info.user, reason ) else: - self.db.jarvis.bans.update_many( - {"user": database_ban_info["id"], "guild": ctx.guild.id}, - {"$set": {"active": False}}, - ) - self.db.jarvis.unbans.insert_one( - { - "user": database_ban_info["user"], - "username": database_ban_info["username"], - "discrim": database_ban_info["discrim"], - "guild": ctx.guild.id, - "admin": ctx.author.id, - "reason": reason, - "time": datetime.now(), - } - ) + database_ban_info.active = False + database_ban_info.update() + _ = Unban( + user=database_ban_info.user, + username=database_ban_info.username, + discrim=database_ban_info.discrim, + guild=ctx.guild.id, + admin=ctx.author.id, + reason=reason, + ).insert() await ctx.send( "Unable to find user in Discord, " + "but removed entry from database." @@ -326,24 +323,23 @@ class AdminCog(commands.Cog): search["active"] = True if type > 0: search["type"] = types[type] - bans = self.db.jarvis.bans.find(search).sort( - [("time", pymongo.DESCENDING)] - ) + bans = Ban.get_many(**search) + bans.sort(key=lambda x: x.created_at, reverse=True) ban_messages = [] db_bans = [] for ban in bans: - if "username" not in ban: - user = await self.bot.fetch_user(ban["user"]) - ban["username"] = user.name if user else "[deleted user]" + if not ban.username: + user = await self.bot.fetch_user(ban.user) + ban.username = user.name if user else "[deleted user]" ban_messages.append( "[{0}] {1} ({2}): {3}".format( - ban["time"].strftime("%d-%m-%Y"), - ban["username"], - ban["user"], - ban["reason"], + ban.created_at.strftime("%d-%m-%Y"), + ban.username, + ban.user, + ban.reason, ) ) - db_bans.append(ban["user"]) + db_bans.append(ban.user) bans = await ctx.guild.bans() for ban in bans: if ban.user.id not in db_bans: @@ -405,15 +401,12 @@ class AdminCog(commands.Cog): f"{user.name} has been kicked from {guild_name}." + f"Reason:\n{reason}" ) - self.db.jarvis.kicks.insert_one( - { - "user": user.id, - "reason": reason, - "admin": ctx.authod.id, - "time": datetime.now(), - "guild": ctx.guild.id, - } - ) + _ = Kick( + user=user.id, + reason=reason, + admin=ctx.author.id, + guild=ctx.guild.id, + ).insert() @cog_ext.cog_slash( name="purge", @@ -438,15 +431,12 @@ class AdminCog(commands.Cog): async for message in channel.history(limit=amount + 1): messages.append(message) await channel.delete_messages(messages) - self.db.jarvis.purges.insert_one( - { - "channel": ctx.channel.id, - "guild": ctx.guild.id, - "admin": ctx.author.id, - "count": amount, - "time": datetime.now(), - } - ) + _ = Purge( + channel=ctx.channel.id, + guild=ctx.guild.id, + admin=ctx.author.id, + count=amount, + ).insert() @cog_ext.cog_slash( name="mute", @@ -485,45 +475,29 @@ class AdminCog(commands.Cog): if len(reason) > 100: await ctx.send("Reason must be < 100 characters", hidden=True) return - mute_setting = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "mute"} - ) + mute_setting = Setting.get(guild=ctx.guild.id, setting="mute") if not mute_setting: await ctx.send( - "Please configure a mute role with /settings mute first", + "Please configure a mute role " + + "with /settings mute first", hidden=True, ) return - role = get(ctx.guild.roles, id=mute_setting["value"]) + role = get(ctx.guild.roles, id=mute_setting.value) if role in user.roles: await ctx.send("User already muted", hidden=True) + return await user.add_roles(role, reason=reason) - time = datetime.now() - expiry = None - if duration < 0: + if duration < 0 or duration > 300: duration = -1 - if duration >= 0: - expiry = time + timedelta(minutes=duration) - self.db.jarvis.mutes.insert_one( - { - "user": user.id, - "reason": reason, - "admin": ctx.author.id, - "time": time, - "guild": ctx.guild.id, - "duration": duration, - "expiry": expiry, - "active": True if duration >= 0 else False, - } - ) - self.db.jarvis.mutes.update_many( - { - "guild": ctx.guild.id, - "user": user.id, - "expiry": {"$lt": expiry}, - }, - {"$set": {"active": False}}, - ) + _ = Mute( + user=user.id, + reason=reason, + admin=ctx.author.id, + guild=ctx.guild.id, + duration=duration, + active=True if duration >= 0 else False, + ).insert() await ctx.send(f"{user.mention} has been muted.\nReason: {reason}") @cog_ext.cog_slash( @@ -540,9 +514,7 @@ class AdminCog(commands.Cog): ) @admin_or_permissions(mute_members=True) async def _unmute(self, ctx: SlashContext, user: Member): - mute_setting = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "mute"} - ) + mute_setting = Setting.get(guild=ctx.guild.id, setting="mute") if not mute_setting: await ctx.send( "Please configure a mute role with " @@ -551,20 +523,17 @@ class AdminCog(commands.Cog): ) return - role = get(ctx.guild.roles, id=mute_setting["value"]) + role = get(ctx.guild.roles, id=mute_setting.value) if role in user.roles: await user.remove_roles(role, reason="Unmute") else: await ctx.send("User is not muted.", hidden=True) return - self.db.jarvis.mutes.update_many( - { - "guild": ctx.guild.id, - "user": user.id, - }, - {"$set": {"active": False}}, - ) + mutes = Mute.get_many(guild=ctx.guild.id, user=user.id) + for mute in mutes: + mute.active = False + mute.update() await ctx.send(f"{user.mention} has been unmuted.") async def _lock_channel( @@ -628,6 +597,12 @@ class AdminCog(commands.Cog): channel: Union[TextChannel, VoiceChannel] = None, ): await ctx.defer(hidden=True) + if duration <= 0: + await ctx.send("Duration must be > 0", hidden=True) + return + elif duration >= 300: + await ctx.send("Duration must be < 5 hours", hidden=True) + return if len(reason) > 100: await ctx.send("Reason must be < 100 characters", hidden=True) return @@ -638,17 +613,13 @@ class AdminCog(commands.Cog): await self._lock_channel(channel, role, ctx.author, reason) except Exception: continue # Just continue on error - self.db.jarvis.locks.insert_one( - { - "channel": channel.id, - "guild": ctx.guild.id, - "admin": ctx.author.id, - "reason": reason, - "duration": duration, - "active": True, - "time": datetime.now(), - } - ) + _ = Lock( + channel=channel.id, + guild=ctx.guild.id, + admin=ctx.author.id, + reason=reason, + duration=duration, + ).insert() await ctx.send(f"{channel.mention} locked for {duration} minute(s)") @cog_ext.cog_slash( @@ -671,9 +642,7 @@ class AdminCog(commands.Cog): ): if not channel: channel = ctx.channel - lock = self.db.jarvis.locks.find_one( - {"guild": ctx.guild.id, "channel": channel.id, "active": True} - ) + lock = Lock.get(guild=ctx.guild.id, channel=channel.id, active=True) if not lock: await ctx.send(f"{channel.mention} not locked.", hidden=True) return @@ -682,13 +651,8 @@ class AdminCog(commands.Cog): await self._unlock_channel(channel, role, ctx.author) except Exception: continue # Just continue on error - self.db.jarvis.locks.update_one( - { - "channel": channel.id, - "guild": ctx.guild.id, - }, - {"$set": {"active": False}}, - ) + lock.active = False + lock.update() await ctx.send(f"{channel.mention} unlocked") @cog_ext.cog_subcommand( @@ -717,7 +681,13 @@ class AdminCog(commands.Cog): reason: str, duration: int = 10, ): - await ctx.defer() + await ctx.defer(hidden=True) + if duration <= 0: + await ctx.send("Duration must be > 0", hidden=True) + return + elif duration >= 300: + await ctx.send("Duration must be < 5 hours", hidden=True) + return channels = ctx.guild.channels roles = ctx.guild.roles updates = [] @@ -736,12 +706,12 @@ class AdminCog(commands.Cog): "reason": reason, "duration": duration, "active": True, - "time": datetime.now(), + "created_at": datetime.utcnow(), } ) ) if updates: - self.db.jarvis.locks.bulk_write(updates) + self.db.locks.bulk_write(updates) await ctx.send(f"Server locked for {duration} minute(s)") @cog_ext.cog_subcommand( @@ -757,9 +727,7 @@ class AdminCog(commands.Cog): channels = ctx.guild.channels roles = ctx.guild.roles updates = [] - locks = list( - self.db.jarvis.locks.find({"guild": ctx.guild.id, "active": True}) - ) + locks = Lock.get_many(guild=ctx.guild.id, active=True) if not locks: await ctx.send("No lockdown detected.", hidden=True) return @@ -781,7 +749,7 @@ class AdminCog(commands.Cog): ) ) if updates: - self.db.jarvis.locks.bulk_write(updates) + self.db.locks.bulk_write(updates) await ctx.send("Server unlocked") @cog_ext.cog_slash( @@ -815,25 +783,20 @@ class AdminCog(commands.Cog): if len(reason) > 100: await ctx.send("Reason must be < 100 characters", hidden=True) return + if duration <= 0: + await ctx.send("Duration must be > 0", hidden=True) + return + elif duration >= 120: + await ctx.send("Duration must be < 5 days", hidden=True) + return await ctx.defer() - self.db.jarvis.warns.insert_one( - { - "user": user.id, - "reason": reason, - "admin": ctx.author.id, - "time": datetime.now(), - "guild": ctx.guild.id, - "duration": duration, - "active": True, - } - ) - count = len( - list( - self.db.jarvis.warns.find( - {"user": user.id, "guild": ctx.guild.id, "active": True} - ) - ) - ) + _ = Warning( + user=user.id, + reason=reason, + admin=ctx.author.id, + guild=ctx.guild.id, + duration=duration, + ).insert() fields = [Field("Reason", reason, False)] embed = build_embed( title="Warning", @@ -863,18 +826,11 @@ class AdminCog(commands.Cog): @commands.has_permissions(administrator=True) async def _warnings(self, ctx: SlashContext, user: User): await ctx.defer() - warnings = list( - self.db.jarvis.warns.find( - { - "user": user.id, - "guild": ctx.guild.id, - } - ) - ) + warnings = Warning.get_many(user=user.id, guild=ctx.guild.id) active = ( [ - f'`{y["time"].strftime("%Y-%m-%d %H:%M:%S")}` - {y["reason"]}' - for y in list(filter(lambda x: x["active"], warnings)) + f'`{y.created_at.strftime("%Y-%m-%d %H:%M:%S")}` - {y.reason}' + for y in list(filter(lambda x: x.active, warnings)) ] if warnings else ["None"] @@ -918,23 +874,17 @@ class AdminCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _roleping_block(self, ctx: SlashContext, role: Role): - roles = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "roleping"} - ) + roles = Setting.get(guild=ctx.guild.id, setting="roleping") if not roles: - roles = {"guild": ctx.guild.id, "setting": "roleping", "value": []} + roles = Setting(guild=ctx.guild.id, setting="roleping", value=[]) - if role.id in roles["value"]: + if role.id in roles.value: await ctx.send( f"Role `{role.name}` already in blocklist.", hidden=True ) return - roles["value"].append(role.id) - self.db.jarvis.settings.update_one( - {"guild": ctx.guild.id, "setting": "roleping"}, - {"$set": roles}, - upsert=True, - ) + roles.value.append(role.id) + roles.update() await ctx.send(f"Role `{role.name}` added to blocklist.") @cog_ext.cog_subcommand( @@ -952,24 +902,18 @@ class AdminCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _roleping_allow(self, ctx: SlashContext, role: Role): - roles = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "roleping"} - ) + roles = Setting.get(guild=ctx.guild.id, setting="roleping") if not roles: await ctx.send("No blocklist configured.", hidden=True) return - if role.id not in roles["value"]: + if role.id not in roles.value: await ctx.send( f"Role `{role.name}` not in blocklist.", hidden=True ) return - roles["value"].delete(role.id) - self.db.jarvis.settings.update_one( - {"guild": ctx.guild.id, "setting": "roleping"}, - {"$set": roles}, - upsert=True, - ) + roles.value.delete(role.id) + roles.update() await ctx.send(f"Role `{role.name}` removed blocklist.") @cog_ext.cog_subcommand( @@ -978,18 +922,16 @@ class AdminCog(commands.Cog): description="List all blocklisted roles", ) async def _roleping_list(self, ctx: SlashContext): - roles = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "roleping"} - ) + roles = Setting.get(guild=ctx.guild.id, setting="roleping") if not roles: await ctx.send("No blocklist configured.", hidden=True) return message = "Blocklisted Roles:\n```\n" - if not roles["value"]: + if not roles.value: await ctx.send("No roles blocklisted.", hidden=True) return - for role in roles["value"]: + for role in roles.value: role = ctx.guild.get_role(role) if not role: continue @@ -1023,20 +965,23 @@ class AdminCog(commands.Cog): if not isinstance(channel, TextChannel): await ctx.send("Channel must be a TextChannel", hidden=True) return - autopurge = self.db.jarvis.autopurge.find( - {"guild": ctx.guild.id, "channel": channel.id} - ) + if delay <= 0: + await ctx.send("Delay must be > 0", hidden=True) + return + elif delay > 300: + await ctx.send("Delay must be < 5 minutes", hidden=True) + return + autopurge = Autopurge(guild=ctx.guild.id, channel=channel.id) if autopurge: await ctx.send("Autopurge already exists.", hidden=True) return - autopurge = { - "guild": ctx.guild.id, - "channel": channel.id, - "admin": ctx.author.id, - "delay": delay, - "time": datetime.utcnow(), - } - self.db.jarvis.autopurge.insert_one(autopurge) + autopurge = Autopurge( + guild=ctx.guild.id, + channel=channel.id, + admin=ctx.author.id, + delay=delay, + ) + autopurge.insert() await ctx.send( f"Autopurge set up on {channel.mention}, " + f"delay is {delay} seconds" @@ -1057,13 +1002,11 @@ class AdminCog(commands.Cog): ) @admin_or_permissions(manage_messages=True) async def _autopurge_remove(self, ctx: SlashContext, channel: TextChannel): - autopurge = self.db.jarvis.autopurge.find( - {"guild": ctx.guild.id, "channel": channel.id} - ) + autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id) if not autopurge: await ctx.send("Autopurge does not exist.", hidden=True) return - self.db.jarvis.autopurge.delete_one({"_id": autopurge["_id"]}) + autopurge.delete() await ctx.send(f"Autopurge removed from {channel.mention}.") @cog_ext.cog_subcommand( @@ -1089,15 +1032,12 @@ class AdminCog(commands.Cog): async def _autopurge_update( self, ctx: SlashContext, channel: TextChannel, delay: int ): - autopurge = self.db.jarvis.autopurge.find_one( - {"guild": ctx.guild.id, "channel": channel.id} - ) + autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id) if not autopurge: await ctx.send("Autopurge does not exist.", hidden=True) return - self.db.jarvis.autopurge.update_one( - {"_id": autopurge["_id"]}, {"$set": {"delay": delay}} - ) + autopurge.delay = delay + autopurge.update() await ctx.send( f"Autopurge delay updated to {delay} seconds on {channel.mention}." ) diff --git a/jarvis/cogs/autoreact.py b/jarvis/cogs/autoreact.py index 4708cdd..ac83e11 100644 --- a/jarvis/cogs/autoreact.py +++ b/jarvis/cogs/autoreact.py @@ -1,5 +1,4 @@ import re -from datetime import datetime from discord import TextChannel from discord.ext import commands @@ -7,17 +6,13 @@ from discord.utils import find from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -import jarvis -from jarvis.config import get_config from jarvis.data.unicode import emoji_list -from jarvis.utils.db import DBManager +from jarvis.db.types import Autoreact class AutoReactCog(commands.Cog): def __init__(self, bot): self.bot = bot - config = get_config() - self.db = DBManager(config.mongo).mongo self.custom_emote = re.compile(r"^<:\w+:(\d+)>$") @cog_ext.cog_subcommand( @@ -38,23 +33,20 @@ class AutoReactCog(commands.Cog): if not isinstance(channel, TextChannel): await ctx.send("Channel must be a text channel", hidden=True) return - exists = self.db.jarvis.autoreact.find_one( - {"guild": ctx.guild.id, "channel": channel.id} - ) + exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) if exists: await ctx.send( f"Autoreact already exists for {channel.mention}.", hidden=True ) return - autoreact = { - "guild": ctx.guild.id, - "channel": channel.id, - "reactions": [], - "admin": ctx.author.id, - "time": datetime.now(), - } - self.db.jarvis.autoreact.insert_one(autoreact) + autoreact = Autoreact( + guild=ctx.guild.id, + channel=channel.id, + reactions=[], + admin=ctx.author.id, + ) + autoreact.insert() await ctx.send(f"Autoreact created for {channel.mention}!") @cog_ext.cog_subcommand( @@ -71,9 +63,7 @@ class AutoReactCog(commands.Cog): ], ) async def _autoreact_delete(self, ctx, channel: TextChannel): - exists = self.db.jarvis.autoreact.delete_one( - {"guild": channel.guild.id, "channel": channel.id} - ) + exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id).delete() if exists: await ctx.send(f"Autoreact removed from {channel.mention}") else: @@ -119,32 +109,28 @@ class AutoReactCog(commands.Cog): "Please use a custom emote from this server.", hidden=True ) return - exists = self.db.jarvis.autoreact.find_one( - {"guild": ctx.guild.id, "channel": channel.id} - ) + exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) if not exists: await ctx.send( "Please create autoreact first with " + f"/autoreact create {channel.mention}" ) return - if emote in exists["reactions"]: + if emote in exists.reactions: await ctx.send( f"Emote already added to {channel.mention} autoreactions.", hidden=True, ) return - if len(exists["reactions"]) >= 5: + if len(exists.reactions) >= 5: await ctx.send( - "Max number of reactions hit. Remove a different one to add this one", + "Max number of reactions hit. " + + "Remove a different one to add this one", hidden=True, ) return - exists["reactions"].append(emote) - self.db.jarvis.autoreact.update_one( - {"_id": exists["_id"]}, - {"$set": {"reactions": exists["reactions"]}}, - ) + exists.reactions.append(emote) + exists.update() await ctx.send(f"Added {emote} to {channel.mention} autoreact.") @cog_ext.cog_subcommand( @@ -168,9 +154,7 @@ class AutoReactCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str): - exists = self.db.jarvis.autoreact.find_one( - {"guild": ctx.guild.id, "channel": channel.id} - ) + exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) if not exists: await ctx.send( "Please create autoreact first with " @@ -178,17 +162,14 @@ class AutoReactCog(commands.Cog): hidden=True, ) return - if emote not in exists["reactions"]: + if emote not in exists.reactions: await ctx.send( f"{emote} not used in {channel.mention} autoreactions.", hidden=True, ) return - exists["reactions"].remove(emote) - self.db.jarvis.autoreact.update_one( - {"_id": exists["_id"]}, - {"$set": {"reactions": exists["reactions"]}}, - ) + exists.reactions.remove(emote) + exists.update() await ctx.send(f"Removed {emote} from {channel.mention} autoreact.") @cog_ext.cog_subcommand( @@ -206,9 +187,7 @@ class AutoReactCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _autoreact_list(self, ctx, channel: TextChannel): - exists = self.db.jarvis.autoreact.find_one( - {"guild": ctx.guild.id, "channel": channel.id} - ) + exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) if not exists: await ctx.send( "Please create autoreact first with " @@ -217,10 +196,10 @@ class AutoReactCog(commands.Cog): ) return message = "" - if len(exists["reactions"]) > 0: + if len(exists.reactions) > 0: message = ( f"Current active autoreacts on {channel.mention}:\n" - + "\n".join(exists["reactions"]) + + "\n".join(exists.reactions) ) else: message = f"No reactions set on {channel.mention}" diff --git a/jarvis/cogs/ctc2.py b/jarvis/cogs/ctc2.py index 2a07c0a..bcd8fda 100644 --- a/jarvis/cogs/ctc2.py +++ b/jarvis/cogs/ctc2.py @@ -2,9 +2,8 @@ import aiohttp from discord.ext import commands from discord_slash import cog_ext -import jarvis from jarvis.config import get_config -from jarvis.utils.db import DBManager +from jarvis.db import DBManager guild_ids = [578757004059738142, 520021794380447745, 862402786116763668] diff --git a/jarvis/cogs/dbrand.py b/jarvis/cogs/dbrand.py index 2d9823f..1aa63e8 100644 --- a/jarvis/cogs/dbrand.py +++ b/jarvis/cogs/dbrand.py @@ -5,7 +5,6 @@ from discord.ext import commands from discord_slash import cog_ext from discord_slash.utils.manage_commands import create_option -import jarvis from jarvis.config import get_config from jarvis.data.dbrand import shipping_lookup from jarvis.utils import build_embed @@ -134,7 +133,8 @@ class DbrandCog(commands.Cog): ( create_option( name="search", - description="Country search query (2 character code, country name, emoji)", + description="Country search query (2 character code, " + + "country name, emoji)", option_type=3, required=True, ) @@ -203,10 +203,19 @@ class DbrandCog(commands.Cog): x for x in data["country"].split(" ") if x != "the" ) country_urlsafe = country.replace("-", "%20") - description = f"Click the link above to see shipping time to {data['country']}." - description += "\n[View all shipping destinations](https://dbrand.com/shipping)" + description = ( + "Click the link above to " + + "see shipping time to {data['country']}." + ) + description += ( + "\n[View all shipping destinations]" + + "(https://dbrand.com/shipping)" + ) description += " | [Check shipping status]" - description += f"(https://dbrand.com/status#main-content:~:text={country_urlsafe})" + description += ( + "(https://dbrand.com/status" + + f"#main-content:~:text={country_urlsafe})" + ) embed = build_embed( title="Shipping to {}".format(data["country"]), description=description, diff --git a/jarvis/cogs/dev.py b/jarvis/cogs/dev.py index cc1d31e..9c02717 100644 --- a/jarvis/cogs/dev.py +++ b/jarvis/cogs/dev.py @@ -4,13 +4,11 @@ import re import subprocess import uuid -import discord import ulid from bson import ObjectId from discord.ext import commands from discord_slash import cog_ext -import jarvis from jarvis.utils import build_embed, convert_bytesize from jarvis.utils.field import Field diff --git a/jarvis/cogs/error.py b/jarvis/cogs/error.py index 615dda5..dc80603 100644 --- a/jarvis/cogs/error.py +++ b/jarvis/cogs/error.py @@ -1,7 +1,5 @@ from discord.ext import commands -import jarvis - class ErrorHandlerCog(commands.Cog): def __init__(self, bot): diff --git a/jarvis/cogs/image.py b/jarvis/cogs/image.py index 6dddda1..fc69cdb 100644 --- a/jarvis/cogs/image.py +++ b/jarvis/cogs/image.py @@ -7,7 +7,6 @@ import numpy as np from discord import File from discord.ext import commands -import jarvis from jarvis.utils import build_embed, convert_bytesize, unconvert_bytesize from jarvis.utils.field import Field @@ -98,7 +97,7 @@ class ImageCog(commands.Cog): Field("Accuracy", f"{accuracy:.02f}%", False), ] embed = build_embed(title=filename, description="", fields=fields) - embed.set_image(url=f"attachment://resized.png") + embed.set_image(url="attachment://resized.png") await ctx.send( embed=embed, file=File(bufio, filename="resized.png"), diff --git a/jarvis/cogs/jokes.py b/jarvis/cogs/jokes.py index e1deabe..f3b1fd7 100644 --- a/jarvis/cogs/jokes.py +++ b/jarvis/cogs/jokes.py @@ -4,13 +4,12 @@ import traceback from datetime import datetime from random import randint -import discord from discord.ext import commands from discord_slash import cog_ext import jarvis +from jarvis.db import DBManager from jarvis.utils import build_embed -from jarvis.utils.db import DBManager from jarvis.utils.field import Field @@ -24,7 +23,7 @@ class JokeCog(commands.Cog): def __init__(self, bot): self.bot = bot config = jarvis.config.get_config() - self.db = DBManager(config.mongo) + self.db = DBManager(config.mongo).mongo.jarvis # TODO: Make this a command group with subcommands async def _joke(self, ctx, id: str = None): @@ -34,7 +33,7 @@ class JokeCog(commands.Cog): return # TODO: Add this as a parameter that can be passed in threshold = 500 # Minimum score - coll = self.db.mongo.jarvis.jokes + coll = self.db.jokes result = None if id: result = coll.find_one({"id": id}) @@ -124,10 +123,6 @@ class JokeCog(commands.Cog): ) # await ctx.send(f"**{result['title']}**\n\n{result['body']}") - @commands.command(name="joke", help="Hear a joke") - async def _joke_pref(self, ctx, id: str = None): - await self._joke(ctx, id) - @cog_ext.cog_slash( name="joke", description="Hear a joke", diff --git a/jarvis/cogs/modlog.py b/jarvis/cogs/modlog.py index 16f2a6f..a6ee596 100644 --- a/jarvis/cogs/modlog.py +++ b/jarvis/cogs/modlog.py @@ -2,16 +2,14 @@ import asyncio from datetime import datetime, timedelta import discord -import pymongo from discord import DMChannel from discord.ext import commands from discord.utils import find from discord_slash import SlashContext -import jarvis from jarvis.config import get_config +from jarvis.db.types import Ban, Kick, MongoSort, Mute, Setting from jarvis.utils import build_embed -from jarvis.utils.db import DBManager from jarvis.utils.field import Field @@ -22,7 +20,6 @@ class ModlogCog(commands.Cog): def __init__(self, bot: discord.ext.commands.Bot): self.bot = bot - self.db = DBManager(get_config().mongo).mongo def get_latest_log(self, auditlog, target): before = datetime.utcnow() - timedelta(seconds=10) @@ -66,11 +63,9 @@ class ModlogCog(commands.Cog): @commands.Cog.listener() async def on_member_ban(self, guild: discord.Guild, user: discord.User): - modlog = self.db.jarvis.settings.find_one( - {"guild": guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=guild.id, setting="modlog") if modlog: - channel = guild.get_channel(modlog["value"]) + channel = guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log auditlog = await guild.audit_logs( limit=50, @@ -81,15 +76,13 @@ class ModlogCog(commands.Cog): log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = self.db.jarvis.bans.find_one( - { - "guild": guild.id, - "user": user.id, - "active": True, - }, - sort=[("time", pymongo.DESCENDING)], + ban = Ban.get( + guild=guild.id, + user=user.id, + active=True, + sort=MongoSort(key="created_at", type="desc"), ) - admin = guild.get_member(mute["admin"]) + admin = guild.get_member(ban.admin) embed = await self.modlog_embed( user, admin, @@ -102,11 +95,9 @@ class ModlogCog(commands.Cog): @commands.Cog.listener() async def on_member_unban(self, guild: discord.Guild, user: discord.User): - modlog = self.db.jarvis.settings.find_one( - {"guild": guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=guild.id, setting="modlog") if modlog: - channel = guild.get_channel(modlog["value"]) + channel = guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log auditlog = await guild.audit_logs( limit=50, @@ -117,15 +108,13 @@ class ModlogCog(commands.Cog): log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = self.db.jarvis.bans.find_one( - { - "guild": guild.id, - "user": user.id, - "active": True, - }, - sort=[("time", pymongo.DESCENDING)], + ban = Ban.get( + guild=guild.id, + user=user.id, + active=True, + sort=MongoSort(key="created_at", type="desc"), ) - admin = guild.get_member(mute["admin"]) + admin = guild.get_member(ban.admin) embed = await self.modlog_embed( user, admin, @@ -138,11 +127,9 @@ class ModlogCog(commands.Cog): @commands.Cog.listener() async def on_member_remove(self, user: discord.User): - modlog = self.db.jarvis.settings.find_one( - {"guild": user.guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=user.guild.id, setting="modlog") if modlog: - channel = user.guild.get_channel(modlog["value"]) + channel = user.guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log auditlog = await user.guild.audit_logs( limit=50, @@ -153,14 +140,12 @@ class ModlogCog(commands.Cog): log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = self.db.jarvis.kicks.find_one( - { - "guild": user.guild.id, - "user": user.id, - }, - sort=[("time", pymongo.DESCENDING)], + kick = Kick.get( + guild=user.guild.id, + user=user.id, + sort=MongoSort(key="created_at", type="desc"), ) - admin = user.guild.get_member(mute["admin"]) + admin = user.guild.get_member(kick.admin) embed = await self.modlog_embed( user, admin, @@ -181,15 +166,13 @@ class ModlogCog(commands.Cog): log: discord.AuditLogEntry = self.get_latest_log(auditlog, before) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = self.db.jarvis.mutes.find_one( - { - "guild": before.guild.id, - "user": before.id, - "active": True, - }, - sort=[("time", pymongo.DESCENDING)], + mute = Mute.get( + guild=before.guild.id, + user=before.id, + active=True, + sort=MongoSort(key="created_at", type="desc"), ) - admin = before.guild.get_member(mute["admin"]) + admin = before.guild.get_member(mute.admin) return await self.modlog_embed( member=before, admin=admin, @@ -208,15 +191,14 @@ class ModlogCog(commands.Cog): log: discord.AuditLogEntry = self.get_latest_log(auditlog, before) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = self.db.jarvis.mutes.find_one( - { - "guild": before.guild.id, - "user": before.id, - "active": True, - }, - sort=[("time", pymongo.DESCENDING)], + mute = Mute.get( + guild=before.guild.id, + user=before.id, + active=True, + sort=MongoSort(key="created_at", type="desc"), ) - admin = before.guild.get_member(mute["admin"]) + mute = Mute(**mute) + admin = before.guild.get_member(mute.admin) return await self.modlog_embed( member=before, admin=admin, @@ -273,28 +255,21 @@ class ModlogCog(commands.Cog): async def on_member_update( self, before: discord.User, after: discord.User ): - modlog = self.db.jarvis.settings.find_one( - {"guild": after.guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=before.guild.id, setting="modlog") if modlog: channel = after.guild.get_channel(modlog["value"]) await asyncio.sleep(0.5) # Need to wait for audit log embed = None - mute = self.db.jarvis.settings.find_one( - {"guild": before.guild.id, "setting": "mute"} - ) - verified = self.db.jarvis.settings.find_one( - {"guild": before.guild.id, "setting": "verified"} - ) - if mute and before.guild.get_role(mute["value"]) in after.roles: + mute = Setting(guild=before.guild.id, setting="mute") + verified = Setting(guild=before.guild.id, setting="verified") + if mute and before.guild.get_role(mute.value) in after.roles: embed = await self.process_mute(before, after) - elif mute and before.guild.get_role(mute["value"]) in before.roles: + elif mute and before.guild.get_role(mute.value) in before.roles: embed = await self.process_unmute(before, after) elif ( verified - and before.guild.get_role(verified["value"]) - not in before.roles - and after.guild.get_role(verified["value"]) in after.roles + and before.guild.get_role(verified.value) not in before.roles + and after.guild.get_role(verified.value) in after.roles ): embed = await self.process_verify(before, after) elif before.nick != after.nick: @@ -358,13 +333,11 @@ class ModlogCog(commands.Cog): self, before: discord.Message, after: discord.Message ): if before.author != get_config().client_id: - modlog = self.db.jarvis.settings.find_one( - {"guild": after.guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=after.guild.id, setting="modlog") if modlog: if before.content == after.content or before.content is None: return - channel = before.guild.get_channel(modlog["value"]) + channel = before.guild.get_channel(modlog.value) fields = [ Field( "Original Message", @@ -398,12 +371,10 @@ class ModlogCog(commands.Cog): @commands.Cog.listener() async def on_message_delete(self, message: discord.Message): - modlog = self.db.jarvis.settings.find_one( - {"guild": message.guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=message.guild.id, setting="modlog") if modlog: fields = [Field("Original Message", message.content, False)] - channel = message.guild.get_channel(modlog["value"]) + channel = message.guild.get_channel(modlog.value) embed = build_embed( title="Message Deleted", description=f"{message.author.mention}'s message was deleted", @@ -424,11 +395,9 @@ class ModlogCog(commands.Cog): @commands.Cog.listener() async def on_slash_command(self, ctx: SlashContext): if not isinstance(ctx.channel, DMChannel): - modlog = self.db.jarvis.settings.find_one( - {"guild": ctx.guild.id, "setting": "modlog"} - ) + modlog = Setting.get(guild=ctx.guild.id, setting="modlog") if modlog: - channel = ctx.guild.get_channel(modlog["value"]) + channel = ctx.guild.get_channel(modlog.value) fields = [ Field("Command", ctx.name), ] diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py index 0677ef1..7ea6cb8 100644 --- a/jarvis/cogs/owner.py +++ b/jarvis/cogs/owner.py @@ -9,8 +9,9 @@ from discord import DMChannel, User from discord.ext import commands import jarvis -from jarvis.config import get_config, reload_config -from jarvis.utils import db, update +from jarvis.config import reload_config +from jarvis.db.types import Config +from jarvis.utils import update from jarvis.utils.permissions import user_is_bot_admin @@ -23,9 +24,7 @@ class OwnerCog(commands.Cog): def __init__(self, bot): self.bot = bot - self.config = get_config() - self.admins = self.config.admins - self.db = db.DBManager(self.config.mongo) + self.admins = Config.get(key="admins") @commands.command(name="load", hidden=True) @user_is_bot_admin() @@ -96,7 +95,7 @@ class OwnerCog(commands.Cog): async def _system(self, ctx): if ctx.invoked_subcommand is None: await ctx.send( - f"Usage: `system `\n" + "Usage: `system `\n" + "Subcommands: `restart`, `update`" ) @@ -173,13 +172,11 @@ class OwnerCog(commands.Cog): @_admin.command(name="add", hidden=True) @commands.is_owner() async def _add(self, ctx, user: User): - if user.id in self.admins: + if user.id in self.admins.value: await ctx.send(f"{user.mention} is already an admin.") return - self.admins.append(user.id) - self.db.mongo.jarvis.config.update_one( - {"key": "admins"}, {"$set": {"value": self.admins}} - ) + self.admins.value.append(user.id) + self.admins.update() reload_config() await ctx.send( f"{user.mention} is now an admin. Use this power carefully." @@ -188,88 +185,84 @@ class OwnerCog(commands.Cog): @_admin.command(name="remove", hidden=True) @commands.is_owner() async def _remove(self, ctx, user: User): - if user.id not in self.admins: + if user.id not in self.admins.value: await ctx.send(f"{user.mention} is not an admin.") return - self.admins.remove(user.id) - self.db.mongo.jarvis.config.update_one( - {"key": "admins"}, {"$set": {"value": self.admins}} - ) + self.admins.value.remove(user.id) + self.admins.update() reload_config() await ctx.send(f"{user.mention} is no longer an admin.") - def resolve_variable(self, variable): - if hasattr(variable, "__iter__"): - var_length = len(list(variable)) - if (var_length > 100) and (not isinstance(variable, str)): - return f"" - elif not var_length: - return f"" + def resolve_variable(self, variable): + if hasattr(variable, "__iter__"): + var_length = len(list(variable)) + if (var_length > 100) and (not isinstance(variable, str)): + return ( + f"" + ) + elif not var_length: + return f"" - if (not variable) and (not isinstance(variable, bool)): - return f"" - return ( - variable - if (len(f"{variable}") <= 1000) - else f"" - ) - - def prepare(self, string): - arr = ( - string.strip("```") - .replace("py\n", "") - .replace("python\n", "") - .split("\n") - ) - if not arr[::-1][0].replace(" ", "").startswith("return"): - arr[len(arr) - 1] = "return " + arr[::-1][0] - return "".join(f"\n\t{i}" for i in arr) - - @commands.command( - pass_context=True, aliases=["eval", "exec", "evaluate"] + if (not variable) and (not isinstance(variable, bool)): + return f"" + return ( + variable + if (len(f"{variable}") <= 1000) + else f"" ) - @user_is_bot_admin() - async def _eval(self, ctx, *, code: str): - if not isinstance(ctx.message.channel, DMChannel): + + def prepare(self, string): + arr = ( + string.strip("```") + .replace("py\n", "") + .replace("python\n", "") + .split("\n") + ) + if not arr[::-1][0].replace(" ", "").startswith("return"): + arr[len(arr) - 1] = "return " + arr[::-1][0] + return "".join(f"\n\t{i}" for i in arr) + + @commands.command(pass_context=True, aliases=["eval", "exec", "evaluate"]) + @user_is_bot_admin() + async def _eval(self, ctx, *, code: str): + if not isinstance(ctx.message.channel, DMChannel): + return + code = self.prepare(code) + args = { + "discord": discord, + "sauce": getsource, + "sys": sys, + "os": os, + "imp": __import__, + "this": self, + "ctx": ctx, + } + + try: + exec( + f"async def func():{code}", + globals().update(args), + locals(), + ) + a = time() + response = await eval("func()", globals().update(args), locals()) + if response is None or isinstance(response, discord.Message): + del args, code return - code = self.prepare(code) - args = { - "discord": discord, - "sauce": getsource, - "sys": sys, - "os": os, - "imp": __import__, - "this": self, - "ctx": ctx, - } - try: - exec( - f"async def func():{code}", - globals().update(args), - locals(), - ) - a = time() - response = await eval( - "func()", globals().update(args), locals() - ) - if response is None or isinstance(response, discord.Message): - del args, code - return + if isinstance(response, str): + response = response.replace("`", "") - if isinstance(response, str): - response = response.replace("`", "") + await ctx.send( + f"```py\n{self.resolve_variable(response)}```" + + f"`{type(response).__name__} | {(time() - a) / 1000} ms`" + ) + except Exception: + await ctx.send(f"Error occurred:```\n{traceback.format_exc()}```") - await ctx.send( - f"```py\n{self.resolve_variable(response)}```" - + f"`{type(response).__name__} | {(time() - a) / 1000} ms`" - ) - except Exception: - await ctx.send( - f"Error occurred:```\n{traceback.format_exc()}```" - ) - - del args, code + del args, code def setup(bot): diff --git a/jarvis/cogs/settings.py b/jarvis/cogs/settings.py index f285fdc..215e0c8 100644 --- a/jarvis/cogs/settings.py +++ b/jarvis/cogs/settings.py @@ -1,29 +1,20 @@ -import discord from discord import Role, TextChannel from discord.ext import commands from discord_slash import cog_ext from discord_slash.utils.manage_commands import create_option -import jarvis -from jarvis.config import get_config -from jarvis.utils import db +from jarvis.db.types import Setting from jarvis.utils.permissions import admin_or_permissions class SettingsCog(commands.Cog): def __init__(self, bot): self.bot = bot - config = get_config() - self.db = db.DBManager(config.mongo).mongo self.cache = {} def update_settings(self, setting, value, guild): - settings = {"setting": setting, "value": value, "guild": guild} - updated = self.db.jarvis.settings.update_one( - {"setting": setting, "guild": guild}, - {"$set": settings}, - upsert=True, - ) + setting = Setting(setting=setting, value=value, guild=guild) + updated = setting.update() return updated is not None diff --git a/jarvis/cogs/starboard.py b/jarvis/cogs/starboard.py index 1649c2a..d2a95e3 100644 --- a/jarvis/cogs/starboard.py +++ b/jarvis/cogs/starboard.py @@ -1,19 +1,10 @@ -from datetime import datetime - -import aiohttp -import discord -from discord import Message, TextChannel +from discord import TextChannel from discord.ext import commands -from discord.utils import find from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -import jarvis -from jarvis.config import get_config +from jarvis.db.types import Star, Starboard from jarvis.utils import build_embed -from jarvis.utils.db import DBManager -from jarvis.utils.field import Field -from jarvis.utils.permissions import admin_or_permissions supported_images = [ "image/png", @@ -27,7 +18,6 @@ supported_images = [ class StarboardCog(commands.Cog): def __init__(self, bot): self.bot = bot - self.db = DBManager(get_config().mongo).mongo @cog_ext.cog_subcommand( base="starboard", @@ -36,13 +26,11 @@ class StarboardCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _list(self, ctx): - starboards = [ - x for x in self.db.jarvis.starboard.find({"guild": ctx.guild.id}) - ] + starboards = Starboard.get_many(guild=ctx.guild.id) if starboards != []: message = "Available Starboards:\n" for s in starboards: - message += f"<#{s['target']}>\n" + message += f"<#{s.channel}>\n" await ctx.send(message) else: await ctx.send("No Starboards available.") @@ -53,42 +41,39 @@ class StarboardCog(commands.Cog): description="Create a starboard", options=[ create_option( - name="target", - description="Target channel", + name="channel", + description="Starboard channel", option_type=7, required=True, ), ], ) @commands.has_permissions(administrator=True) - async def _create(self, ctx, target: TextChannel): - if target not in ctx.guild.channels: + async def _create(self, ctx, channel: TextChannel): + if channel not in ctx.guild.channels: await ctx.send( - "Target channel not in guild. Choose an existing channel.", + "Channel not in guild. Choose an existing channel.", hidden=True, ) return - if not isinstance(target, TextChannel): - await ctx.send("Target must be a TextChannel", hidden=True) + if not isinstance(channel, TextChannel): + await ctx.send("Channel must be a TextChannel", hidden=True) return - exists = self.db.jarvis.starboard.find_one( - {"target": target.id, "guild": ctx.guild.id} - ) + exists = Starboard.get(channel=channel.id, guild=ctx.guild.id) if exists: await ctx.send( - f"Starboard already exists at {target.mention}.", hidden=True + f"Starboard already exists at {channel.mention}.", hidden=True ) return - self.db.jarvis.starboard.insert_one( - { - "guild": ctx.guild.id, - "target": target.id, - "admin": ctx.author.id, - "time": datetime.now(), - } + _ = Starboard( + guild=ctx.guild.id, + channel=channel.id, + admin=ctx.author.id, + ).insert() + await ctx.send( + f"Starboard created. Check it out at {channel.mention}." ) - await ctx.send(f"Starboard created. Check it out at {target.mention}.") @cog_ext.cog_subcommand( base="starboard", @@ -105,14 +90,11 @@ class StarboardCog(commands.Cog): ) @commands.has_permissions(administrator=True) async def _delete(self, ctx, channel: TextChannel): - deleted = self.db.jarvis.starboard.delete_one( - { - "target": channel.id, - "guild": ctx.guild.id, - } - ) + deleted = Starboard.get( + channel=channel.id, guild=ctx.guild.id + ).delete() if deleted: - self.db.jarvis.stars.delete_many({"starboard": channel.id}) + _ = Star.delete_many(starboard=channel.id) await ctx.send( f"Starboard deleted from {channel.mention}.", hidden=True ) @@ -157,9 +139,7 @@ class StarboardCog(commands.Cog): ): if not channel: channel = ctx.channel - exists = self.db.jarvis.starboard.find_one( - {"target": starboard.id, "guild": ctx.guild.id} - ) + exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id) if not exists: await ctx.send( f"Starboard does not exist in {starboard.mention}. " @@ -170,13 +150,11 @@ class StarboardCog(commands.Cog): message = await channel.fetch_message(int(message)) - exists = self.db.jarvis.stars.find_one( - { - "message": message.id, - "channel": message.channel.id, - "guild": message.guild.id, - "starboard": starboard.id, - } + exists = Star.get( + message=message.id, + channel=message.channel.id, + guild=message.guild.id, + starboard=starboard.id, ) if exists: @@ -186,9 +164,9 @@ class StarboardCog(commands.Cog): ) return - count = self.db.jarvis.stars.find( - {"guild": message.guild.id, "starboard": starboard.id} - ).count() + count = len( + Star.get_many(guild=message.guild.id, starboard=starboard.id) + ) content = message.content attachments = message.attachments @@ -221,18 +199,15 @@ class StarboardCog(commands.Cog): star = await starboard.send(embed=embed) - self.db.jarvis.stars.insert_one( - { - "index": count, - "message": message.id, - "channel": message.channel.id, - "guild": message.guild.id, - "starboard": starboard.id, - "admin": ctx.author.id, - "time": datetime.now(), - "star": star.id, - } - ) + _ = Star( + index=count, + message=message.id, + channel=message.channel.id, + guild=message.guild.id, + starboard=starboard.id, + admin=ctx.author.id, + star=star.id, + ).insert() await ctx.send( "Message saved to Starboard.\n" + f"See it in {starboard.mention}" @@ -268,9 +243,7 @@ class StarboardCog(commands.Cog): id: int, starboard: TextChannel, ): - exists = self.db.jarvis.starboard.find_one( - {"target": starboard.id, "guild": ctx.guild.id} - ) + exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id) if not exists: await ctx.send( f"Starboard does not exist in {starboard.mention}. " @@ -279,26 +252,22 @@ class StarboardCog(commands.Cog): ) return - star = self.db.jarvis.stars.find_one( - { - "starboard": starboard.id, - "id": id, - "guild": ctx.guild.id, - "active": True, - } + star = Star.get( + starboard=starboard.id, + id=id, + guild=ctx.guild.id, + active=True, ) if not star: await ctx.send(f"No star exists with id {id}", hidden=True) return - message = await starboard.fetch_message(star["star"]) + message = await starboard.fetch_message(star.star) if message: await message.delete() - self.db.jarvis.stars.update_one( - {"starboard": starboard.id, "id": id, "guild": ctx.guild.id}, - {"$set": {"active": False}}, - ) + star.active = False + star.update() await ctx.send(f"Star {id} deleted") diff --git a/jarvis/cogs/verify.py b/jarvis/cogs/verify.py index 83d4d9f..8391218 100644 --- a/jarvis/cogs/verify.py +++ b/jarvis/cogs/verify.py @@ -1,15 +1,11 @@ from random import randint -from discord import Role from discord.ext import commands from discord_slash import ComponentContext, SlashContext, cog_ext from discord_slash.model import ButtonStyle -from discord_slash.utils import manage_commands, manage_components -from discord_slash.utils.manage_commands import create_option +from discord_slash.utils import manage_components -import jarvis -from jarvis.config import get_config -from jarvis.utils.db import DBManager +from jarvis.db.types import Setting def create_layout(): @@ -33,7 +29,6 @@ def create_layout(): class VerifyCog(commands.Cog): def __init__(self, bot): self.bot = bot - self.db = DBManager(get_config().mongo).mongo @cog_ext.cog_slash( name="verify", @@ -41,15 +36,13 @@ class VerifyCog(commands.Cog): ) async def _verify(self, ctx: SlashContext): await ctx.defer() - role = self.db.jarvis.settings.find_one( - {"setting": "verified", "guild": ctx.guild.id} - ) + role = Setting.get(guild=ctx.guild.id, setting="verified") if not role: await ctx.send( "This guild has not enabled verification", delete_after=5 ) return - if ctx.guild.get_role(role["value"]) in ctx.author.roles: + if ctx.guild.get_role(role.value) in ctx.author.roles: await ctx.send("You are already verified.", delete_after=5) return components = create_layout() @@ -74,21 +67,18 @@ class VerifyCog(commands.Cog): for c in components: for c2 in c["components"]: c2["disabled"] = True - setting = self.db.jarvis.settings.find_one( - {"setting": "verified", "guild": ctx.guild.id} - ) - role = ctx.guild.get_role(setting["value"]) + setting = Setting.get(guild=ctx.guild.id, setting="verified") + role = ctx.guild.get_role(setting.value) await ctx.author.add_roles(role, reason="Verification passed") - setting = self.db.jarvis.settings.find_one( - {"setting": "unverified", "guild": ctx.guild.id} - ) + setting = Setting.get(guild=ctx.guild.id, setting="unverified") if setting: - role = ctx.guild.get_role(setting["value"]) + role = ctx.guild.get_role(setting.value) await ctx.author.remove_roles( role, reason="Verification passed" ) await ctx.edit_origin( - content=f"Welcome, {ctx.author.mention}. Please enjoy your stay.", + content=f"Welcome, {ctx.author.mention}. " + + "Please enjoy your stay.", components=manage_components.spread_to_rows( *components, max_in_row=5 ), diff --git a/jarvis/config.py b/jarvis/config.py index 5ce6e86..0ba4eb3 100644 --- a/jarvis/config.py +++ b/jarvis/config.py @@ -1,6 +1,6 @@ from yaml import load -from jarvis.utils.db import DBManager +from jarvis.db import DBManager try: from yaml import CLoader as Loader diff --git a/jarvis/utils/db.py b/jarvis/db/__init__.py similarity index 100% rename from jarvis/utils/db.py rename to jarvis/db/__init__.py diff --git a/jarvis/db/types.py b/jarvis/db/types.py new file mode 100644 index 0000000..ea65425 --- /dev/null +++ b/jarvis/db/types.py @@ -0,0 +1,539 @@ +import logging +from dataclasses import asdict, dataclass, field +from datetime import datetime +from typing import Any, Optional + +from bson import ObjectId +from pymongo import ASCENDING, DESCENDING +from pymongo.collection import Collection + +from jarvis.config import get_config +from jarvis.db import DBManager + +logger = logging.getLogger("mongodb") + +sort_lookup = { + "asc": ASCENDING, + "ascending": ASCENDING, + "desc": DESCENDING, + "descending": DESCENDING, +} + +coll_lookup = { + "Autoreact": "autoreact", + "Ban": "bans", + "Config": "config", + "Lock": "locks", + "Kick": "kicks", + "Mute": "mutes", + "Purge": "purges", + "Setting": "settings", + "Starboard": "starboard", + "Star": "stars", + "Unban": "unbans", + "Warning": "warns", +} + +db_instance = DBManager(get_config().mongo).mongo.narvis + + +################# +# Core Classes # +################# + + +@dataclass +class MongoSort: + direction: int + key: Any + + def __init__(self, direction, key): + if type(direction) is str: + direction = sort_lookup[direction] + self.direction = direction + self.key = key + + def as_tuple(self): + return (self.key, self.direction) + + +@dataclass +class MongoObject: + """ + A MongoDB object + + :param _id: MongoDB ObjectId, not initialized by default + """ + + _id: ObjectId = field(default=None, init=False) + + def __post_init__(self): + self.coll = db_instance[coll_lookup[type(self).__name__]] + + def to_dict(self): + return asdict(self) + + def insert(self) -> ObjectId: + """ + Inserts the object into the database + + :param collection: Collection to insert object into + :return: Inserted object ObjectId + """ + to_insert = self.to_dict() + if not to_insert["_id"]: + _ = to_insert.pop("_id") + try: + result = self.coll.insert_one(to_insert) + except Exception as e: + logger.error(f"Failed to insert {type(self).__name__}", e) + else: + id = result.inserted_id + self._id = id + return id + + def delete(self) -> bool: + """ + Deletes an object from the database + + :param collection: Collection to delete object from + :return: If delete was successful + """ + to_delete = self.to_dict() + search = {"_id": to_delete["_id"]} + try: + result = self.coll.delete_one(search) + except Exception as e: + logger.error(f"Failed to delete {type(self).__name__}", e) + else: + return result.deleted_count > 0 + + @classmethod + def delete_many(cls, **kwargs) -> bool: + try: + coll = db_instance[coll_lookup[cls.__name__]] + result = coll.delete_many(kwargs) + except Exception as e: + logger.error(f"Failed to delete {cls.__name__}s", e) + else: + return result.deleted_count > 0 + + def update(self) -> bool: + """ + Updates an object in the database + + :param collection: Collection to update + :return: If update was successful + """ + to_update = self.to_dict() + search = {"_id": to_update["_id"]} + try: + result = self.coll.update_one( + search, {"$set": to_update}, upsert=True + ) + except Exception as e: + logger.error(f"Failed to update {type(self).__name__}", e) + else: + return result.modified_count > 0 + + @classmethod + def get(cls, **kwargs) -> Optional[object]: + """ + Get an object from MongoDB + + :param collection: Collection to query + :return: Optional object + """ + try: + sort = None + if "sort" in kwargs: + sort = kwargs.pop("sort") + if type(sort) is list: + sort = [x.as_tuple() for x in sort] + else: + sort = [sort.as_tuple()] + coll = db_instance[coll_lookup[cls.__name__]] + result = coll.find_one(kwargs, sort=sort) + except Exception as e: + logger.error(f"Failed to get {cls.__name__}", e) + else: + if result: + _id = result.pop("_id") + r = cls(**result) + r._id = _id + return r + return None + + @classmethod + def get_many(cls, **kwargs) -> Optional[list]: + """ + Gets objects from MongoDB + + :param collection: Collection to query + :return: Optional object + """ + global db_instance + try: + sort = None + if "sort" in kwargs: + sort = kwargs.pop("sort") + if type(sort) is list: + sort = [x.as_tuple() for x in sort] + else: + sort = [sort.as_tuple()] + coll = db_instance[coll_lookup[cls.__name__]] + result = coll.find(kwargs, sort=sort) + except Exception as e: + logger.error(f"Failed to get {cls.__name__}", e) + else: + if result: + r_list = [] + for r_item in result: + _id = r_item.pop("_id") + r = cls(**r_item) + r._id = _id + r_list.append(r) + return r_list + return [] + + +@dataclass +class ActiveObject: + """ + A type of Mongo object that can be active + + :param active: If object is active + """ + + active: bool + + @classmethod + def get_active(cls, **kwargs) -> list: + """ + Gets a list of active objects + + :param collection: Collection to query + :return: List of active objects + """ + kwargs.update({"active": True}) + try: + sort = None + if "sort" in kwargs: + sort = kwargs.pop("sort") + if type(sort) is list: + sort = [x.as_tuple() for x in sort] + else: + sort = [sort.as_tuple()] + coll = db_instance[coll_lookup[cls.__name__]] + result = coll.find(kwargs, sort=sort) + except Exception as e: + logger.error(f"Failed to get {cls.__name__}", e) + else: + if result: + r_list = [] + for r_item in result: + _id = r_item.pop("_id") + r = cls(**r_item) + r._id = _id + r_list.append(r) + return r_list + return [] + + +#################### +# Database Objects # +#################### + + +@dataclass +class Autopurge(MongoObject): + """ + Channel Autoreact object + + :param _id: MongoDB ID + :param guild: ID of guild + :param channel: ID of channel + :param delay: Purge delay + :param admin: ID of admin who added autoreact + :param created_at: Time the autoreact was created + """ + + guild: int + channel: int + delay: int + admin: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Autoreact(MongoObject): + """ + Channel Autoreact object + + :param _id: MongoDB ID + :param guild: ID of guild + :param channel: ID of channel + :param reactions: List of reactions + :param admin: ID of admin who added autoreact + :param created_at: Time the autoreact was created + """ + + guild: int + channel: int + reactions: list + admin: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Ban(MongoObject, ActiveObject): + """ + User Ban object + + :param _id: MongoDB ID + :param active: If the ban is active + :param admin: ID of admin who banned the user + :param user: ID of banned user + :param username: Username of banned user + :param discrim: Discriminator of banned user + :param duration: Duration of ban + :param guild: ID of guild that user belonged to + :param type: Type of ban + :param reason: Reason for the ban + :param created_at: Time the ban happened + """ + + admin: int + user: int + username: str + discrim: int + duration: int + guild: int + type: str + reason: str + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Config(MongoObject): + """ + J.A.R.V.I.S. Config object + + :param _id: MongoDB ID + :param key: Config key + :param value: Config value + """ + + key: str + value: any + + +@dataclass +class Joke(MongoObject): + """ + Joke object + + :param _id: MongoDB ID + :param id: Reddit ID + :param body: Joke Body + :param title: Joke Title + :param created_utc: Created at + :param over_18: Is NSFW + :param score: Reddit Score + """ + + id: str + body: str + title: str + created_utc: datetime + over_18: bool + score: int + + +@dataclass +class Kick(MongoObject): + """ + User Kick object + + :param _id: MongoDB ID + :param user: Kicked User ID + :param reason: Kick reason + :param admin: ID of admin who kicked user + :param created_at: Time user was kicked + :param guild: Guild ID + """ + + admin: int + guild: int + reason: str + user: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Lock(MongoObject, ActiveObject): + """ + Channel Lock object + + :param _id: MongoDB ID + :param active: If the lock is active + :param admin: ID of admin who locked channel + :param channel: ID of locked channel + :param duration: Duration of lock + :param guild: ID of guild that channel belongs to + :param reason: Reason for the lock + :param created_at: Time the lock happened + """ + + admin: int + channel: int + duration: int + guild: int + reason: str + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Mute(MongoObject, ActiveObject): + """ + User Mute object + + :param _id: MongoDB ID + :param active: If the mute is active + :param admin: ID of admin who muted the user + :param user: ID of muted user + :param duration: Duration of mute + :param guild: ID of guild that user belongs to + :param reason: Reason for the mute + :param created_at: Time the mute happened + """ + + admin: int + user: int + duration: int + guild: int + reason: str + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Purge(MongoObject): + """ + Channel Purge object + + :param _id: MongoDB ID + :param admin: ID of admin who purged messages + :param channel: ID of purged channel + :param guild: ID of guild that channel belongs to + :param count: Number of purged messages + :param created_at: Time the purge happened + """ + + admin: int + channel: int + guild: int + count: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Setting(MongoObject): + """ + Guild Setting object + + :param _id: MongoDB ID + :param guild: ID of guild + :param setting: Setting key + :param value: Setting value + """ + + guild: int + setting: str + value: any + + +@dataclass +class Star(MongoObject): + """ + Starboard Star object + + :param _id: MongoDB ID + :param index: Starboard star index + :param message: Star Message ID + :param channel: Starboard Channel ID + :param guild: Starboard Guild ID + :param admin: ID of admin who created star + :param created_at: Time created + """ + + index: int + message: int + channel: int + guild: int + admin: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Starboard(MongoObject): + """ + Channel Starboard object + + :param _id: MongoDB ID + :param channel: Starboard Channel ID + :param guild: Starboard Guild ID + :param admin: ID of admin who created starboard + :param created_at: Time created + """ + + channel: int + guild: int + admin: int + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Unban(MongoObject): + """ + Guild Unban object + + :param _id: MongoDB ID + :param user: User ID + :param username: User Username + :param discrim: User Discriminator + :param guild: Guild ID + :param admin: Admin who unbanned user + :param reason: Reason for unban + :param created_at: Time created + """ + + user: int + username: str + discrim: int + guild: int + admin: int + reason: str + created_at: datetime = field(default_factory=datetime.utcnow) + + +@dataclass +class Warning(MongoObject, ActiveObject): + """ + User Warning object + + :param _id: MongoDB ID + :param active: If the warning is active + :param admin: ID of admin who warned the user + :param created_at: Time the warning happened + :param duration: Duration of warning + :param guild: ID of guild that user belongs to + :param reason: Reason for the warning + :param user: ID of warned user + """ + + admin: int + user: int + duration: int + guild: int + reason: str + created_at: datetime = field(default_factory=datetime.utcnow) diff --git a/jarvis/utils/__init__.py b/jarvis/utils/__init__.py index b34862c..a0b1b68 100644 --- a/jarvis/utils/__init__.py +++ b/jarvis/utils/__init__.py @@ -7,7 +7,7 @@ from discord.ext import commands import jarvis.cogs import jarvis.config -import jarvis.utils.db +import jarvis.db __all__ = ["field", "db"]