jarvis/jarvis-bot
jarvis/jarvis-bot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'db-rewrite' into 'main'

Browse source 
Massive Database Access Re-write

Closes #26, Closes #61

See merge request stark-industries/j.a.r.v.i.s.!16
This commit is contained in:
Zeva Rose 2021-07-22 14:15:00 +00:00
commit 14e0a865a7
18 changed files with 1033 additions and 677 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

198
jarvis/__init__.py
View file

@ -7,14 +7,15 @@ import pymongo
from discord import DMChannel, Intents, Member, Message from discord import DMChannel, Intents, Member, Message
from discord.ext import commands from discord.ext import commands
from discord.ext.tasks import loop from discord.ext.tasks import loop
from discord.utils import find, get from discord.utils import find
from discord_slash import SlashCommand from discord_slash import SlashCommand
from psutil import Process from psutil import Process
from jarvis import logo, utils from jarvis import logo, utils
from jarvis.config import get_config 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 import build_embed
from jarvis.utils.db import DBManager
from jarvis.utils.field import Field from jarvis.utils.field import Field
if asyncio.get_event_loop().is_closed(): 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) slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True)
jarvis_self = Process() jarvis_self = Process()
__version__ = "1.0.1" __version__ = "1.1.0"
db = DBManager(get_config().mongo).mongo db = DBManager(get_config().mongo).mongo
jarvis_db = db.jarvis
@jarvis.event @jarvis.event
@ -67,25 +69,16 @@ async def on_ready():
@jarvis.event @jarvis.event
async def on_member_join(user: Member): async def on_member_join(user: Member):
guild = user.guild guild = user.guild
db = DBManager(get_config().mongo).mongo mutes = Mute.get_active(guild=guild.id)
mutes = list(
db.jarvis.mutes.find(
{"active": True, "user": user.id, "guild": guild.id}
)
)
if mutes and len(mutes) >= 1: if mutes and len(mutes) >= 1:
mute_role = db.jarvis.settings.find_one( mute_role = Setting.get(guild=guild.id, setting="mute")
{"guild": guild.id, "setting": "mute"} role = guild.get_role(mute_role.value)
)
role = guild.get_role(mute_role["value"])
await user.add_roles( await user.add_roles(
role, reason="User is muted still muted from prior mute" role, reason="User is muted still muted from prior mute"
) )
unverified = db.jarvis.settings.find_one( unverified = Setting.get(guild=guild.id, setting="unverified")
{"guild": guild.id, "setting": "unverified"}
)
if 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") 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) not isinstance(message.channel, DMChannel)
and message.author.id != jarvis.user.id and message.author.id != jarvis.user.id
): ):
autoreact = db.jarvis.autoreact.find_one( autoreact = Autoreact.get(
{"guild": message.guild.id, "channel": message.channel.id} guild=message.guild.id,
channel=message.channel.id,
) )
if autoreact: if autoreact:
for reaction in autoreact["reactions"]: for reaction in autoreact.reactions:
await message.add_reaction(reaction) await message.add_reaction(reaction)
massmention = db.jarvis.settings.find_one( massmention = Setting.get(
{"guild": message.guild.id, "setting": "massmention"} guild=message.guild.id,
setting="massmention",
) )
if ( if (
massmention["value"] > 0 massmention.value > 0
and len(message.mentions) and len(message.mentions)
- (1 if message.author in message.mentions else 0) - (1 if message.author in message.mentions else 0)
> massmention["value"] > massmention.value
): ):
db.jarvis.warns.insert_one( warning = Warning(
{ active=True,
"user": message.author.id, admin=get_config().client_id,
"reason": "Mass Mention", duration=24,
"admin": get_config().client_id, guild=message.guild.id,
"time": datetime.now(), reason="Mass Mention",
"guild": message.guild.id, user=message.author.id,
"duration": 24,
"active": True,
}
) )
warning.insert()
fields = [Field("Reason", "Mass Mention", False)] fields = [Field("Reason", "Mass Mention", False)]
embed = build_embed( embed = build_embed(
title="Warning", title="Warning",
@ -138,9 +131,7 @@ async def on_message(message: Message):
+ f"| {message.author.id}" + f"| {message.author.id}"
) )
await message.channel.send(embed=embed) await message.channel.send(embed=embed)
roleping = db.jarvis.settings.find_one( roleping = Setting.get(guild=message.guild.id, setting="roleping")
{"guild": message.guild.id, "setting": "roleping"}
)
roles = [] roles = []
for mention in message.role_mentions: for mention in message.role_mentions:
roles.append(mention.id) roles.append(mention.id)
@ -149,22 +140,18 @@ async def on_message(message: Message):
roles.append(role.id) roles.append(role.id)
if ( if (
roleping roleping
and any(x in roleping["value"] for x in roles) and any(x in roleping.value for x in roles)
and not any( and not any(x.id in roleping.value for x in message.author.roles)
x.id in roleping["value"] for x in message.author.roles
)
): ):
db.jarvis.warns.insert_one( warning = Warning(
{ active=True,
"user": message.author.id, admin=get_config().client_id,
"reason": "Pinged a blocked role/user with a blocked role", duration=24,
"admin": get_config().client_id, guild=message.guild.id,
"time": datetime.now(), reason="Pinged a blocked role/user with a blocked role",
"guild": message.guild.id, user=message.author.id,
"duration": 24,
"active": True,
}
) )
warning.insert()
fields = [ fields = [
Field( Field(
"Reason", "Reason",
@ -188,11 +175,9 @@ async def on_message(message: Message):
+ f"| {message.author.id}" + f"| {message.author.id}"
) )
await message.channel.send(embed=embed) await message.channel.send(embed=embed)
autopurge = db.jarvis.autopurge.find_one( autopurge = Setting.get(guild=message.guild.id, setting="autopurge")
{"guild": message.guild.id, "channel": message.channel.id}
)
if autopurge: if autopurge:
await message.delete(delay=autopurge["delay"]) await message.delete(delay=autopurge.delay)
content = re.sub(r"\s+", "", message.content) content = re.sub(r"\s+", "", message.content)
match = invites.search(content) match = invites.search(content)
if match: if match:
@ -203,17 +188,15 @@ async def on_message(message: Message):
] ]
if match.group(1) not in allowed: if match.group(1) not in allowed:
await message.delete() await message.delete()
db.jarvis.warns.insert_one( warning = Warning(
{ active=True,
"user": message.author.id, admin=get_config().client_id,
"reason": "Sent an invite link", duration=24,
"admin": get_config().client_id, guild=message.guild.id,
"time": datetime.now(), reason="Sent an invite link",
"guild": message.guild.id, user=message.author.id,
"duration": 24,
"active": True,
}
) )
warning.insert()
fields = [ fields = [
Field( Field(
"Reason", "Reason",
@ -255,56 +238,57 @@ async def on_guild_join(guild):
await general.send("Importing all preferences from home interface...") await general.send("Importing all preferences from home interface...")
# Set some default settings # Set some default settings
db = DBManager(get_config().mongo).mongo setting = Setting(guild=guild.id, setting="massmention", value=5)
db.jarvis.settings.insert_one( setting.insert()
{"guild": guild.id, "setting": "massmention", "value": 5}
)
await general.send("Systems are now fully operational") await general.send("Systems are now fully operational")
@loop(minutes=1) @loop(minutes=1)
async def unmute(): async def unmute():
db = DBManager(get_config().mongo).mongo mutes = Mute.get_active(duration={"$gt": 0})
mutes = list(db.jarvis.mutes.find({"active": True, "length": {"$gt": 0}})) mute_roles = Setting.get(setting="mute")
mute_roles = list(
db.jarvis.settings.find({"setting": "mute"}).sort(
[("guild", pymongo.ASCENDING)]
)
)
updates = [] updates = []
for mute in mutes: for mute in mutes:
if mute["time"] + timedelta(minutes=mute["length"]) < datetime.now(): if (
mute_role = [ mute.created_at + timedelta(minutes=mute.duration)
x["value"] for x in mute_roles if x["guild"] == mute["guild"] < datetime.utcnow()
][0] ):
guild = await jarvis.fetch_guild(mute["guild"]) 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) role = guild.get_role(mute_role)
user = await guild.fetch_member(mute["user"]) user = await guild.fetch_member(mute.user)
if user: if user:
if role in user.roles: if role in user.roles:
await user.remove_roles(role, reason="Mute expired") await user.remove_roles(role, reason="Mute expired")
# Objects can't handle bulk_write, so handle it via raw methods
updates.append( updates.append(
pymongo.UpdateOne( 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}}, {"$set": {"active": False}},
) )
) )
if updates: if updates:
db.jarvis.mutes.bulk_write(updates) jarvis_db.mutes.bulk_write(updates)
@loop(minutes=10) @loop(minutes=10)
async def unban(): async def unban():
db = DBManager(get_config().mongo).mongo bans = Ban.get_active(type="temp")
bans = list(db.jarvis.bans.find({"active": True, "type": "temp"}))
updates = [] updates = []
for ban in bans: for ban in bans:
if ban["time"] + timedelta( if ban.created_at + timedelta(
hours=ban["length"] hours=ban.duration
) < datetime.now() + timedelta(minutes=10): ) < datetime.utcnow() + timedelta(minutes=10):
guild = await jarvis.fetch_guild(ban["guild"]) guild = await jarvis.fetch_guild(ban.guild)
user = await jarvis.fetch_user(ban["user"]) user = await jarvis.fetch_user(ban.user)
if user: if user:
guild.unban(user) guild.unban(user)
updates.append( updates.append(
@ -312,25 +296,27 @@ async def unban():
{ {
"user": user.id, "user": user.id,
"guild": guild.id, "guild": guild.id,
"time": ban["time"], "created_at": ban.created_at,
"type": "temp", "type": "temp",
}, },
{"$set": {"active": False}}, {"$set": {"active": False}},
) )
) )
if updates: if updates:
db.jarvis.bans.bulk_write(updates) jarvis_db.bans.bulk_write(updates)
@loop(minutes=1) @loop(minutes=1)
async def unlock(): async def unlock():
db = DBManager(get_config().mongo).mongo locks = Lock.get_active()
locks = list(db.jarvis.locks.find({"active": True}))
updates = [] updates = []
for lock in locks: for lock in locks:
if lock["time"] + timedelta(minutes=lock["duration"]) < datetime.now(): if (
guild = await jarvis.fetch_guild(lock["guild"]) lock.created_at + timedelta(minutes=lock.duration)
channel = await jarvis.fetch_channel(lock["channel"]) < datetime.utcnow()
):
guild = await jarvis.fetch_guild(lock.guild)
channel = await jarvis.fetch_channel(lock.channel)
if channel: if channel:
roles = await guild.fetch_roles() roles = await guild.fetch_roles()
for role in roles: for role in roles:
@ -344,29 +330,31 @@ async def unlock():
{ {
"channel": channel.id, "channel": channel.id,
"guild": guild.id, "guild": guild.id,
"time": lock["time"], "created_at": lock.created_at,
}, },
{"$set": {"active": False}}, {"$set": {"active": False}},
) )
) )
if updates: if updates:
db.jarvis.locks.bulk_write(updates) jarvis_db.locks.bulk_write(updates)
@loop(hours=1) @loop(hours=1)
async def unwarn(): async def unwarn():
db = DBManager(get_config().mongo).mongo warns = Warning.get_active()
warns = list(db.jarvis.warns.find({"active": True}))
updates = [] updates = []
for warn in warns: 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( updates.append(
pymongo.UpdateOne( pymongo.UpdateOne(
{"_id": warn["_id"]}, {"$set": {"active": False}} {"_id": warn._id}, {"$set": {"active": False}}
) )
) )
if updates: if updates:
db.jarvis.warns.bulk_write(updates) jarvis_db.warns.bulk_write(updates)
def run(ctx=None): def run(ctx=None):

394
jarvis/cogs/admin.py
View file

@ -1,5 +1,5 @@
import re import re
from datetime import datetime, timedelta from datetime import datetime
from typing import Union from typing import Union
import pymongo import pymongo
@ -10,8 +10,19 @@ from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_choice, create_option from discord_slash.utils.manage_commands import create_choice, create_option
import jarvis 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 import build_embed
from jarvis.utils.db import DBManager
from jarvis.utils.field import Field from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions from jarvis.utils.permissions import admin_or_permissions
@ -26,7 +37,7 @@ class AdminCog(commands.Cog):
def __init__(self, bot: commands.Bot): def __init__(self, bot: commands.Bot):
self.bot = bot self.bot = bot
config = jarvis.config.get_config() config = jarvis.config.get_config()
self.db = DBManager(config.mongo).mongo self.db = DBManager(config.mongo).mongo.jarvis
@cog_ext.cog_slash( @cog_ext.cog_slash(
name="ban", name="ban",
@ -83,6 +94,11 @@ class AdminCog(commands.Cog):
"You cannot set a temp ban to < 0 hours.", hidden=True "You cannot set a temp ban to < 0 hours.", hidden=True
) )
return 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: if len(reason) > 100:
await ctx.send("Reason must be < 100 characters", hidden=True) await ctx.send("Reason must be < 100 characters", hidden=True)
return return
@ -100,11 +116,8 @@ class AdminCog(commands.Cog):
f"You have been {mtype}banned from {guild_name}." f"You have been {mtype}banned from {guild_name}."
+ f" Reason:\n{reason}" + f" Reason:\n{reason}"
) )
time = datetime.now()
expiry = None
if mtype == "temp": if mtype == "temp":
user_message += f"\nDuration: {duration} hours" user_message += f"\nDuration: {duration} hours"
expiry = time + timedelta(hours=duration)
try: try:
await ctx.guild.ban(user, reason=reason) await ctx.guild.ban(user, reason=reason)
@ -124,41 +137,30 @@ class AdminCog(commands.Cog):
if type == "soft": if type == "soft":
active = False active = False
self.db.jarvis.bans.insert_one( _ = Ban(
{ user=user.id,
"user": user.id, username=user.name,
"username": user.name, discrim=user.discriminator,
"discrim": user.discriminator, reason=reason,
"reason": reason, admin=ctx.author.id,
"admin": ctx.author.id, guild=ctx.guild.id,
"time": datetime.now(), type=type,
"guild": ctx.guild.id, duration=duration,
"type": type, active=active,
"duration": duration, ).insert()
"expiry": expiry,
"active": active,
}
)
async def discord_apply_unban( async def discord_apply_unban(
self, ctx: SlashContext, user: User, reason: str self, ctx: SlashContext, user: User, reason: str
): ):
await ctx.guild.unban(user, reason=reason) await ctx.guild.unban(user, reason=reason)
self.db.jarvis.unbans.insert_one( _ = Unban(
{ user=user.id,
"user": user.id, username=user.name,
"username": user.name, discrim=user.discriminator,
"discrim": user.discriminator, guild=ctx.guild.id,
"guild": ctx.guild.id, admin=ctx.author.id,
"admin": ctx.author.id, reason=reason,
"reason": reason, ).insert()
"time": datetime.now(),
}
)
_ = self.db.jarvis.bans.update(
{"user": user.id, "guild": ctx.guild.id},
{"$set": {"active": False}},
)
await ctx.send("User successfully unbanned.\nReason: " + reason) await ctx.send("User successfully unbanned.\nReason: " + reason)
@cog_ext.cog_slash( @cog_ext.cog_slash(
@ -239,8 +241,8 @@ class AdminCog(commands.Cog):
# We take advantage of the previous checks to save CPU cycles # We take advantage of the previous checks to save CPU cycles
if not discord_ban_info: if not discord_ban_info:
if isinstance(user, int): if isinstance(user, int):
database_ban_info = self.db.jarvis.bans.find_one( database_ban_info = Ban.get(
{"guild": ctx.guild.id, "user": user, "active": True} guild=ctx.guild.id, user=user, active=True
) )
else: else:
search = { search = {
@ -250,7 +252,7 @@ class AdminCog(commands.Cog):
} }
if discrim: if discrim:
search["discrim"] = 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: if not discord_ban_info and not database_ban_info:
await ctx.send(f"Unable to find user {orig_user}", hidden=True) 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 ctx, discord_ban_info.user, reason
) )
else: else:
self.db.jarvis.bans.update_many( database_ban_info.active = False
{"user": database_ban_info["id"], "guild": ctx.guild.id}, database_ban_info.update()
{"$set": {"active": False}}, _ = Unban(
) user=database_ban_info.user,
self.db.jarvis.unbans.insert_one( username=database_ban_info.username,
{ discrim=database_ban_info.discrim,
"user": database_ban_info["user"], guild=ctx.guild.id,
"username": database_ban_info["username"], admin=ctx.author.id,
"discrim": database_ban_info["discrim"], reason=reason,
"guild": ctx.guild.id, ).insert()
"admin": ctx.author.id,
"reason": reason,
"time": datetime.now(),
}
)
await ctx.send( await ctx.send(
"Unable to find user in Discord, " "Unable to find user in Discord, "
+ "but removed entry from database." + "but removed entry from database."
@ -326,24 +323,23 @@ class AdminCog(commands.Cog):
search["active"] = True search["active"] = True
if type > 0: if type > 0:
search["type"] = types[type] search["type"] = types[type]
bans = self.db.jarvis.bans.find(search).sort( bans = Ban.get_many(**search)
[("time", pymongo.DESCENDING)] bans.sort(key=lambda x: x.created_at, reverse=True)
)
ban_messages = [] ban_messages = []
db_bans = [] db_bans = []
for ban in bans: for ban in bans:
if "username" not in ban: if not ban.username:
user = await self.bot.fetch_user(ban["user"]) user = await self.bot.fetch_user(ban.user)
ban["username"] = user.name if user else "[deleted user]" ban.username = user.name if user else "[deleted user]"
ban_messages.append( ban_messages.append(
"[{0}] {1} ({2}): {3}".format( "[{0}] {1} ({2}): {3}".format(
ban["time"].strftime("%d-%m-%Y"), ban.created_at.strftime("%d-%m-%Y"),
ban["username"], ban.username,
ban["user"], ban.user,
ban["reason"], ban.reason,
) )
) )
db_bans.append(ban["user"]) db_bans.append(ban.user)
bans = await ctx.guild.bans() bans = await ctx.guild.bans()
for ban in bans: for ban in bans:
if ban.user.id not in db_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"{user.name} has been kicked from {guild_name}."
+ f"Reason:\n{reason}" + f"Reason:\n{reason}"
) )
self.db.jarvis.kicks.insert_one( _ = Kick(
{ user=user.id,
"user": user.id, reason=reason,
"reason": reason, admin=ctx.author.id,
"admin": ctx.authod.id, guild=ctx.guild.id,
"time": datetime.now(), ).insert()
"guild": ctx.guild.id,
}
)
@cog_ext.cog_slash( @cog_ext.cog_slash(
name="purge", name="purge",
@ -438,15 +431,12 @@ class AdminCog(commands.Cog):
async for message in channel.history(limit=amount + 1): async for message in channel.history(limit=amount + 1):
messages.append(message) messages.append(message)
await channel.delete_messages(messages) await channel.delete_messages(messages)
self.db.jarvis.purges.insert_one( _ = Purge(
{ channel=ctx.channel.id,
"channel": ctx.channel.id, guild=ctx.guild.id,
"guild": ctx.guild.id, admin=ctx.author.id,
"admin": ctx.author.id, count=amount,
"count": amount, ).insert()
"time": datetime.now(),
}
)
@cog_ext.cog_slash( @cog_ext.cog_slash(
name="mute", name="mute",
@ -485,45 +475,29 @@ class AdminCog(commands.Cog):
if len(reason) > 100: if len(reason) > 100:
await ctx.send("Reason must be < 100 characters", hidden=True) await ctx.send("Reason must be < 100 characters", hidden=True)
return return
mute_setting = self.db.jarvis.settings.find_one( mute_setting = Setting.get(guild=ctx.guild.id, setting="mute")
{"guild": ctx.guild.id, "setting": "mute"}
)
if not mute_setting: if not mute_setting:
await ctx.send( await ctx.send(
"Please configure a mute role with /settings mute <role> first", "Please configure a mute role "
+ "with /settings mute <role> first",
hidden=True, hidden=True,
) )
return return
role = get(ctx.guild.roles, id=mute_setting["value"]) role = get(ctx.guild.roles, id=mute_setting.value)
if role in user.roles: if role in user.roles:
await ctx.send("User already muted", hidden=True) await ctx.send("User already muted", hidden=True)
return
await user.add_roles(role, reason=reason) await user.add_roles(role, reason=reason)
time = datetime.now() if duration < 0 or duration > 300:
expiry = None
if duration < 0:
duration = -1 duration = -1
if duration >= 0: _ = Mute(
expiry = time + timedelta(minutes=duration) user=user.id,
self.db.jarvis.mutes.insert_one( reason=reason,
{ admin=ctx.author.id,
"user": user.id, guild=ctx.guild.id,
"reason": reason, duration=duration,
"admin": ctx.author.id, active=True if duration >= 0 else False,
"time": time, ).insert()
"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}},
)
await ctx.send(f"{user.mention} has been muted.\nReason: {reason}") await ctx.send(f"{user.mention} has been muted.\nReason: {reason}")
@cog_ext.cog_slash( @cog_ext.cog_slash(
@ -540,9 +514,7 @@ class AdminCog(commands.Cog):
) )
@admin_or_permissions(mute_members=True) @admin_or_permissions(mute_members=True)
async def _unmute(self, ctx: SlashContext, user: Member): async def _unmute(self, ctx: SlashContext, user: Member):
mute_setting = self.db.jarvis.settings.find_one( mute_setting = Setting.get(guild=ctx.guild.id, setting="mute")
{"guild": ctx.guild.id, "setting": "mute"}
)
if not mute_setting: if not mute_setting:
await ctx.send( await ctx.send(
"Please configure a mute role with " "Please configure a mute role with "
@ -551,20 +523,17 @@ class AdminCog(commands.Cog):
) )
return return
role = get(ctx.guild.roles, id=mute_setting["value"]) role = get(ctx.guild.roles, id=mute_setting.value)
if role in user.roles: if role in user.roles:
await user.remove_roles(role, reason="Unmute") await user.remove_roles(role, reason="Unmute")
else: else:
await ctx.send("User is not muted.", hidden=True) await ctx.send("User is not muted.", hidden=True)
return return
self.db.jarvis.mutes.update_many( mutes = Mute.get_many(guild=ctx.guild.id, user=user.id)
{ for mute in mutes:
"guild": ctx.guild.id, mute.active = False
"user": user.id, mute.update()
},
{"$set": {"active": False}},
)
await ctx.send(f"{user.mention} has been unmuted.") await ctx.send(f"{user.mention} has been unmuted.")
async def _lock_channel( async def _lock_channel(
@ -628,6 +597,12 @@ class AdminCog(commands.Cog):
channel: Union[TextChannel, VoiceChannel] = None, channel: Union[TextChannel, VoiceChannel] = None,
): ):
await ctx.defer(hidden=True) 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: if len(reason) > 100:
await ctx.send("Reason must be < 100 characters", hidden=True) await ctx.send("Reason must be < 100 characters", hidden=True)
return return
@ -638,17 +613,13 @@ class AdminCog(commands.Cog):
await self._lock_channel(channel, role, ctx.author, reason) await self._lock_channel(channel, role, ctx.author, reason)
except Exception: except Exception:
continue # Just continue on error continue # Just continue on error
self.db.jarvis.locks.insert_one( _ = Lock(
{ channel=channel.id,
"channel": channel.id, guild=ctx.guild.id,
"guild": ctx.guild.id, admin=ctx.author.id,
"admin": ctx.author.id, reason=reason,
"reason": reason, duration=duration,
"duration": duration, ).insert()
"active": True,
"time": datetime.now(),
}
)
await ctx.send(f"{channel.mention} locked for {duration} minute(s)") await ctx.send(f"{channel.mention} locked for {duration} minute(s)")
@cog_ext.cog_slash( @cog_ext.cog_slash(
@ -671,9 +642,7 @@ class AdminCog(commands.Cog):
): ):
if not channel: if not channel:
channel = ctx.channel channel = ctx.channel
lock = self.db.jarvis.locks.find_one( lock = Lock.get(guild=ctx.guild.id, channel=channel.id, active=True)
{"guild": ctx.guild.id, "channel": channel.id, "active": True}
)
if not lock: if not lock:
await ctx.send(f"{channel.mention} not locked.", hidden=True) await ctx.send(f"{channel.mention} not locked.", hidden=True)
return return
@ -682,13 +651,8 @@ class AdminCog(commands.Cog):
await self._unlock_channel(channel, role, ctx.author) await self._unlock_channel(channel, role, ctx.author)
except Exception: except Exception:
continue # Just continue on error continue # Just continue on error
self.db.jarvis.locks.update_one( lock.active = False
{ lock.update()
"channel": channel.id,
"guild": ctx.guild.id,
},
{"$set": {"active": False}},
)
await ctx.send(f"{channel.mention} unlocked") await ctx.send(f"{channel.mention} unlocked")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -717,7 +681,13 @@ class AdminCog(commands.Cog):
reason: str, reason: str,
duration: int = 10, 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 channels = ctx.guild.channels
roles = ctx.guild.roles roles = ctx.guild.roles
updates = [] updates = []
@ -736,12 +706,12 @@ class AdminCog(commands.Cog):
"reason": reason, "reason": reason,
"duration": duration, "duration": duration,
"active": True, "active": True,
"time": datetime.now(), "created_at": datetime.utcnow(),
} }
) )
) )
if updates: 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)") await ctx.send(f"Server locked for {duration} minute(s)")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -757,9 +727,7 @@ class AdminCog(commands.Cog):
channels = ctx.guild.channels channels = ctx.guild.channels
roles = ctx.guild.roles roles = ctx.guild.roles
updates = [] updates = []
locks = list( locks = Lock.get_many(guild=ctx.guild.id, active=True)
self.db.jarvis.locks.find({"guild": ctx.guild.id, "active": True})
)
if not locks: if not locks:
await ctx.send("No lockdown detected.", hidden=True) await ctx.send("No lockdown detected.", hidden=True)
return return
@ -781,7 +749,7 @@ class AdminCog(commands.Cog):
) )
) )
if updates: if updates:
self.db.jarvis.locks.bulk_write(updates) self.db.locks.bulk_write(updates)
await ctx.send("Server unlocked") await ctx.send("Server unlocked")
@cog_ext.cog_slash( @cog_ext.cog_slash(
@ -815,25 +783,20 @@ class AdminCog(commands.Cog):
if len(reason) > 100: if len(reason) > 100:
await ctx.send("Reason must be < 100 characters", hidden=True) await ctx.send("Reason must be < 100 characters", hidden=True)
return 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() await ctx.defer()
self.db.jarvis.warns.insert_one( _ = Warning(
{ user=user.id,
"user": user.id, reason=reason,
"reason": reason, admin=ctx.author.id,
"admin": ctx.author.id, guild=ctx.guild.id,
"time": datetime.now(), duration=duration,
"guild": ctx.guild.id, ).insert()
"duration": duration,
"active": True,
}
)
count = len(
list(
self.db.jarvis.warns.find(
{"user": user.id, "guild": ctx.guild.id, "active": True}
)
)
)
fields = [Field("Reason", reason, False)] fields = [Field("Reason", reason, False)]
embed = build_embed( embed = build_embed(
title="Warning", title="Warning",
@ -863,18 +826,11 @@ class AdminCog(commands.Cog):
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _warnings(self, ctx: SlashContext, user: User): async def _warnings(self, ctx: SlashContext, user: User):
await ctx.defer() await ctx.defer()
warnings = list( warnings = Warning.get_many(user=user.id, guild=ctx.guild.id)
self.db.jarvis.warns.find(
{
"user": user.id,
"guild": ctx.guild.id,
}
)
)
active = ( active = (
[ [
f'`{y["time"].strftime("%Y-%m-%d %H:%M:%S")}` - {y["reason"]}' f'`{y.created_at.strftime("%Y-%m-%d %H:%M:%S")}` - {y.reason}'
for y in list(filter(lambda x: x["active"], warnings)) for y in list(filter(lambda x: x.active, warnings))
] ]
if warnings if warnings
else ["None"] else ["None"]
@ -918,23 +874,17 @@ class AdminCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _roleping_block(self, ctx: SlashContext, role: Role): async def _roleping_block(self, ctx: SlashContext, role: Role):
roles = self.db.jarvis.settings.find_one( roles = Setting.get(guild=ctx.guild.id, setting="roleping")
{"guild": ctx.guild.id, "setting": "roleping"}
)
if not roles: 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( await ctx.send(
f"Role `{role.name}` already in blocklist.", hidden=True f"Role `{role.name}` already in blocklist.", hidden=True
) )
return return
roles["value"].append(role.id) roles.value.append(role.id)
self.db.jarvis.settings.update_one( roles.update()
{"guild": ctx.guild.id, "setting": "roleping"},
{"$set": roles},
upsert=True,
)
await ctx.send(f"Role `{role.name}` added to blocklist.") await ctx.send(f"Role `{role.name}` added to blocklist.")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -952,24 +902,18 @@ class AdminCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _roleping_allow(self, ctx: SlashContext, role: Role): async def _roleping_allow(self, ctx: SlashContext, role: Role):
roles = self.db.jarvis.settings.find_one( roles = Setting.get(guild=ctx.guild.id, setting="roleping")
{"guild": ctx.guild.id, "setting": "roleping"}
)
if not roles: if not roles:
await ctx.send("No blocklist configured.", hidden=True) await ctx.send("No blocklist configured.", hidden=True)
return return
if role.id not in roles["value"]: if role.id not in roles.value:
await ctx.send( await ctx.send(
f"Role `{role.name}` not in blocklist.", hidden=True f"Role `{role.name}` not in blocklist.", hidden=True
) )
return return
roles["value"].delete(role.id) roles.value.delete(role.id)
self.db.jarvis.settings.update_one( roles.update()
{"guild": ctx.guild.id, "setting": "roleping"},
{"$set": roles},
upsert=True,
)
await ctx.send(f"Role `{role.name}` removed blocklist.") await ctx.send(f"Role `{role.name}` removed blocklist.")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -978,18 +922,16 @@ class AdminCog(commands.Cog):
description="List all blocklisted roles", description="List all blocklisted roles",
) )
async def _roleping_list(self, ctx: SlashContext): async def _roleping_list(self, ctx: SlashContext):
roles = self.db.jarvis.settings.find_one( roles = Setting.get(guild=ctx.guild.id, setting="roleping")
{"guild": ctx.guild.id, "setting": "roleping"}
)
if not roles: if not roles:
await ctx.send("No blocklist configured.", hidden=True) await ctx.send("No blocklist configured.", hidden=True)
return return
message = "Blocklisted Roles:\n```\n" message = "Blocklisted Roles:\n```\n"
if not roles["value"]: if not roles.value:
await ctx.send("No roles blocklisted.", hidden=True) await ctx.send("No roles blocklisted.", hidden=True)
return return
for role in roles["value"]: for role in roles.value:
role = ctx.guild.get_role(role) role = ctx.guild.get_role(role)
if not role: if not role:
continue continue
@ -1023,20 +965,23 @@ class AdminCog(commands.Cog):
if not isinstance(channel, TextChannel): if not isinstance(channel, TextChannel):
await ctx.send("Channel must be a TextChannel", hidden=True) await ctx.send("Channel must be a TextChannel", hidden=True)
return return
autopurge = self.db.jarvis.autopurge.find( if delay <= 0:
{"guild": ctx.guild.id, "channel": channel.id} 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: if autopurge:
await ctx.send("Autopurge already exists.", hidden=True) await ctx.send("Autopurge already exists.", hidden=True)
return return
autopurge = { autopurge = Autopurge(
"guild": ctx.guild.id, guild=ctx.guild.id,
"channel": channel.id, channel=channel.id,
"admin": ctx.author.id, admin=ctx.author.id,
"delay": delay, delay=delay,
"time": datetime.utcnow(), )
} autopurge.insert()
self.db.jarvis.autopurge.insert_one(autopurge)
await ctx.send( await ctx.send(
f"Autopurge set up on {channel.mention}, " f"Autopurge set up on {channel.mention}, "
+ f"delay is {delay} seconds" + f"delay is {delay} seconds"
@ -1057,13 +1002,11 @@ class AdminCog(commands.Cog):
) )
@admin_or_permissions(manage_messages=True) @admin_or_permissions(manage_messages=True)
async def _autopurge_remove(self, ctx: SlashContext, channel: TextChannel): async def _autopurge_remove(self, ctx: SlashContext, channel: TextChannel):
autopurge = self.db.jarvis.autopurge.find( autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if not autopurge: if not autopurge:
await ctx.send("Autopurge does not exist.", hidden=True) await ctx.send("Autopurge does not exist.", hidden=True)
return return
self.db.jarvis.autopurge.delete_one({"_id": autopurge["_id"]}) autopurge.delete()
await ctx.send(f"Autopurge removed from {channel.mention}.") await ctx.send(f"Autopurge removed from {channel.mention}.")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -1089,15 +1032,12 @@ class AdminCog(commands.Cog):
async def _autopurge_update( async def _autopurge_update(
self, ctx: SlashContext, channel: TextChannel, delay: int self, ctx: SlashContext, channel: TextChannel, delay: int
): ):
autopurge = self.db.jarvis.autopurge.find_one( autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if not autopurge: if not autopurge:
await ctx.send("Autopurge does not exist.", hidden=True) await ctx.send("Autopurge does not exist.", hidden=True)
return return
self.db.jarvis.autopurge.update_one( autopurge.delay = delay
{"_id": autopurge["_id"]}, {"$set": {"delay": delay}} autopurge.update()
)
await ctx.send( await ctx.send(
f"Autopurge delay updated to {delay} seconds on {channel.mention}." f"Autopurge delay updated to {delay} seconds on {channel.mention}."
) )

69
jarvis/cogs/autoreact.py
View file

@ -1,5 +1,4 @@
import re import re
from datetime import datetime
from discord import TextChannel from discord import TextChannel
from discord.ext import commands from discord.ext import commands
@ -7,17 +6,13 @@ from discord.utils import find
from discord_slash import SlashContext, cog_ext from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option 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.data.unicode import emoji_list
from jarvis.utils.db import DBManager from jarvis.db.types import Autoreact
class AutoReactCog(commands.Cog): class AutoReactCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
config = get_config()
self.db = DBManager(config.mongo).mongo
self.custom_emote = re.compile(r"^<:\w+:(\d+)>$") self.custom_emote = re.compile(r"^<:\w+:(\d+)>$")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -38,23 +33,20 @@ class AutoReactCog(commands.Cog):
if not isinstance(channel, TextChannel): if not isinstance(channel, TextChannel):
await ctx.send("Channel must be a text channel", hidden=True) await ctx.send("Channel must be a text channel", hidden=True)
return return
exists = self.db.jarvis.autoreact.find_one( exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if exists: if exists:
await ctx.send( await ctx.send(
f"Autoreact already exists for {channel.mention}.", hidden=True f"Autoreact already exists for {channel.mention}.", hidden=True
) )
return return
autoreact = { autoreact = Autoreact(
"guild": ctx.guild.id, guild=ctx.guild.id,
"channel": channel.id, channel=channel.id,
"reactions": [], reactions=[],
"admin": ctx.author.id, admin=ctx.author.id,
"time": datetime.now(), )
} autoreact.insert()
self.db.jarvis.autoreact.insert_one(autoreact)
await ctx.send(f"Autoreact created for {channel.mention}!") await ctx.send(f"Autoreact created for {channel.mention}!")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -71,9 +63,7 @@ class AutoReactCog(commands.Cog):
], ],
) )
async def _autoreact_delete(self, ctx, channel: TextChannel): async def _autoreact_delete(self, ctx, channel: TextChannel):
exists = self.db.jarvis.autoreact.delete_one( exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id).delete()
{"guild": channel.guild.id, "channel": channel.id}
)
if exists: if exists:
await ctx.send(f"Autoreact removed from {channel.mention}") await ctx.send(f"Autoreact removed from {channel.mention}")
else: else:
@ -119,32 +109,28 @@ class AutoReactCog(commands.Cog):
"Please use a custom emote from this server.", hidden=True "Please use a custom emote from this server.", hidden=True
) )
return return
exists = self.db.jarvis.autoreact.find_one( exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if not exists: if not exists:
await ctx.send( await ctx.send(
"Please create autoreact first with " "Please create autoreact first with "
+ f"/autoreact create {channel.mention}" + f"/autoreact create {channel.mention}"
) )
return return
if emote in exists["reactions"]: if emote in exists.reactions:
await ctx.send( await ctx.send(
f"Emote already added to {channel.mention} autoreactions.", f"Emote already added to {channel.mention} autoreactions.",
hidden=True, hidden=True,
) )
return return
if len(exists["reactions"]) >= 5: if len(exists.reactions) >= 5:
await ctx.send( 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, hidden=True,
) )
return return
exists["reactions"].append(emote) exists.reactions.append(emote)
self.db.jarvis.autoreact.update_one( exists.update()
{"_id": exists["_id"]},
{"$set": {"reactions": exists["reactions"]}},
)
await ctx.send(f"Added {emote} to {channel.mention} autoreact.") await ctx.send(f"Added {emote} to {channel.mention} autoreact.")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -168,9 +154,7 @@ class AutoReactCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str): async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str):
exists = self.db.jarvis.autoreact.find_one( exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if not exists: if not exists:
await ctx.send( await ctx.send(
"Please create autoreact first with " "Please create autoreact first with "
@ -178,17 +162,14 @@ class AutoReactCog(commands.Cog):
hidden=True, hidden=True,
) )
return return
if emote not in exists["reactions"]: if emote not in exists.reactions:
await ctx.send( await ctx.send(
f"{emote} not used in {channel.mention} autoreactions.", f"{emote} not used in {channel.mention} autoreactions.",
hidden=True, hidden=True,
) )
return return
exists["reactions"].remove(emote) exists.reactions.remove(emote)
self.db.jarvis.autoreact.update_one( exists.update()
{"_id": exists["_id"]},
{"$set": {"reactions": exists["reactions"]}},
)
await ctx.send(f"Removed {emote} from {channel.mention} autoreact.") await ctx.send(f"Removed {emote} from {channel.mention} autoreact.")
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
@ -206,9 +187,7 @@ class AutoReactCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _autoreact_list(self, ctx, channel: TextChannel): async def _autoreact_list(self, ctx, channel: TextChannel):
exists = self.db.jarvis.autoreact.find_one( exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
{"guild": ctx.guild.id, "channel": channel.id}
)
if not exists: if not exists:
await ctx.send( await ctx.send(
"Please create autoreact first with " "Please create autoreact first with "
@ -217,10 +196,10 @@ class AutoReactCog(commands.Cog):
) )
return return
message = "" message = ""
if len(exists["reactions"]) > 0: if len(exists.reactions) > 0:
message = ( message = (
f"Current active autoreacts on {channel.mention}:\n" f"Current active autoreacts on {channel.mention}:\n"
+ "\n".join(exists["reactions"]) + "\n".join(exists.reactions)
) )
else: else:
message = f"No reactions set on {channel.mention}" message = f"No reactions set on {channel.mention}"

3
jarvis/cogs/ctc2.py
View file

@ -2,9 +2,8 @@ import aiohttp
from discord.ext import commands from discord.ext import commands
from discord_slash import cog_ext from discord_slash import cog_ext
import jarvis
from jarvis.config import get_config from jarvis.config import get_config
from jarvis.utils.db import DBManager from jarvis.db import DBManager
guild_ids = [578757004059738142, 520021794380447745, 862402786116763668] guild_ids = [578757004059738142, 520021794380447745, 862402786116763668]

19
jarvis/cogs/dbrand.py
View file

@ -5,7 +5,6 @@ from discord.ext import commands
from discord_slash import cog_ext from discord_slash import cog_ext
from discord_slash.utils.manage_commands import create_option from discord_slash.utils.manage_commands import create_option
import jarvis
from jarvis.config import get_config from jarvis.config import get_config
from jarvis.data.dbrand import shipping_lookup from jarvis.data.dbrand import shipping_lookup
from jarvis.utils import build_embed from jarvis.utils import build_embed
@ -134,7 +133,8 @@ class DbrandCog(commands.Cog):
( (
create_option( create_option(
name="search", 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, option_type=3,
required=True, required=True,
) )
@ -203,10 +203,19 @@ class DbrandCog(commands.Cog):
x for x in data["country"].split(" ") if x != "the" x for x in data["country"].split(" ") if x != "the"
) )
country_urlsafe = country.replace("-", "%20") country_urlsafe = country.replace("-", "%20")
description = f"Click the link above to see shipping time to {data['country']}." description = (
description += "\n[View all shipping destinations](https://dbrand.com/shipping)" "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 += " | [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( embed = build_embed(
title="Shipping to {}".format(data["country"]), title="Shipping to {}".format(data["country"]),
description=description, description=description,

2
jarvis/cogs/dev.py
View file

@ -4,13 +4,11 @@ import re
import subprocess import subprocess
import uuid import uuid
import discord
import ulid import ulid
from bson import ObjectId from bson import ObjectId
from discord.ext import commands from discord.ext import commands
from discord_slash import cog_ext from discord_slash import cog_ext
import jarvis
from jarvis.utils import build_embed, convert_bytesize from jarvis.utils import build_embed, convert_bytesize
from jarvis.utils.field import Field from jarvis.utils.field import Field

2
jarvis/cogs/error.py
View file

@ -1,7 +1,5 @@
from discord.ext import commands from discord.ext import commands
import jarvis
class ErrorHandlerCog(commands.Cog): class ErrorHandlerCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):

3
jarvis/cogs/image.py
View file

@ -7,7 +7,6 @@ import numpy as np
from discord import File from discord import File
from discord.ext import commands from discord.ext import commands
import jarvis
from jarvis.utils import build_embed, convert_bytesize, unconvert_bytesize from jarvis.utils import build_embed, convert_bytesize, unconvert_bytesize
from jarvis.utils.field import Field from jarvis.utils.field import Field
@ -98,7 +97,7 @@ class ImageCog(commands.Cog):
Field("Accuracy", f"{accuracy:.02f}%", False), Field("Accuracy", f"{accuracy:.02f}%", False),
] ]
embed = build_embed(title=filename, description="", fields=fields) 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( await ctx.send(
embed=embed, embed=embed,
file=File(bufio, filename="resized.png"), file=File(bufio, filename="resized.png"),

11
jarvis/cogs/jokes.py
View file

@ -4,13 +4,12 @@ import traceback
from datetime import datetime from datetime import datetime
from random import randint from random import randint
import discord
from discord.ext import commands from discord.ext import commands
from discord_slash import cog_ext from discord_slash import cog_ext
import jarvis import jarvis
from jarvis.db import DBManager
from jarvis.utils import build_embed from jarvis.utils import build_embed
from jarvis.utils.db import DBManager
from jarvis.utils.field import Field from jarvis.utils.field import Field
@ -24,7 +23,7 @@ class JokeCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
config = jarvis.config.get_config() 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 # TODO: Make this a command group with subcommands
async def _joke(self, ctx, id: str = None): async def _joke(self, ctx, id: str = None):
@ -34,7 +33,7 @@ class JokeCog(commands.Cog):
return return
# TODO: Add this as a parameter that can be passed in # TODO: Add this as a parameter that can be passed in
threshold = 500 # Minimum score threshold = 500 # Minimum score
coll = self.db.mongo.jarvis.jokes coll = self.db.jokes
result = None result = None
if id: if id:
result = coll.find_one({"id": 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']}") # 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( @cog_ext.cog_slash(
name="joke", name="joke",
description="Hear a joke", description="Hear a joke",

131
jarvis/cogs/modlog.py
View file

@ -2,16 +2,14 @@ import asyncio
from datetime import datetime, timedelta from datetime import datetime, timedelta
import discord import discord
import pymongo
from discord import DMChannel from discord import DMChannel
from discord.ext import commands from discord.ext import commands
from discord.utils import find from discord.utils import find
from discord_slash import SlashContext from discord_slash import SlashContext
import jarvis
from jarvis.config import get_config 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 import build_embed
from jarvis.utils.db import DBManager
from jarvis.utils.field import Field from jarvis.utils.field import Field
@ -22,7 +20,6 @@ class ModlogCog(commands.Cog):
def __init__(self, bot: discord.ext.commands.Bot): def __init__(self, bot: discord.ext.commands.Bot):
self.bot = bot self.bot = bot
self.db = DBManager(get_config().mongo).mongo
def get_latest_log(self, auditlog, target): def get_latest_log(self, auditlog, target):
before = datetime.utcnow() - timedelta(seconds=10) before = datetime.utcnow() - timedelta(seconds=10)
@ -66,11 +63,9 @@ class ModlogCog(commands.Cog):
@commands.Cog.listener() @commands.Cog.listener()
async def on_member_ban(self, guild: discord.Guild, user: discord.User): async def on_member_ban(self, guild: discord.Guild, user: discord.User):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=guild.id, setting="modlog")
{"guild": guild.id, "setting": "modlog"}
)
if 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 await asyncio.sleep(0.5) # Need to wait for audit log
auditlog = await guild.audit_logs( auditlog = await guild.audit_logs(
limit=50, limit=50,
@ -81,15 +76,13 @@ class ModlogCog(commands.Cog):
log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) log: discord.AuditLogEntry = self.get_latest_log(auditlog, user)
admin: discord.User = log.user admin: discord.User = log.user
if admin.id == get_config().client_id: if admin.id == get_config().client_id:
mute = self.db.jarvis.bans.find_one( ban = Ban.get(
{ guild=guild.id,
"guild": guild.id, user=user.id,
"user": user.id, active=True,
"active": True, sort=MongoSort(key="created_at", type="desc"),
},
sort=[("time", pymongo.DESCENDING)],
) )
admin = guild.get_member(mute["admin"]) admin = guild.get_member(ban.admin)
embed = await self.modlog_embed( embed = await self.modlog_embed(
user, user,
admin, admin,
@ -102,11 +95,9 @@ class ModlogCog(commands.Cog):
@commands.Cog.listener() @commands.Cog.listener()
async def on_member_unban(self, guild: discord.Guild, user: discord.User): async def on_member_unban(self, guild: discord.Guild, user: discord.User):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=guild.id, setting="modlog")
{"guild": guild.id, "setting": "modlog"}
)
if 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 await asyncio.sleep(0.5) # Need to wait for audit log
auditlog = await guild.audit_logs( auditlog = await guild.audit_logs(
limit=50, limit=50,
@ -117,15 +108,13 @@ class ModlogCog(commands.Cog):
log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) log: discord.AuditLogEntry = self.get_latest_log(auditlog, user)
admin: discord.User = log.user admin: discord.User = log.user
if admin.id == get_config().client_id: if admin.id == get_config().client_id:
mute = self.db.jarvis.bans.find_one( ban = Ban.get(
{ guild=guild.id,
"guild": guild.id, user=user.id,
"user": user.id, active=True,
"active": True, sort=MongoSort(key="created_at", type="desc"),
},
sort=[("time", pymongo.DESCENDING)],
) )
admin = guild.get_member(mute["admin"]) admin = guild.get_member(ban.admin)
embed = await self.modlog_embed( embed = await self.modlog_embed(
user, user,
admin, admin,
@ -138,11 +127,9 @@ class ModlogCog(commands.Cog):
@commands.Cog.listener() @commands.Cog.listener()
async def on_member_remove(self, user: discord.User): async def on_member_remove(self, user: discord.User):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=user.guild.id, setting="modlog")
{"guild": user.guild.id, "setting": "modlog"}
)
if 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 await asyncio.sleep(0.5) # Need to wait for audit log
auditlog = await user.guild.audit_logs( auditlog = await user.guild.audit_logs(
limit=50, limit=50,
@ -153,14 +140,12 @@ class ModlogCog(commands.Cog):
log: discord.AuditLogEntry = self.get_latest_log(auditlog, user) log: discord.AuditLogEntry = self.get_latest_log(auditlog, user)
admin: discord.User = log.user admin: discord.User = log.user
if admin.id == get_config().client_id: if admin.id == get_config().client_id:
mute = self.db.jarvis.kicks.find_one( kick = Kick.get(
{ guild=user.guild.id,
"guild": user.guild.id, user=user.id,
"user": user.id, sort=MongoSort(key="created_at", type="desc"),
},
sort=[("time", pymongo.DESCENDING)],
) )
admin = user.guild.get_member(mute["admin"]) admin = user.guild.get_member(kick.admin)
embed = await self.modlog_embed( embed = await self.modlog_embed(
user, user,
admin, admin,
@ -181,15 +166,13 @@ class ModlogCog(commands.Cog):
log: discord.AuditLogEntry = self.get_latest_log(auditlog, before) log: discord.AuditLogEntry = self.get_latest_log(auditlog, before)
admin: discord.User = log.user admin: discord.User = log.user
if admin.id == get_config().client_id: if admin.id == get_config().client_id:
mute = self.db.jarvis.mutes.find_one( mute = Mute.get(
{ guild=before.guild.id,
"guild": before.guild.id, user=before.id,
"user": before.id, active=True,
"active": True, sort=MongoSort(key="created_at", type="desc"),
},
sort=[("time", pymongo.DESCENDING)],
) )
admin = before.guild.get_member(mute["admin"]) admin = before.guild.get_member(mute.admin)
return await self.modlog_embed( return await self.modlog_embed(
member=before, member=before,
admin=admin, admin=admin,
@ -208,15 +191,14 @@ class ModlogCog(commands.Cog):
log: discord.AuditLogEntry = self.get_latest_log(auditlog, before) log: discord.AuditLogEntry = self.get_latest_log(auditlog, before)
admin: discord.User = log.user admin: discord.User = log.user
if admin.id == get_config().client_id: if admin.id == get_config().client_id:
mute = self.db.jarvis.mutes.find_one( mute = Mute.get(
{ guild=before.guild.id,
"guild": before.guild.id, user=before.id,
"user": before.id, active=True,
"active": True, sort=MongoSort(key="created_at", type="desc"),
},
sort=[("time", pymongo.DESCENDING)],
) )
admin = before.guild.get_member(mute["admin"]) mute = Mute(**mute)
admin = before.guild.get_member(mute.admin)
return await self.modlog_embed( return await self.modlog_embed(
member=before, member=before,
admin=admin, admin=admin,
@ -273,28 +255,21 @@ class ModlogCog(commands.Cog):
async def on_member_update( async def on_member_update(
self, before: discord.User, after: discord.User self, before: discord.User, after: discord.User
): ):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=before.guild.id, setting="modlog")
{"guild": after.guild.id, "setting": "modlog"}
)
if modlog: if modlog:
channel = after.guild.get_channel(modlog["value"]) channel = after.guild.get_channel(modlog["value"])
await asyncio.sleep(0.5) # Need to wait for audit log await asyncio.sleep(0.5) # Need to wait for audit log
embed = None embed = None
mute = self.db.jarvis.settings.find_one( mute = Setting(guild=before.guild.id, setting="mute")
{"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:
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:
embed = await self.process_mute(before, after) 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) embed = await self.process_unmute(before, after)
elif ( elif (
verified verified
and before.guild.get_role(verified["value"]) and before.guild.get_role(verified.value) not in before.roles
not in before.roles and after.guild.get_role(verified.value) in after.roles
and after.guild.get_role(verified["value"]) in after.roles
): ):
embed = await self.process_verify(before, after) embed = await self.process_verify(before, after)
elif before.nick != after.nick: elif before.nick != after.nick:
@ -358,13 +333,11 @@ class ModlogCog(commands.Cog):
self, before: discord.Message, after: discord.Message self, before: discord.Message, after: discord.Message
): ):
if before.author != get_config().client_id: if before.author != get_config().client_id:
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=after.guild.id, setting="modlog")
{"guild": after.guild.id, "setting": "modlog"}
)
if modlog: if modlog:
if before.content == after.content or before.content is None: if before.content == after.content or before.content is None:
return return
channel = before.guild.get_channel(modlog["value"]) channel = before.guild.get_channel(modlog.value)
fields = [ fields = [
Field( Field(
"Original Message", "Original Message",
@ -398,12 +371,10 @@ class ModlogCog(commands.Cog):
@commands.Cog.listener() @commands.Cog.listener()
async def on_message_delete(self, message: discord.Message): async def on_message_delete(self, message: discord.Message):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=message.guild.id, setting="modlog")
{"guild": message.guild.id, "setting": "modlog"}
)
if modlog: if modlog:
fields = [Field("Original Message", message.content, False)] fields = [Field("Original Message", message.content, False)]
channel = message.guild.get_channel(modlog["value"]) channel = message.guild.get_channel(modlog.value)
embed = build_embed( embed = build_embed(
title="Message Deleted", title="Message Deleted",
description=f"{message.author.mention}'s message was deleted", description=f"{message.author.mention}'s message was deleted",
@ -424,11 +395,9 @@ class ModlogCog(commands.Cog):
@commands.Cog.listener() @commands.Cog.listener()
async def on_slash_command(self, ctx: SlashContext): async def on_slash_command(self, ctx: SlashContext):
if not isinstance(ctx.channel, DMChannel): if not isinstance(ctx.channel, DMChannel):
modlog = self.db.jarvis.settings.find_one( modlog = Setting.get(guild=ctx.guild.id, setting="modlog")
{"guild": ctx.guild.id, "setting": "modlog"}
)
if modlog: if modlog:
channel = ctx.guild.get_channel(modlog["value"]) channel = ctx.guild.get_channel(modlog.value)
fields = [ fields = [
Field("Command", ctx.name), Field("Command", ctx.name),
] ]

157
jarvis/cogs/owner.py
View file

@ -9,8 +9,9 @@ from discord import DMChannel, User
from discord.ext import commands from discord.ext import commands
import jarvis import jarvis
from jarvis.config import get_config, reload_config from jarvis.config import reload_config
from jarvis.utils import db, update from jarvis.db.types import Config
from jarvis.utils import update
from jarvis.utils.permissions import user_is_bot_admin from jarvis.utils.permissions import user_is_bot_admin
@ -23,9 +24,7 @@ class OwnerCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
self.config = get_config() self.admins = Config.get(key="admins")
self.admins = self.config.admins
self.db = db.DBManager(self.config.mongo)
@commands.command(name="load", hidden=True) @commands.command(name="load", hidden=True)
@user_is_bot_admin() @user_is_bot_admin()
@ -96,7 +95,7 @@ class OwnerCog(commands.Cog):
async def _system(self, ctx): async def _system(self, ctx):
if ctx.invoked_subcommand is None: if ctx.invoked_subcommand is None:
await ctx.send( await ctx.send(
f"Usage: `system <subcommand>`\n" "Usage: `system <subcommand>`\n"
+ "Subcommands: `restart`, `update`" + "Subcommands: `restart`, `update`"
) )
@ -173,13 +172,11 @@ class OwnerCog(commands.Cog):
@_admin.command(name="add", hidden=True) @_admin.command(name="add", hidden=True)
@commands.is_owner() @commands.is_owner()
async def _add(self, ctx, user: User): 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.") await ctx.send(f"{user.mention} is already an admin.")
return return
self.admins.append(user.id) self.admins.value.append(user.id)
self.db.mongo.jarvis.config.update_one( self.admins.update()
{"key": "admins"}, {"$set": {"value": self.admins}}
)
reload_config() reload_config()
await ctx.send( await ctx.send(
f"{user.mention} is now an admin. Use this power carefully." 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) @_admin.command(name="remove", hidden=True)
@commands.is_owner() @commands.is_owner()
async def _remove(self, ctx, user: User): 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.") await ctx.send(f"{user.mention} is not an admin.")
return return
self.admins.remove(user.id) self.admins.value.remove(user.id)
self.db.mongo.jarvis.config.update_one( self.admins.update()
{"key": "admins"}, {"$set": {"value": self.admins}}
)
reload_config() reload_config()
await ctx.send(f"{user.mention} is no longer an admin.") await ctx.send(f"{user.mention} is no longer an admin.")
def resolve_variable(self, variable): def resolve_variable(self, variable):
if hasattr(variable, "__iter__"): if hasattr(variable, "__iter__"):
var_length = len(list(variable)) var_length = len(list(variable))
if (var_length > 100) and (not isinstance(variable, str)): if (var_length > 100) and (not isinstance(variable, str)):
return f"<a {type(variable).__name__} iterable with more than 100 values ({var_length})>" return (
elif not var_length: f"<a {type(variable).__name__} iterable "
return f"<an empty {type(variable).__name__} iterable>" + f"with more than 100 values ({var_length})>"
)
elif not var_length:
return f"<an empty {type(variable).__name__} iterable>"
if (not variable) and (not isinstance(variable, bool)): if (not variable) and (not isinstance(variable, bool)):
return f"<an empty {type(variable).__name__} object>" return f"<an empty {type(variable).__name__} object>"
return ( return (
variable variable
if (len(f"{variable}") <= 1000) if (len(f"{variable}") <= 1000)
else f"<a long {type(variable).__name__} object with the length of {len(f'{variable}'):,}>" else f"<a long {type(variable).__name__} object "
) + f"with the length of {len(f'{variable}'):,}>"
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): def prepare(self, string):
if not isinstance(ctx.message.channel, DMChannel): 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 return
code = self.prepare(code)
args = {
"discord": discord,
"sauce": getsource,
"sys": sys,
"os": os,
"imp": __import__,
"this": self,
"ctx": ctx,
}
try: if isinstance(response, str):
exec( response = response.replace("`", "")
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): await ctx.send(
response = response.replace("`", "") 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( del args, code
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
def setup(bot): def setup(bot):

15
jarvis/cogs/settings.py
View file

@ -1,29 +1,20 @@
import discord
from discord import Role, TextChannel from discord import Role, TextChannel
from discord.ext import commands from discord.ext import commands
from discord_slash import cog_ext from discord_slash import cog_ext
from discord_slash.utils.manage_commands import create_option from discord_slash.utils.manage_commands import create_option
import jarvis from jarvis.db.types import Setting
from jarvis.config import get_config
from jarvis.utils import db
from jarvis.utils.permissions import admin_or_permissions from jarvis.utils.permissions import admin_or_permissions
class SettingsCog(commands.Cog): class SettingsCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
config = get_config()
self.db = db.DBManager(config.mongo).mongo
self.cache = {} self.cache = {}
def update_settings(self, setting, value, guild): def update_settings(self, setting, value, guild):
settings = {"setting": setting, "value": value, "guild": guild} setting = Setting(setting=setting, value=value, guild=guild)
updated = self.db.jarvis.settings.update_one( updated = setting.update()
{"setting": setting, "guild": guild},
{"$set": settings},
upsert=True,
)
return updated is not None return updated is not None

133
jarvis/cogs/starboard.py
View file

@ -1,19 +1,10 @@
from datetime import datetime from discord import TextChannel
import aiohttp
import discord
from discord import Message, TextChannel
from discord.ext import commands from discord.ext import commands
from discord.utils import find
from discord_slash import SlashContext, cog_ext from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option from discord_slash.utils.manage_commands import create_option
import jarvis from jarvis.db.types import Star, Starboard
from jarvis.config import get_config
from jarvis.utils import build_embed 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 = [ supported_images = [
"image/png", "image/png",
@ -27,7 +18,6 @@ supported_images = [
class StarboardCog(commands.Cog): class StarboardCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
self.db = DBManager(get_config().mongo).mongo
@cog_ext.cog_subcommand( @cog_ext.cog_subcommand(
base="starboard", base="starboard",
@ -36,13 +26,11 @@ class StarboardCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _list(self, ctx): async def _list(self, ctx):
starboards = [ starboards = Starboard.get_many(guild=ctx.guild.id)
x for x in self.db.jarvis.starboard.find({"guild": ctx.guild.id})
]
if starboards != []: if starboards != []:
message = "Available Starboards:\n" message = "Available Starboards:\n"
for s in starboards: for s in starboards:
message += f"<#{s['target']}>\n" message += f"<#{s.channel}>\n"
await ctx.send(message) await ctx.send(message)
else: else:
await ctx.send("No Starboards available.") await ctx.send("No Starboards available.")
@ -53,42 +41,39 @@ class StarboardCog(commands.Cog):
description="Create a starboard", description="Create a starboard",
options=[ options=[
create_option( create_option(
name="target", name="channel",
description="Target channel", description="Starboard channel",
option_type=7, option_type=7,
required=True, required=True,
), ),
], ],
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _create(self, ctx, target: TextChannel): async def _create(self, ctx, channel: TextChannel):
if target not in ctx.guild.channels: if channel not in ctx.guild.channels:
await ctx.send( await ctx.send(
"Target channel not in guild. Choose an existing channel.", "Channel not in guild. Choose an existing channel.",
hidden=True, hidden=True,
) )
return return
if not isinstance(target, TextChannel): if not isinstance(channel, TextChannel):
await ctx.send("Target must be a TextChannel", hidden=True) await ctx.send("Channel must be a TextChannel", hidden=True)
return return
exists = self.db.jarvis.starboard.find_one( exists = Starboard.get(channel=channel.id, guild=ctx.guild.id)
{"target": target.id, "guild": ctx.guild.id}
)
if exists: if exists:
await ctx.send( await ctx.send(
f"Starboard already exists at {target.mention}.", hidden=True f"Starboard already exists at {channel.mention}.", hidden=True
) )
return return
self.db.jarvis.starboard.insert_one( _ = Starboard(
{ guild=ctx.guild.id,
"guild": ctx.guild.id, channel=channel.id,
"target": target.id, admin=ctx.author.id,
"admin": ctx.author.id, ).insert()
"time": datetime.now(), 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( @cog_ext.cog_subcommand(
base="starboard", base="starboard",
@ -105,14 +90,11 @@ class StarboardCog(commands.Cog):
) )
@commands.has_permissions(administrator=True) @commands.has_permissions(administrator=True)
async def _delete(self, ctx, channel: TextChannel): async def _delete(self, ctx, channel: TextChannel):
deleted = self.db.jarvis.starboard.delete_one( deleted = Starboard.get(
{ channel=channel.id, guild=ctx.guild.id
"target": channel.id, ).delete()
"guild": ctx.guild.id,
}
)
if deleted: if deleted:
self.db.jarvis.stars.delete_many({"starboard": channel.id}) _ = Star.delete_many(starboard=channel.id)
await ctx.send( await ctx.send(
f"Starboard deleted from {channel.mention}.", hidden=True f"Starboard deleted from {channel.mention}.", hidden=True
) )
@ -157,9 +139,7 @@ class StarboardCog(commands.Cog):
): ):
if not channel: if not channel:
channel = ctx.channel channel = ctx.channel
exists = self.db.jarvis.starboard.find_one( exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id)
{"target": starboard.id, "guild": ctx.guild.id}
)
if not exists: if not exists:
await ctx.send( await ctx.send(
f"Starboard does not exist in {starboard.mention}. " f"Starboard does not exist in {starboard.mention}. "
@ -170,13 +150,11 @@ class StarboardCog(commands.Cog):
message = await channel.fetch_message(int(message)) message = await channel.fetch_message(int(message))
exists = self.db.jarvis.stars.find_one( exists = Star.get(
{ message=message.id,
"message": message.id, channel=message.channel.id,
"channel": message.channel.id, guild=message.guild.id,
"guild": message.guild.id, starboard=starboard.id,
"starboard": starboard.id,
}
) )
if exists: if exists:
@ -186,9 +164,9 @@ class StarboardCog(commands.Cog):
) )
return return
count = self.db.jarvis.stars.find( count = len(
{"guild": message.guild.id, "starboard": starboard.id} Star.get_many(guild=message.guild.id, starboard=starboard.id)
).count() )
content = message.content content = message.content
attachments = message.attachments attachments = message.attachments
@ -221,18 +199,15 @@ class StarboardCog(commands.Cog):
star = await starboard.send(embed=embed) star = await starboard.send(embed=embed)
self.db.jarvis.stars.insert_one( _ = Star(
{ index=count,
"index": count, message=message.id,
"message": message.id, channel=message.channel.id,
"channel": message.channel.id, guild=message.guild.id,
"guild": message.guild.id, starboard=starboard.id,
"starboard": starboard.id, admin=ctx.author.id,
"admin": ctx.author.id, star=star.id,
"time": datetime.now(), ).insert()
"star": star.id,
}
)
await ctx.send( await ctx.send(
"Message saved to Starboard.\n" + f"See it in {starboard.mention}" "Message saved to Starboard.\n" + f"See it in {starboard.mention}"
@ -268,9 +243,7 @@ class StarboardCog(commands.Cog):
id: int, id: int,
starboard: TextChannel, starboard: TextChannel,
): ):
exists = self.db.jarvis.starboard.find_one( exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id)
{"target": starboard.id, "guild": ctx.guild.id}
)
if not exists: if not exists:
await ctx.send( await ctx.send(
f"Starboard does not exist in {starboard.mention}. " f"Starboard does not exist in {starboard.mention}. "
@ -279,26 +252,22 @@ class StarboardCog(commands.Cog):
) )
return return
star = self.db.jarvis.stars.find_one( star = Star.get(
{ starboard=starboard.id,
"starboard": starboard.id, id=id,
"id": id, guild=ctx.guild.id,
"guild": ctx.guild.id, active=True,
"active": True,
}
) )
if not star: if not star:
await ctx.send(f"No star exists with id {id}", hidden=True) await ctx.send(f"No star exists with id {id}", hidden=True)
return return
message = await starboard.fetch_message(star["star"]) message = await starboard.fetch_message(star.star)
if message: if message:
await message.delete() await message.delete()
self.db.jarvis.stars.update_one( star.active = False
{"starboard": starboard.id, "id": id, "guild": ctx.guild.id}, star.update()
{"$set": {"active": False}},
)
await ctx.send(f"Star {id} deleted") await ctx.send(f"Star {id} deleted")

30
jarvis/cogs/verify.py
View file

@ -1,15 +1,11 @@
from random import randint from random import randint
from discord import Role
from discord.ext import commands from discord.ext import commands
from discord_slash import ComponentContext, SlashContext, cog_ext from discord_slash import ComponentContext, SlashContext, cog_ext
from discord_slash.model import ButtonStyle from discord_slash.model import ButtonStyle
from discord_slash.utils import manage_commands, manage_components from discord_slash.utils import manage_components
from discord_slash.utils.manage_commands import create_option
import jarvis from jarvis.db.types import Setting
from jarvis.config import get_config
from jarvis.utils.db import DBManager
def create_layout(): def create_layout():
@ -33,7 +29,6 @@ def create_layout():
class VerifyCog(commands.Cog): class VerifyCog(commands.Cog):
def __init__(self, bot): def __init__(self, bot):
self.bot = bot self.bot = bot
self.db = DBManager(get_config().mongo).mongo
@cog_ext.cog_slash( @cog_ext.cog_slash(
name="verify", name="verify",
@ -41,15 +36,13 @@ class VerifyCog(commands.Cog):
) )
async def _verify(self, ctx: SlashContext): async def _verify(self, ctx: SlashContext):
await ctx.defer() await ctx.defer()
role = self.db.jarvis.settings.find_one( role = Setting.get(guild=ctx.guild.id, setting="verified")
{"setting": "verified", "guild": ctx.guild.id}
)
if not role: if not role:
await ctx.send( await ctx.send(
"This guild has not enabled verification", delete_after=5 "This guild has not enabled verification", delete_after=5
) )
return 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) await ctx.send("You are already verified.", delete_after=5)
return return
components = create_layout() components = create_layout()
@ -74,21 +67,18 @@ class VerifyCog(commands.Cog):
for c in components: for c in components:
for c2 in c["components"]: for c2 in c["components"]:
c2["disabled"] = True c2["disabled"] = True
setting = self.db.jarvis.settings.find_one( setting = Setting.get(guild=ctx.guild.id, setting="verified")
{"setting": "verified", "guild": ctx.guild.id} role = ctx.guild.get_role(setting.value)
)
role = ctx.guild.get_role(setting["value"])
await ctx.author.add_roles(role, reason="Verification passed") await ctx.author.add_roles(role, reason="Verification passed")
setting = self.db.jarvis.settings.find_one( setting = Setting.get(guild=ctx.guild.id, setting="unverified")
{"setting": "unverified", "guild": ctx.guild.id}
)
if setting: if setting:
role = ctx.guild.get_role(setting["value"]) role = ctx.guild.get_role(setting.value)
await ctx.author.remove_roles( await ctx.author.remove_roles(
role, reason="Verification passed" role, reason="Verification passed"
) )
await ctx.edit_origin( 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=manage_components.spread_to_rows(
*components, max_in_row=5 *components, max_in_row=5
), ),

2
jarvis/config.py
View file

@ -1,6 +1,6 @@
from yaml import load from yaml import load
from jarvis.utils.db import DBManager from jarvis.db import DBManager
try: try:
from yaml import CLoader as Loader from yaml import CLoader as Loader

0
jarvis/utils/db.py → jarvis/db/__init__.py
View file

539
jarvis/db/types.py Normal file
View file

@ -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)

2
jarvis/utils/__init__.py
View file

@ -7,7 +7,7 @@ from discord.ext import commands
import jarvis.cogs import jarvis.cogs
import jarvis.config import jarvis.config
import jarvis.utils.db import jarvis.db
__all__ = ["field", "db"] __all__ = ["field", "db"]