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

Merge branch 'mongoengine'

Browse source
This commit is contained in:
Zeva Rose 2021-08-05 17:01:29 -06:00
commit a2dba0d437
34 changed files with 690 additions and 1121 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
jarvis/__init__.py
View file

@ -5,11 +5,11 @@ from discord import Intents
from discord.ext import commands
from discord.utils import find
from discord_slash import SlashCommand
from mongoengine import connect
from psutil import Process
from jarvis import logo, tasks, utils
from jarvis.config import get_config
from jarvis.db import DBManager
from jarvis.events import guild, member, message
if asyncio.get_event_loop().is_closed():
@ -25,11 +25,9 @@ jarvis = commands.Bot(
)
slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True)
jarvis_self = Process()
__version__ = "1.9.6"
__version__ = "1.10.0"
jconfig = get_config()
db = DBManager(jconfig.mongo["connect"]).mongo
jarvis_db = db[jconfig.mongo["database"]]
@jarvis.event
@ -59,28 +57,39 @@ def run(ctx=None):
global restart_ctx
if ctx:
restart_ctx = ctx
connect(
db="ctc2",
alias="ctc2",
authentication_source="admin",
**jconfig.mongo["connect"],
)
connect(
db=jconfig.mongo["database"],
alias="main",
authentication_source="admin",
**jconfig.mongo["connect"],
)
jconfig.get_db_config()
for extension in utils.get_extensions():
jarvis.load_extension(extension)
config = get_config()
print(
" https://discord.com/api/oauth2/authorize?client_id="
+ "{}&permissions=8&scope=bot%20applications.commands".format(
config.client_id
jconfig.client_id
)
)
jarvis.max_messages = config.max_messages
jarvis.max_messages = jconfig.max_messages
tasks.init()
# Add event listeners
if config.events:
if jconfig.events:
listeners = [
guild.GuildEventHandler(jarvis),
member.MemberEventHandler(jarvis),
message.MessageEventHandler(jarvis),
]
jarvis.run(config.token, bot=True, reconnect=True)
jarvis.run(jconfig.token, bot=True, reconnect=True)
for cog in jarvis.cogs:
session = getattr(cog, "_session", None)
if session:

19
jarvis/cogs/admin/ban.py
View file

@ -9,7 +9,7 @@ from discord_slash import SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from discord_slash.utils.manage_commands import create_choice, create_option
from jarvis.db.types import Ban, Unban
from jarvis.db.models import Ban, Unban
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -40,7 +40,7 @@ class BanCog(CacheCog):
type=type,
duration=duration,
active=active,
).insert()
).save()
embed = build_embed(
title="User Banned",
@ -68,7 +68,7 @@ class BanCog(CacheCog):
guild=ctx.guild.id,
admin=ctx.author.id,
reason=reason,
).insert()
).save()
embed = build_embed(
title="User Unbanned",
@ -280,9 +280,9 @@ class BanCog(CacheCog):
# We take advantage of the previous checks to save CPU cycles
if not discord_ban_info:
if isinstance(user, int):
database_ban_info = Ban.get(
database_ban_info = Ban.objects(
guild=ctx.guild.id, user=user, active=True
)
).first()
else:
search = {
"guild": ctx.guild.id,
@ -291,7 +291,7 @@ class BanCog(CacheCog):
}
if discrim:
search["discrim"] = discrim
database_ban_info = Ban.get(**search)
database_ban_info = Ban.objects(**search).first()
if not discord_ban_info and not database_ban_info:
await ctx.send(f"Unable to find user {orig_user}", hidden=True)
@ -308,7 +308,7 @@ class BanCog(CacheCog):
)
else:
database_ban_info.active = False
database_ban_info.update()
database_ban_info.save()
_ = Unban(
user=database_ban_info.user,
username=database_ban_info.username,
@ -316,7 +316,7 @@ class BanCog(CacheCog):
guild=ctx.guild.id,
admin=ctx.author.id,
reason=reason,
).insert()
).save()
await ctx.send(
"Unable to find user in Discord, "
+ "but removed entry from database."
@ -371,8 +371,7 @@ class BanCog(CacheCog):
search["active"] = True
if type > 0:
search["type"] = types[type]
bans = Ban.get_many(**search)
bans.sort(key=lambda x: x.created_at, reverse=True)
bans = Ban.objects(**search).order_by("-created_at")
db_bans = []
fields = []
for ban in bans:

4
jarvis/cogs/admin/kick.py
View file

@ -2,7 +2,7 @@ from discord import User
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Kick
from jarvis.db.models import Kick
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -86,4 +86,4 @@ class KickCog(CacheCog):
reason=reason,
admin=ctx.author.id,
guild=ctx.guild.id,
).insert()
).save()

10
jarvis/cogs/admin/lock.py
View file

@ -5,7 +5,7 @@ from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Lock
from jarvis.db.models import Lock
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.permissions import admin_or_permissions
@ -97,7 +97,7 @@ class LockCog(CacheCog):
admin=ctx.author.id,
reason=reason,
duration=duration,
).insert()
).save()
await ctx.send(f"{channel.mention} locked for {duration} minute(s)")
@cog_ext.cog_slash(
@ -120,7 +120,9 @@ class LockCog(CacheCog):
):
if not channel:
channel = ctx.channel
lock = Lock.get(guild=ctx.guild.id, channel=channel.id, active=True)
lock = Lock.objects(
guild=ctx.guild.id, channel=channel.id, active=True
).first()
if not lock:
await ctx.send(f"{channel.mention} not locked.", hidden=True)
return
@ -130,5 +132,5 @@ class LockCog(CacheCog):
except Exception:
continue # Just continue on error
lock.active = False
lock.update()
lock.save()
await ctx.send(f"{channel.mention} unlocked")

48
jarvis/cogs/admin/lockdown.py
View file

@ -1,13 +1,10 @@
from datetime import datetime
import pymongo
from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.config import get_config
from jarvis.db import DBManager
from jarvis.db.types import Lock
from jarvis.db.models import Lock
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.permissions import admin_or_permissions
@ -15,10 +12,6 @@ from jarvis.utils.permissions import admin_or_permissions
class LockdownCog(CacheCog):
def __init__(self, bot: commands.Bot):
super().__init__(bot)
config = get_config()
self.db = DBManager(config.mongo["connect"]).mongo[
config.mongo["database"]
]
@cog_ext.cog_subcommand(
base="lockdown",
@ -63,20 +56,18 @@ class LockdownCog(CacheCog):
except Exception:
continue # Just continue on error
updates.append(
pymongo.InsertOne(
{
"channel": channel.id,
"guild": ctx.guild.id,
"admin": ctx.author.id,
"reason": reason,
"duration": duration,
"active": True,
"created_at": datetime.utcnow(),
}
Lock(
channel=channel.id,
guild=ctx.guild.id,
admin=ctx.author.id,
reason=reason,
duration=duration,
active=True,
created_at=datetime.utcnow(),
)
)
if updates:
self.db.locks.bulk_write(updates)
Lock.objects().insert(updates)
await ctx.send(f"Server locked for {duration} minute(s)")
@cog_ext.cog_subcommand(
@ -91,8 +82,8 @@ class LockdownCog(CacheCog):
):
channels = ctx.guild.channels
roles = ctx.guild.roles
updates = []
locks = Lock.get_many(guild=ctx.guild.id, active=True)
update = False
locks = Lock.objects(guild=ctx.guild.id, active=True)
if not locks:
await ctx.send("No lockdown detected.", hidden=True)
return
@ -103,16 +94,9 @@ class LockdownCog(CacheCog):
await self._unlock_channel(channel, role, ctx.author)
except Exception:
continue # Just continue on error
updates.append(
pymongo.UpdateOne(
{
"channel": channel.id,
"guild": ctx.guild.id,
"admin": ctx.author.id,
},
{"$set": {"active": False}},
update = True
if update:
Lock.objects(guild=ctx.guild.id, active=True).update(
set__active=False
)
)
if updates:
self.db.locks.bulk_write(updates)
await ctx.send("Server unlocked")

19
jarvis/cogs/admin/mute.py
View file

@ -4,7 +4,7 @@ from discord.utils import get
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Mute, Setting
from jarvis.db.models import Mute, Setting
from jarvis.utils import build_embed
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions
@ -32,7 +32,7 @@ class MuteCog(commands.Cog):
),
create_option(
name="duration",
description="Mute duration",
description="Duration of mute in minutes, default 30",
option_type=4,
required=False,
),
@ -73,7 +73,7 @@ class MuteCog(commands.Cog):
guild=ctx.guild.id,
duration=duration,
active=True if duration >= 0 else False,
).insert()
).save()
embed = build_embed(
title="User Muted",
@ -102,7 +102,9 @@ class MuteCog(commands.Cog):
)
@admin_or_permissions(mute_members=True)
async def _unmute(self, ctx: SlashContext, user: Member):
mute_setting = Setting.get(guild=ctx.guild.id, setting="mute")
mute_setting = Setting.objects(
guild=ctx.guild.id, setting="mute"
).first()
if not mute_setting:
await ctx.send(
"Please configure a mute role with "
@ -118,12 +120,11 @@ class MuteCog(commands.Cog):
await ctx.send("User is not muted.", hidden=True)
return
mutes = Mute.get_many(guild=ctx.guild.id, user=user.id)
for mute in mutes:
mute.active = False
mute.update()
_ = Mute.objects(guild=ctx.guild.id, user=user.id).update(
set__active=False
)
embed = build_embed(
title="User Unmwaruted",
title="User Unmuted",
description=f"{user.mention} has been unmuted",
fields=[],
)

17
jarvis/cogs/admin/purge.py
View file

@ -39,7 +39,7 @@ class PurgeCog(commands.Cog):
guild=ctx.guild.id,
admin=ctx.author.id,
count=amount,
).insert()
).save()
@cog_ext.cog_subcommand(
base="autopurge",
@ -73,17 +73,18 @@ class PurgeCog(commands.Cog):
elif delay > 300:
await ctx.send("Delay must be < 5 minutes", hidden=True)
return
autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id)
autopurge = Autopurge.objects(
guild=ctx.guild.id, channel=channel.id
).first()
if autopurge:
await ctx.send("Autopurge already exists.", hidden=True)
return
autopurge = Autopurge(
_ = Autopurge(
guild=ctx.guild.id,
channel=channel.id,
admin=ctx.author.id,
delay=delay,
)
autopurge.insert()
).save()
await ctx.send(
f"Autopurge set up on {channel.mention}, "
+ f"delay is {delay} seconds"
@ -104,7 +105,7 @@ class PurgeCog(commands.Cog):
)
@admin_or_permissions(manage_messages=True)
async def _autopurge_remove(self, ctx: SlashContext, channel: TextChannel):
autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id)
autopurge = Autopurge.objects(guild=ctx.guild.id, channel=channel.id)
if not autopurge:
await ctx.send("Autopurge does not exist.", hidden=True)
return
@ -134,12 +135,12 @@ class PurgeCog(commands.Cog):
async def _autopurge_update(
self, ctx: SlashContext, channel: TextChannel, delay: int
):
autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id)
autopurge = Autopurge.objects(guild=ctx.guild.id, channel=channel.id)
if not autopurge:
await ctx.send("Autopurge does not exist.", hidden=True)
return
autopurge.delay = delay
autopurge.update()
autopurge.save()
await ctx.send(
f"Autopurge delay updated to {delay} seconds on {channel.mention}."
)

43
jarvis/cogs/admin/roleping.py
View file

@ -2,12 +2,11 @@ from datetime import datetime, timedelta
from ButtonPaginator import Paginator
from discord import Member, Role
from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Roleping
from jarvis.db.models import Roleping
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -33,23 +32,19 @@ class RolepingCog(CacheCog):
)
@admin_or_permissions(manage_guild=True)
async def _roleping_add(self, ctx: SlashContext, role: Role):
roleping = Roleping.get(guild=ctx.guild.id, role=role.id)
if not roleping:
roleping = Roleping(
roleping = Roleping.objects(guild=ctx.guild.id, role=role.id).first()
if roleping:
await ctx.send(
f"Role `{role.name}` already in roleping.", hidden=True
)
return
_ = Roleping(
role=role.id,
guild=ctx.guild.id,
admin=ctx.author.id,
active=True,
bypass={"roles": [], "users": []},
)
else:
await ctx.send(
f"Role `{role.name}` already in roleping.", hidden=True
)
return
roleping.insert()
).save()
await ctx.send(f"Role `{role.name}` added to roleping.")
@cog_ext.cog_subcommand(
@ -67,7 +62,7 @@ class RolepingCog(CacheCog):
)
@admin_or_permissions(manage_guild=True)
async def _roleping_remove(self, ctx: SlashContext, role: Role):
roleping = Roleping.get(guild=ctx.guild.id, role=role.id)
roleping = Roleping.objects(guild=ctx.guild.id, role=role.id)
if not roleping:
await ctx.send("Roleping does not exist", hidden=True)
return
@ -91,7 +86,7 @@ class RolepingCog(CacheCog):
)
return
rolepings = Roleping.get_many(guild=ctx.guild.id)
rolepings = Roleping.objects(guild=ctx.guild.id)
if not rolepings:
await ctx.send("No rolepings configured", hidden=True)
return
@ -199,7 +194,7 @@ class RolepingCog(CacheCog):
async def _roleping_bypass_user(
self, ctx: SlashContext, user: Member, rping: Role
):
roleping = Roleping.get(guild=ctx.guild.id, role=rping.id)
roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first()
if not roleping:
await ctx.send(
f"Roleping not configured for {rping.mention}", hidden=True
@ -231,7 +226,7 @@ class RolepingCog(CacheCog):
return
roleping.bypass["users"].append(user.id)
roleping.update()
roleping.save()
await ctx.send(
f"{user.nick or user.name} user bypass added for `{rping.name}`"
)
@ -262,7 +257,7 @@ class RolepingCog(CacheCog):
async def _roleping_bypass_role(
self, ctx: SlashContext, role: Role, rping: Role
):
roleping = Roleping.get(guild=ctx.guild.id, role=rping.id)
roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first()
if not roleping:
await ctx.send(
f"Roleping not configured for {rping.mention}", hidden=True
@ -282,7 +277,7 @@ class RolepingCog(CacheCog):
return
roleping.bypass["roles"].append(role.id)
roleping.update()
roleping.save()
await ctx.send(f"{role.name} role bypass added for `{rping.name}`")
@cog_ext.cog_subcommand(
@ -311,7 +306,7 @@ class RolepingCog(CacheCog):
async def _roleping_restore_user(
self, ctx: SlashContext, user: Member, rping: Role
):
roleping = Roleping.get(guild=ctx.guild.id, role=rping.id)
roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first()
if not roleping:
await ctx.send(
f"Roleping not configured for {rping.mention}", hidden=True
@ -323,7 +318,7 @@ class RolepingCog(CacheCog):
return
roleping.bypass["users"].delete(user.id)
roleping.update()
roleping.save()
await ctx.send(
f"{user.nick or user.name} user bypass removed for `{rping.name}`"
)
@ -354,7 +349,7 @@ class RolepingCog(CacheCog):
async def _roleping_restore_role(
self, ctx: SlashContext, role: Role, rping: Role
):
roleping = Roleping.get(guild=ctx.guild.id, role=rping.id)
roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first()
if not roleping:
await ctx.send(
f"Roleping not configured for {rping.mention}", hidden=True
@ -374,5 +369,5 @@ class RolepingCog(CacheCog):
return
roleping.bypass["roles"].append(role.id)
roleping.update()
roleping.save()
await ctx.send(f"{role.name} role bypass added for `{rping.name}`")

26
jarvis/cogs/admin/warning.py
View file

@ -2,12 +2,11 @@ from datetime import datetime, timedelta
from ButtonPaginator import Paginator
from discord import User
from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from discord_slash.utils.manage_commands import create_choice, create_option
from jarvis.db.types import MongoSort, Warning
from jarvis.db.models import Warning
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -63,7 +62,7 @@ class WarningCog(CacheCog):
guild=ctx.guild.id,
duration=duration,
active=True,
).insert()
).save()
fields = [Field("Reason", reason, False)]
embed = build_embed(
title="Warning",
@ -112,19 +111,20 @@ class WarningCog(CacheCog):
hidden=True,
)
return
warnings = Warning.get_many(
warnings = Warning.objects(
user=user.id,
guild=ctx.guild.id,
sort=MongoSort(direction="desc", key="created_at"),
)
active_warns = list(filter(lambda x: x.active, warnings))
).order_by("-created_at")
active_warns = Warning.objects(
user=user.id, guild=ctx.guild.id, active=False
).order_by("-created_at")
pages = []
if active:
if len(active_warns) == 0:
if active_warns.count() == 0:
embed = build_embed(
title="Warnings",
description=f"{len(warnings)} total | 0 currently active",
description=f"{warnings.count()} total | 0 currently active",
fields=[],
)
embed.set_author(name=user.name, icon_url=user.avatar_url)
@ -151,8 +151,8 @@ class WarningCog(CacheCog):
for i in range(0, len(fields), 5):
embed = build_embed(
title="Warnings",
description=f"{len(warnings)} total | "
+ f"{len(active_warns)} currently active",
description=f"{warnings.count()} total | "
+ f"{active_warns.count()} currently active",
fields=fields[i : i + 5],
)
embed.set_author(
@ -179,8 +179,8 @@ class WarningCog(CacheCog):
for i in range(0, len(fields), 5):
embed = build_embed(
title="Warnings",
description=f"{len(warnings)} total | "
+ f"{len(active_warns)} currently active",
description=f"{warnings.count()} total | "
+ f"{active_warns.count()} currently active",
fields=fields[i : i + 5],
)
embed.set_author(

31
jarvis/cogs/autoreact.py
View file

@ -7,7 +7,7 @@ from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.data.unicode import emoji_list
from jarvis.db.types import Autoreact
from jarvis.db.models import Autoreact
from jarvis.utils.permissions import admin_or_permissions
@ -34,20 +34,21 @@ class AutoReactCog(commands.Cog):
if not isinstance(channel, TextChannel):
await ctx.send("Channel must be a text channel", hidden=True)
return
exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
exists = Autoreact.objects(
guild=ctx.guild.id, channel=channel.id
).first()
if exists:
await ctx.send(
f"Autoreact already exists for {channel.mention}.", hidden=True
)
return
autoreact = Autoreact(
_ = Autoreact(
guild=ctx.guild.id,
channel=channel.id,
reactions=[],
admin=ctx.author.id,
)
autoreact.insert()
).save()
await ctx.send(f"Autoreact created for {channel.mention}!")
@cog_ext.cog_subcommand(
@ -65,7 +66,9 @@ class AutoReactCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _autoreact_delete(self, ctx, channel: TextChannel):
exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id).delete()
exists = Autoreact.objects(
guild=ctx.guild.id, channel=channel.id
).delete()
if exists:
await ctx.send(f"Autoreact removed from {channel.mention}")
else:
@ -111,7 +114,9 @@ class AutoReactCog(commands.Cog):
"Please use a custom emote from this server.", hidden=True
)
return
exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
exists = Autoreact.objects(
guild=ctx.guild.id, channel=channel.id
).first()
if not exists:
await ctx.send(
"Please create autoreact first with "
@ -132,7 +137,7 @@ class AutoReactCog(commands.Cog):
)
return
exists.reactions.append(emote)
exists.update()
exists.save()
await ctx.send(f"Added {emote} to {channel.mention} autoreact.")
@cog_ext.cog_subcommand(
@ -156,7 +161,9 @@ class AutoReactCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str):
exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
exists = Autoreact.objects(
guild=ctx.guild.id, channel=channel.id
).first()
if not exists:
await ctx.send(
"Please create autoreact first with "
@ -171,7 +178,7 @@ class AutoReactCog(commands.Cog):
)
return
exists.reactions.remove(emote)
exists.update()
exists.save()
await ctx.send(f"Removed {emote} from {channel.mention} autoreact.")
@cog_ext.cog_subcommand(
@ -189,7 +196,9 @@ class AutoReactCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _autoreact_list(self, ctx, channel: TextChannel):
exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id)
exists = Autoreact.objects(
guild=ctx.guild.id, channel=channel.id
).first()
if not exists:
await ctx.send(
"Please create autoreact first with "

20
jarvis/cogs/ctc2.py
View file

@ -11,8 +11,7 @@ from discord.utils import find
from discord_slash import SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from jarvis.config import get_config
from jarvis.db import DBManager
from jarvis.db.models import Guess
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -29,8 +28,6 @@ invites = re.compile(
class CTCCog(CacheCog):
def __init__(self, bot):
super().__init__(bot)
mconf = get_config().mongo
self.db = DBManager(mconf["connect"]).mongo
self._session = aiohttp.ClientSession()
self.url = "https://completethecodetwo.cards/pw"
@ -73,7 +70,7 @@ class CTCCog(CacheCog):
hidden=True,
)
return
guessed = self.db.ctc2.guesses.find_one({"guess": guess})
guessed = Guess.objects(guess=guess).first()
if guessed:
await ctx.send("Already guessed, dipshit.", hidden=True)
return
@ -86,9 +83,7 @@ class CTCCog(CacheCog):
correct = True
else:
await ctx.send("Nope.", hidden=True)
self.db.ctc2.guesses.insert_one(
{"guess": guess, "user": ctx.author.id, "correct": correct}
)
_ = Guess(guess=guess, user=ctx.author.id, correct=correct).save()
@cog_ext.cog_subcommand(
base="ctc2",
@ -107,17 +102,16 @@ class CTCCog(CacheCog):
hidden=True,
)
return
guesses = self.db.ctc2.guesses.find().sort(
[("correct", pymongo.DESCENDING), ("_id", pymongo.DESCENDING)]
)
guesses = Guess.objects().order_by("-correct", "-id")
fields = []
for guess in guesses:
user = ctx.guild.get_member(guess["user"])
if not user:
user = self.bot.fetch_user(guess["user"])
user = await self.bot.fetch_user(guess["user"])
if not user:
user = "[redacted]"
if isinstance(user, User) or isinstance(user, Member):
if isinstance(user, (Member, User)):
user = user.name + "#" + user.discriminator
name = "Correctly" if guess["correct"] else "Incorrectly"
name += " guessed by: " + user

283
jarvis/cogs/dev.py
View file

@ -2,12 +2,13 @@ import base64
import hashlib
import re
import subprocess
import uuid
import uuid as uuidpy
import ulid
import ulid as ulidpy
from bson import ObjectId
from discord.ext import commands
from discord_slash import cog_ext
from discord_slash.utils.manage_commands import create_choice, create_option
from jarvis.utils import build_embed, convert_bytesize
from jarvis.utils.field import Field
@ -39,7 +40,12 @@ UUID_VERIFY = re.compile(
re.IGNORECASE,
)
UUID_GET = {3: uuid.uuid3, 5: uuid.uuid5}
invites = re.compile(
r"(?:https?://)?(?:www.)?(?:discord.(?:gg|io|me|li)|discord(?:app)?.com/invite)/([^\s/]+?)(?=\b)",
flags=re.IGNORECASE,
)
UUID_GET = {3: uuidpy.uuid3, 5: uuidpy.uuid5}
def hash_obj(hash, data, text: bool = True) -> str:
@ -64,24 +70,36 @@ class DevCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
async def _hash(self, ctx, method: str, *, data: str = None):
if method not in supported_hashes:
algo_txt = ", ".join(f"`{x}`" for x in supported_hashes)
@cog_ext.cog_slash(
name="hash",
description="Hash some data",
options=[
create_option(
name="method",
description="Hash method",
option_type=3,
required=True,
choices=[
create_choice(name=x, value=x) for x in supported_hashes
],
),
create_option(
name="data",
description="Data to hash",
option_type=3,
required=True,
),
],
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _hash(self, ctx, method: str, data: str):
if not data:
await ctx.send(
"Unsupported hash algorithm. Supported:\n" + algo_txt,
"No data to hash",
hidden=True,
)
return
if not data and len(ctx.message.attachments) == 0:
await ctx.send(
"No data to hash. Either attach a file or send text to hash"
)
return
text = True
if len(ctx.message.attachments) > 0:
text = False
data = await ctx.message.attachments[0].read()
# Default to sha256, just in case
hash = getattr(hashlib, method, hashlib.sha256)()
hex = hash_obj(hash, data, text)
@ -98,206 +116,161 @@ class DevCog(commands.Cog):
)
await ctx.send(embed=embed)
@commands.command(name="hash")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _hash_pref(self, ctx, method: str, *, data: str = None):
await self._hash(ctx, method, data=data)
@cog_ext.cog_slash(
name="hash",
description="Hash some data",
name="uuid",
description="Generate a UUID",
options=[
create_option(
name="version",
description="UUID version",
option_type=3,
required=True,
choices=[
create_choice(name=x, value=x) for x in ["3", "4", "5"]
],
),
create_option(
name="data",
description="Data for UUID version 3,5",
option_type=3,
required=False,
),
],
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _hash_slash(self, ctx, method: str, *, data: str = None):
await self._hash(ctx, method, data=data)
async def _uuid(self, ctx, version: str = None, data: str = None):
if not version:
await ctx.send("Supported UUID versions: 3, 4, 5", hidden=True)
return
if version not in ["3", "4", "5"]:
await ctx.send("Supported UUID versions: 3, 4, 5", hidden=True)
return
async def _uuid(self, ctx, version: str, data: str = None):
version = int(version)
if version in [3, 5] and not data:
await ctx.send(f"UUID{version} requires data.", hidden=True)
return
if version == 4:
await ctx.send(f"UUID4: `{uuid.uuid4()}`")
await ctx.send(f"UUID4: `{uuidpy.uuid4()}`")
else:
to_send = None
if OID_VERIFY.match(data):
to_send = UUID_GET[version](uuid.NAMESPACE_OID, data)
to_send = UUID_GET[version](uuidpy.NAMESPACE_OID, data)
elif URL_VERIFY.match(data):
to_send = UUID_GET[version](uuid.NAMESPACE_URL, data)
to_send = UUID_GET[version](uuidpy.NAMESPACE_URL, data)
elif DN_VERIFY.match(data):
to_send = UUID_GET[version](uuid.NAMESPACE_X500, data)
to_send = UUID_GET[version](uuidpy.NAMESPACE_X500, data)
else:
to_send = UUID_GET[version](uuid.NAMESPACE_DNS, data)
to_send = UUID_GET[version](uuidpy.NAMESPACE_DNS, data)
await ctx.send(f"UUID{version}: `{to_send}`")
@commands.command(name="uuid")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _uuid_pref(self, ctx, version: str = None, data: str = None):
await self._uuid(ctx, version, data)
@cog_ext.cog_slash(
name="uuid",
description="Generate a UUID",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _uuid_slash(self, ctx, version: str = None, data: str = None):
await self._uuid(ctx, version, data)
async def _objectid(self, ctx):
"""Generates new bson.ObjectId"""
await ctx.send(f"ObjectId: `{str(ObjectId())}`")
@commands.command(name="objectid")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _objectid_pref(self, ctx):
await self._objectid(ctx)
@cog_ext.cog_slash(
name="objectid",
description="Generate an ObjectID",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _objectid_slash(self, ctx):
await self._objectid(ctx)
async def _ulid(self, ctx):
"""Generates a new ULID"""
await ctx.send(f"ULID: `{ulid.new().str}`")
@commands.command(name="ulid")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _ulid_pref(self, ctx):
await self._ulid(ctx)
async def _objectid(self, ctx):
"""Generates new bson.ObjectId"""
await ctx.send(f"ObjectId: `{str(ObjectId())}`")
@cog_ext.cog_slash(
name="ulid",
description="Generate a ULID",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _ulid_slash(self, ctx):
await self._ulid(ctx)
async def _uuid2ulid(self, ctx: commands.Context, u):
"""Converts a UUID to a ULID"""
if UUID_VERIFY.match(u):
u = ulid.parse(u)
await ctx.send(f"ULID: `{u.str}`")
else:
await ctx.send("Invalid UUID")
@commands.command(name="uuid2ulid")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _uuid2ulid_pref(self, ctx, u: str):
await self._uuid2ulid(ctx, u)
async def _ulid(self, ctx):
"""Generates a new ULID"""
await ctx.send(f"ULID: `{ulidpy.new().str}`")
@cog_ext.cog_slash(
name="uuid2ulid",
description="Convert a UUID to a ULID",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _uuid2ulid_slash(self, ctx, u: str):
await self._uuid2ulid(ctx, u)
async def _ulid2uuid(self, ctx: commands.Context, u):
"""Converts a ULID to a UUID"""
if ULID_VERIFY.match(u):
u = ulid.parse(u)
await ctx.send(f"UUID: `{u.uuid}`")
async def _uuid2ulid(self, ctx: commands.Context, uuid):
"""Converts a UUID to a ULID"""
if UUID_VERIFY.match(uuid):
u = ulidpy.parse(uuid)
await ctx.send(f"ULID: `{u.str}`")
else:
await ctx.send("Invalid ULID.")
@commands.command(name="ulid2uuid")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _ulid2uuid_pref(self, ctx, u):
await self._ulid2uuid(ctx, u)
await ctx.send("Invalid UUID")
@cog_ext.cog_slash(
name="ulid2uuid",
description="Convert a ULID to a UUID",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _ulid2uuid_slash(self, ctx, u):
await self._ulid2uuid(ctx, u)
async def _ulid2uuid(self, ctx: commands.Context, ulid):
"""Converts a ULID to a UUID"""
if ULID_VERIFY.match(ulid):
ulid = ulidpy.parse(ulid)
await ctx.send(f"UUID: `{ulid.uuid}`")
else:
await ctx.send("Invalid ULID.")
base64_methods = ["b64", "b16", "b32", "a85", "b85"]
@cog_ext.cog_slash(
name="encode",
description="Encode some data",
options=[
create_option(
name="method",
description="Encode method",
option_type=3,
required=True,
choices=[
create_choice(name=x, value=x) for x in base64_methods
],
),
create_option(
name="data",
description="Data to encode",
option_type=3,
required=True,
),
],
)
async def _encode(self, ctx, method: str, data: str):
"""Encodes text with specified encoding method"""
if method not in self.base64_methods:
methods = ", ".join(f"`{x}`" for x in self.base64_methods)
await ctx.send(
"Usage: encode <method> <data>\nSupported methods:\n"
+ methods,
hidden=True,
)
return
method = getattr(base64, method + "encode")
await ctx.send(f"`{method(data.encode('UTF-8')).decode('UTF-8')}`")
@commands.command(name="encode")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _encode_pref(self, ctx, method: str, *, data: str):
await self._encode(ctx, method, data)
@cog_ext.cog_slash(
name="encode",
description="Encode using the base64 module",
name="decode",
description="Decode some data",
options=[
create_option(
name="method",
description="Decode method",
option_type=3,
required=True,
choices=[
create_choice(name=x, value=x) for x in base64_methods
],
),
create_option(
name="data",
description="Data to encode",
option_type=3,
required=True,
),
],
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _encode_slash(self, ctx, method: str, *, data: str):
await self._encode(ctx, method, data)
async def _decode(self, ctx, method: str, data: str):
"""Decodes text with specified encoding method"""
if method not in self.base64_methods:
methods = ", ".join(f"`{x}`" for x in self.base64_methods)
method = getattr(base64, method + "decode")
decoded = method(data.encode("UTF-8")).decode("UTF-8")
if invites.search(decoded):
await ctx.send(
"Usage: decode <method> <data>\nSupported methods:\n"
+ methods,
"Please don't use this to bypass invite restrictions",
hidden=True,
)
return
method = getattr(base64, method + "decode")
await ctx.send(f"`{method(data.encode('UTF-8')).decode('UTF-8')}`")
@commands.command(name="decode")
@commands.cooldown(1, 2, commands.BucketType.user)
async def _decode_pref(self, ctx, method: str, *, data: str):
await self._decode(ctx, method, data)
@cog_ext.cog_slash(
name="decode",
description="Decode using the base64 module",
)
@commands.cooldown(1, 2, commands.BucketType.user)
async def _decode_slash(self, ctx, method: str, *, data: str):
await self._decode(ctx, method, data)
async def _cloc(self, ctx):
output = subprocess.check_output(
["tokei", "-C", "--sort", "code"]
).decode("UTF-8")
await ctx.send(f"```\n{output}\n```")
@commands.command(name="cloc", help="Get J.A.R.V.I.S. lines of code")
@commands.cooldown(1, 30, commands.BucketType.channel)
async def _cloc_pref(self, ctx):
await self._cloc(ctx)
await ctx.send(f"`{decoded}`")
@cog_ext.cog_slash(
name="cloc",
description="Get J.A.R.V.I.S. lines of code",
)
@commands.cooldown(1, 30, commands.BucketType.channel)
async def _cloc_slash(self, ctx):
await ctx.defer()
await self._cloc(ctx)
async def _cloc(self, ctx):
output = subprocess.check_output(
["tokei", "-C", "--sort", "code"]
).decode("UTF-8")
await ctx.send(f"```\n{output}\n```")
def setup(bot):

4
jarvis/cogs/error.py
View file

@ -1,6 +1,8 @@
from discord.ext import commands
from discord_slash import SlashContext
from jarvis import slash
class ErrorHandlerCog(commands.Cog):
def __init__(self, bot):
@ -40,7 +42,7 @@ class ErrorHandlerCog(commands.Cog):
f"Error processing command:\n```{error}```",
hidden=True,
)
ctx.slash.commands[ctx.command].reset_cooldown(ctx)
slash.commands[ctx.command].reset_cooldown(ctx)
def setup(bot):

50
jarvis/cogs/jokes.py
View file

@ -8,8 +8,7 @@ from discord.ext import commands
from discord_slash import cog_ext
import jarvis
from jarvis.config import get_config
from jarvis.db import DBManager
from jarvis.db.models import Joke
from jarvis.utils import build_embed
from jarvis.utils.field import Field
@ -23,10 +22,6 @@ class JokeCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
config = get_config()
self.db = DBManager(config.mongo["connect"]).mongo[
config.mongo["database"]
]
# TODO: Make this a command group with subcommands
async def _joke(self, ctx, id: str = None):
@ -36,48 +31,35 @@ class JokeCog(commands.Cog):
return
# TODO: Add this as a parameter that can be passed in
threshold = 500 # Minimum score
coll = self.db.jokes
result = None
if id:
result = coll.find_one({"id": id})
result = Joke.objects(rid=id).first()
else:
result = list(
coll.aggregate(
[
pipeline = [
{"$match": {"score": {"$gt": threshold}}},
{"$sample": {"size": 1}},
]
)
)[0]
while result["body"] in ["[removed]", "[deleted]"]:
result = list(
coll.aggregate(
[
{"$match": {"score": {"$gt": threshold}}},
{"$sample": {"size": 1}},
]
)
)[0]
result = Joke.objects().aggregate(pipeline)
while result.body in ["[removed]", "[deleted]"]:
result = Joke.objects().aggregate(pipeline)
if result is None:
await ctx.send(
"Humor module failed. Please try again later.", hidden=True
)
return
emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result["body"])
emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result.body)
for match in emotes:
result["body"] = result["body"].replace(
match, html.unescape(match)
)
emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result["title"])
result.body = result.body.replace(match, html.unescape(match))
emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result.title)
for match in emotes:
result["title"] = result["title"].replace(
result.title = result.title.replace(
match, html.unescape(match)
)
body_chunks = []
body = ""
for word in result["body"].split(" "):
for word in result.body.split(" "):
if len(body) + 1 + len(word) > 1024:
body_chunks.append(Field("​", body, False))
body = ""
@ -89,7 +71,7 @@ class JokeCog(commands.Cog):
body += " " + word
desc = ""
title = result["title"]
title = result.title
if len(title) > 256:
new_title = ""
limit = False
@ -106,18 +88,18 @@ class JokeCog(commands.Cog):
body_chunks.append(Field("​", body, False))
fields = body_chunks
fields.append(Field("Score", result["score"]))
fields.append(Field("Score", result.score))
# Field(
# "Created At",
# str(datetime.fromtimestamp(result["created_utc"])),
# ),
fields.append(Field("ID", result["id"]))
fields.append(Field("ID", result.id))
embed = build_embed(
title=title,
description=desc,
fields=fields,
url=f"https://reddit.com/r/jokes/comments/{result['id']}",
timestamp=datetime.fromtimestamp(result["created_utc"]),
url=f"https://reddit.com/r/jokes/comments/{result.rid}",
timestamp=datetime.fromtimestamp(result.created_utc),
)
await ctx.send(embed=embed)
except Exception:

81
jarvis/cogs/modlog/member.py
View file

@ -7,7 +7,7 @@ from discord.utils import find
from jarvis.cogs.modlog.utils import get_latest_log, modlog_embed
from jarvis.config import get_config
from jarvis.db.types import Ban, Kick, MongoSort, Mute, Setting
from jarvis.db.models import Ban, Kick, Mute, Setting, Unban
from jarvis.utils import build_embed
from jarvis.utils.field import Field
@ -18,7 +18,7 @@ class ModlogMemberCog(commands.Cog):
@commands.Cog.listener()
async def on_member_ban(self, guild: discord.Guild, user: discord.User):
modlog = Setting.get(guild=guild.id, setting="modlog")
modlog = Setting.objects(guild=guild.id, setting="modlog").first()
if modlog:
channel = guild.get_channel(modlog.value)
await asyncio.sleep(0.5) # Need to wait for audit log
@ -31,11 +31,14 @@ class ModlogMemberCog(commands.Cog):
log: discord.AuditLogEntry = get_latest_log(auditlog, user)
admin: discord.User = log.user
if admin.id == get_config().client_id:
ban = Ban.get(
ban = (
Ban.objects(
guild=guild.id,
user=user.id,
active=True,
sort=MongoSort(key="created_at", direction="desc"),
)
.sort("-created_at")
.first()
)
admin = guild.get_member(ban.admin)
embed = modlog_embed(
@ -50,7 +53,7 @@ class ModlogMemberCog(commands.Cog):
@commands.Cog.listener()
async def on_member_unban(self, guild: discord.Guild, user: discord.User):
modlog = Setting.get(guild=guild.id, setting="modlog")
modlog = Setting.objects(guild=guild.id, setting="modlog").first()
if modlog:
channel = guild.get_channel(modlog.value)
await asyncio.sleep(0.5) # Need to wait for audit log
@ -63,13 +66,15 @@ class ModlogMemberCog(commands.Cog):
log: discord.AuditLogEntry = get_latest_log(auditlog, user)
admin: discord.User = log.user
if admin.id == get_config().client_id:
ban = Ban.get(
unban = (
Unban.objects(
guild=guild.id,
user=user.id,
active=True,
sort=MongoSort(key="created_at", direction="desc"),
)
admin = guild.get_member(ban.admin)
.sort("-created_at")
.first()
)
admin = guild.get_member(unban.admin)
embed = modlog_embed(
user,
admin,
@ -82,7 +87,7 @@ class ModlogMemberCog(commands.Cog):
@commands.Cog.listener()
async def on_member_remove(self, user: discord.User):
modlog = Setting.get(guild=user.guild.id, setting="modlog")
modlog = Setting.objects(guild=user.guild.id, setting="modlog").first()
if modlog:
channel = user.guild.get_channel(modlog.value)
await asyncio.sleep(0.5) # Need to wait for audit log
@ -95,10 +100,13 @@ class ModlogMemberCog(commands.Cog):
log: discord.AuditLogEntry = get_latest_log(auditlog, user)
admin: discord.User = log.user
if admin.id == get_config().client_id:
kick = Kick.get(
kick = (
Kick.objects(
guild=user.guild.id,
user=user.id,
sort=MongoSort(key="created_at", direction="desc"),
)
.sort("-created_at")
.first()
)
admin = user.guild.get_member(kick.admin)
embed = modlog_embed(
@ -122,11 +130,14 @@ class ModlogMemberCog(commands.Cog):
log: discord.AuditLogEntry = get_latest_log(auditlog, before)
admin: discord.User = log.user
if admin.id == get_config().client_id:
mute = Mute.get(
mute = (
Mute.objects(
guild=before.guild.id,
user=before.id,
active=True,
sort=MongoSort(key="created_at", direction="desc"),
)
.sort("-created_at")
.first()
)
admin = before.guild.get_member(mute.admin)
return modlog_embed(
@ -148,13 +159,15 @@ class ModlogMemberCog(commands.Cog):
log: discord.AuditLogEntry = get_latest_log(auditlog, before)
admin: discord.User = log.user
if admin.id == get_config().client_id:
mute = Mute.get(
mute = (
Mute.objects(
guild=before.guild.id,
user=before.id,
active=True,
sort=MongoSort(key="created_at", direction="desc"),
)
mute = Mute(**mute)
.sort("-created_at")
.first()
)
admin = before.guild.get_member(mute.admin)
return modlog_embed(
member=before,
@ -214,21 +227,41 @@ class ModlogMemberCog(commands.Cog):
async def on_member_update(
self, before: discord.User, after: discord.User
):
modlog = Setting.get(guild=before.guild.id, setting="modlog")
modlog = Setting.objects(
guild=before.guild.id, setting="modlog"
).first()
if modlog:
channel = after.guild.get_channel(modlog.value)
await asyncio.sleep(0.5) # Need to wait for audit log
embed = None
mute = Setting.get(guild=before.guild.id, setting="mute")
verified = Setting.get(guild=before.guild.id, setting="verified")
if mute and before.guild.get_role(mute.value) in after.roles:
mute = Setting.objects(
guild=before.guild.id, setting="mute"
).first()
verified = Setting.objects(
guild=before.guild.id, setting="verified"
).first()
mute_role = None
verified_role = None
if mute:
mute_role = before.guild.get_role(mute.value)
if verified:
verified_role = before.guild.get_role(verified.value)
if (
mute
and mute_role in after.roles
and mute_role not in before.roles
):
embed = await self.process_mute(before, after)
elif mute and before.guild.get_role(mute.value) in before.roles:
elif (
mute
and mute_role in before.roles
and mute_role not in after.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 verified_role not in before.roles
and verified_role in after.roles
):
embed = await self.process_verify(before, after)
elif before.nick != after.nick:

10
jarvis/cogs/modlog/message.py
View file

@ -1,7 +1,7 @@
import discord
from discord.ext import commands
from jarvis.db.types import Setting
from jarvis.db.models import Setting
from jarvis.utils import build_embed
from jarvis.utils.field import Field
@ -15,7 +15,9 @@ class ModlogMessageCog(commands.Cog):
self, before: discord.Message, after: discord.Message
):
if not before.author.bot:
modlog = Setting.get(guild=after.guild.id, setting="modlog")
modlog = Setting.objects(
guild=after.guild.id, setting="modlog"
).first()
if modlog:
if before.content == after.content or before.content is None:
return
@ -53,7 +55,9 @@ class ModlogMessageCog(commands.Cog):
@commands.Cog.listener()
async def on_message_delete(self, message: discord.Message):
modlog = Setting.get(guild=message.guild.id, setting="modlog")
modlog = Setting.objects(
guild=message.guild.id, setting="modlog"
).first()
if modlog:
fields = [
Field("Original Message", message.content or "N/A", False)

18
jarvis/cogs/owner.py
View file

@ -10,7 +10,7 @@ from discord.ext import commands
import jarvis
from jarvis.config import reload_config
from jarvis.db.types import Config
from jarvis.db.models import Config
from jarvis.utils import update
from jarvis.utils.permissions import user_is_bot_admin
@ -24,7 +24,7 @@ class OwnerCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.admins = Config.get(key="admins")
self.admins = Config.objects(key="admins").first()
@commands.command(name="load", hidden=True)
@user_is_bot_admin()
@ -32,7 +32,7 @@ class OwnerCog(commands.Cog):
info = await self.bot.application_info()
if (
ctx.message.author == info.owner
or ctx.message.author.id in self.admins
or ctx.message.author.id in self.admins.value
):
try:
if "jarvis.cogs." not in cog:
@ -58,7 +58,7 @@ class OwnerCog(commands.Cog):
info = await self.bot.application_info()
if (
ctx.message.author == info.owner
or ctx.message.author.id in self.admins
or ctx.message.author.id in self.admins.value
):
try:
if "jarvis.cogs." not in cog:
@ -84,7 +84,7 @@ class OwnerCog(commands.Cog):
info = await self.bot.application_info()
if (
ctx.message.author == info.owner
or ctx.message.author.id in self.admins
or ctx.message.author.id in self.admins.value
):
try:
if "jarvis.cogs." not in cog:
@ -118,7 +118,7 @@ class OwnerCog(commands.Cog):
info = await self.bot.application_info()
if (
ctx.message.author == info.owner
or ctx.message.author.id in self.admins
or ctx.message.author.id in self.admins.value
):
await ctx.send("Restarting core systems...")
if isinstance(ctx.channel, discord.channel.DMChannel):
@ -141,7 +141,7 @@ class OwnerCog(commands.Cog):
info = await self.bot.application_info()
if (
ctx.message.author == info.owner
or ctx.message.author.id in self.admins
or ctx.message.author.id in self.admins.value
):
await ctx.send("Updating core systems...")
status = update()
@ -189,7 +189,7 @@ class OwnerCog(commands.Cog):
await ctx.send(f"{user.mention} is already an admin.")
return
self.admins.value.append(user.id)
self.admins.update()
self.admins.save()
reload_config()
await ctx.send(
f"{user.mention} is now an admin. Use this power carefully."
@ -202,7 +202,7 @@ class OwnerCog(commands.Cog):
await ctx.send(f"{user.mention} is not an admin.")
return
self.admins.value.remove(user.id)
self.admins.update()
self.admins.save()
reload_config()
await ctx.send(f"{user.mention} is no longer an admin.")

20
jarvis/cogs/remindme.py
View file

@ -16,7 +16,7 @@ from discord_slash.utils.manage_components import (
wait_for_component,
)
from jarvis.db.types import Reminder
from jarvis.db.models import Reminder
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
@ -124,8 +124,8 @@ class RemindmeCog(CacheCog):
)
return
reminders = Reminder.get_active(user=ctx.author.id)
if len(reminders) >= 5:
reminders = Reminder.objects(user=ctx.author.id, active=True).count()
if reminders >= 5:
await ctx.send(
"You already have 5 (or more) active reminders. "
+ "Please either remove an old one, or wait for one to pass",
@ -147,7 +147,7 @@ class RemindmeCog(CacheCog):
message=message,
remind_at=remind_at,
active=True,
).insert()
).save()
embed = build_embed(
title="Reminder Set",
@ -210,7 +210,7 @@ class RemindmeCog(CacheCog):
hidden=True,
)
return
reminders = Reminder.get_active(user=ctx.author.id)
reminders = Reminder.objects(user=ctx.author.id, active=True)
if not reminders:
await ctx.send("You have no reminders set.", hidden=True)
return
@ -225,7 +225,7 @@ class RemindmeCog(CacheCog):
description="Delete a reminder",
)
async def _delete(self, ctx: SlashContext):
reminders = Reminder.get_active(user=ctx.author.id)
reminders = Reminder.objects(user=ctx.author.id, active=True)
if not reminders:
await ctx.send("You have no reminders set", hidden=True)
return
@ -263,8 +263,8 @@ class RemindmeCog(CacheCog):
timeout=60 * 5,
)
for to_delete in context.selected_options:
_ = Reminder.get(
user=ctx.author.id, _id=ObjectId(to_delete)
_ = Reminder.objects(
user=ctx.author.id, id=ObjectId(to_delete)
).delete()
for row in components:
@ -307,8 +307,8 @@ class RemindmeCog(CacheCog):
@loop(seconds=15)
async def _remind(self):
reminders = Reminder.get_active(
remind_at={"$lt": datetime.utcnow() + timedelta(seconds=30)}
reminders = Reminder.objects(
remind_at__lte=datetime.utcnow() + timedelta(seconds=30)
)
for reminder in reminders:
if reminder.remind_at <= datetime.utcnow():

26
jarvis/cogs/rolegiver.py
View file

@ -3,7 +3,7 @@ from discord.ext import commands
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Setting
from jarvis.db.models import Setting
from jarvis.utils import build_embed
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions
@ -28,7 +28,9 @@ class RolegiverCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _rolegiver_add(self, ctx: SlashContext, role: Role):
setting = Setting.get(guild=ctx.guild.id, setting="rolegiver")
setting = Setting.objects(
guild=ctx.guild.id, setting="rolegiver"
).first()
if setting and role.id in setting.value:
await ctx.send("Role already in rolegiver", hidden=True)
return
@ -51,7 +53,7 @@ class RolegiverCog(commands.Cog):
Field(name="Existing Role(s)", value=value),
]
setting.value.append(role.id)
setting.update()
setting.save()
embed = build_embed(
title="Rolegiver Updated",
@ -87,7 +89,9 @@ class RolegiverCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _rolegiver_remove(self, ctx: SlashContext, role: Role):
setting = Setting.get(guild=ctx.guild.id, setting="rolegiver")
setting = Setting.objects(
guild=ctx.guild.id, setting="rolegiver"
).first()
if not setting or (setting and not setting.value):
await ctx.send("Rolegiver has no roles", hidden=True)
return
@ -96,7 +100,7 @@ class RolegiverCog(commands.Cog):
return
setting.value.remove(role.id)
setting.update()
setting.save()
roles = []
for role_id in setting.value:
@ -137,7 +141,9 @@ class RolegiverCog(commands.Cog):
description="List roles rolegiver",
)
async def _rolegiver_list(self, ctx: SlashContext):
setting = Setting.get(guild=ctx.guild.id, setting="rolegiver")
setting = Setting.objects(
guild=ctx.guild.id, setting="rolegiver"
).first()
if not setting or (setting and not setting.value):
await ctx.send("Rolegiver has no roles", hidden=True)
return
@ -186,7 +192,9 @@ class RolegiverCog(commands.Cog):
)
@commands.cooldown(1, 10, commands.BucketType.user)
async def _role_get(self, ctx: SlashContext, role: Role):
setting = Setting.get(guild=ctx.guild.id, setting="rolegiver")
setting = Setting.objects(
guild=ctx.guild.id, setting="rolegiver"
).first()
if not setting or (setting and not setting.value):
await ctx.send("Rolegiver has no roles", hidden=True)
return
@ -244,7 +252,9 @@ class RolegiverCog(commands.Cog):
)
@commands.cooldown(1, 10, commands.BucketType.user)
async def _role_forfeit(self, ctx: SlashContext, role: Role):
setting = Setting.get(guild=ctx.guild.id, setting="rolegiver")
setting = Setting.objects(
guild=ctx.guild.id, setting="rolegiver"
).first()
if not setting or (setting and not setting.value):
await ctx.send("Rolegiver has no roles", hidden=True)
return

16
jarvis/cogs/settings.py
View file

@ -4,7 +4,7 @@ from discord.utils import find
from discord_slash import SlashContext, cog_ext
from discord_slash.utils.manage_commands import create_option
from jarvis.db.types import Setting
from jarvis.db.models import Setting
from jarvis.utils import build_embed
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions
@ -15,13 +15,15 @@ class SettingsCog(commands.Cog):
self.bot = bot
def update_settings(self, setting, value, guild):
setting = Setting(setting=setting, value=value, guild=guild)
updated = setting.update()
setting = Setting.objects(setting=setting, guild=guild).first()
if not setting:
setting = Setting(setting=setting, guild=guild, value=value)
updated = setting.save()
return updated is not None
def delete_settings(self, setting, guild):
return Setting.delete_many(setting=setting, guild=guild)
return Setting.objects(setting=setting, guild=guild).delete()
@cog_ext.cog_subcommand(
base="settings",
@ -231,7 +233,7 @@ class SettingsCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _view(self, ctx: SlashContext):
settings = Setting.get_many(guild=ctx.guild.id)
settings = Setting.objects(guild=ctx.guild.id)
fields = []
for setting in settings:
@ -261,8 +263,8 @@ class SettingsCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _clear(self, ctx: SlashContext):
deleted = Setting.delete_many(guild=ctx.guild.id)
await ctx.send(f"Guild settings cleared: `{deleted}`")
deleted = Setting.objects(guild=ctx.guild.id).delete()
await ctx.send(f"Guild settings cleared: `{deleted is not None}`")
def setup(bot):

36
jarvis/cogs/starboard.py
View file

@ -10,7 +10,7 @@ from discord_slash.utils.manage_components import (
wait_for_component,
)
from jarvis.db.types import Star, Starboard
from jarvis.db.models import Star, Starboard
from jarvis.utils import build_embed
from jarvis.utils.permissions import admin_or_permissions
@ -34,7 +34,7 @@ class StarboardCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _list(self, ctx):
starboards = Starboard.get_many(guild=ctx.guild.id)
starboards = Starboard.objects(guild=ctx.guild.id)
if starboards != []:
message = "Available Starboards:\n"
for s in starboards:
@ -68,15 +68,17 @@ class StarboardCog(commands.Cog):
await ctx.send("Channel must be a TextChannel", hidden=True)
return
exists = Starboard.get(channel=channel.id, guild=ctx.guild.id)
exists = Starboard.objects(
channel=channel.id, guild=ctx.guild.id
).first()
if exists:
await ctx.send(
f"Starboard already exists at {channel.mention}.", hidden=True
)
return
count = Starboard.get_many(guild=ctx.guild.id)
if count and len(count) >= 25:
count = Starboard.objects(guild=ctx.guild.id).count()
if count >= 25:
await ctx.send("25 starboard limit reached", hidden=True)
return
@ -84,7 +86,7 @@ class StarboardCog(commands.Cog):
guild=ctx.guild.id,
channel=channel.id,
admin=ctx.author.id,
).insert()
).save()
await ctx.send(
f"Starboard created. Check it out at {channel.mention}."
)
@ -104,11 +106,11 @@ class StarboardCog(commands.Cog):
)
@admin_or_permissions(manage_guild=True)
async def _delete(self, ctx, channel: TextChannel):
deleted = Starboard.get(
deleted = Starboard.objects(
channel=channel.id, guild=ctx.guild.id
).delete()
if deleted:
_ = Star.delete_many(starboard=channel.id)
_ = Star.objects(starboard=channel.id).delete()
await ctx.send(
f"Starboard deleted from {channel.mention}.", hidden=True
)
@ -146,7 +148,7 @@ class StarboardCog(commands.Cog):
):
if not channel:
channel = ctx.channel
starboards = Starboard.get_many(guild=ctx.guild.id)
starboards = Starboard.objects(guild=ctx.guild.id)
if not starboards:
await ctx.send("No starboards exist.", hidden=True)
return
@ -189,12 +191,12 @@ class StarboardCog(commands.Cog):
message = await channel.fetch_message(int(message))
exists = Star.get(
exists = Star.objects(
message=message.id,
channel=message.channel.id,
guild=message.guild.id,
starboard=starboard.id,
)
).first()
if exists:
await ctx.send(
@ -247,7 +249,7 @@ class StarboardCog(commands.Cog):
admin=ctx.author.id,
star=star.id,
active=True,
).insert()
).save()
components[0]["components"][0]["disabled"] = True
@ -286,7 +288,9 @@ class StarboardCog(commands.Cog):
if not isinstance(starboard, TextChannel):
await ctx.send("Channel must be a TextChannel", hidden=True)
return
exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id)
exists = Starboard.objects(
channel=starboard.id, guild=ctx.guild.id
).first()
if not exists:
await ctx.send(
f"Starboard does not exist in {starboard.mention}. "
@ -295,12 +299,12 @@ class StarboardCog(commands.Cog):
)
return
star = Star.get(
star = Star.objects(
starboard=starboard.id,
index=id,
guild=ctx.guild.id,
active=True,
)
).first()
if not star:
await ctx.send(f"No star exists with id {id}", hidden=True)
return
@ -310,7 +314,7 @@ class StarboardCog(commands.Cog):
await message.delete()
star.active = False
star.update()
star.save()
await ctx.send(f"Star {id} deleted")

12
jarvis/cogs/verify.py
View file

@ -5,7 +5,7 @@ from discord_slash import ComponentContext, SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from discord_slash.utils import manage_components
from jarvis.db.types import Setting
from jarvis.db.models import Setting
def create_layout():
@ -37,7 +37,7 @@ class VerifyCog(commands.Cog):
@commands.cooldown(1, 15, commands.BucketType.user)
async def _verify(self, ctx: SlashContext):
await ctx.defer()
role = Setting.get(guild=ctx.guild.id, setting="verified")
role = Setting.objects(guild=ctx.guild.id, setting="verified").first()
if not role:
await ctx.send(
"This guild has not enabled verification", delete_after=5
@ -68,10 +68,14 @@ class VerifyCog(commands.Cog):
for c in components:
for c2 in c["components"]:
c2["disabled"] = True
setting = Setting.get(guild=ctx.guild.id, setting="verified")
setting = Setting.objects(
guild=ctx.guild.id, setting="verified"
).first()
role = ctx.guild.get_role(setting.value)
await ctx.author.add_roles(role, reason="Verification passed")
setting = Setting.get(guild=ctx.guild.id, setting="unverified")
setting = Setting.objects(
guild=ctx.guild.id, setting="unverified"
).first()
if setting:
role = ctx.guild.get_role(setting.value)
await ctx.author.remove_roles(

11
jarvis/config.py
View file

@ -1,6 +1,6 @@
from yaml import load
from jarvis.db import DBManager
from jarvis.db.models import Config as DBConfig
try:
from yaml import CLoader as Loader
@ -38,10 +38,11 @@ class Config(object):
self.events = events
self.max_messages = max_messages
self.gitlab_token = gitlab_token
db = DBManager(config=mongo["connect"]).mongo[mongo["database"]].config
db_config = db.find()
for item in db_config:
setattr(self, item["key"], item["value"])
def get_db_config(self):
config = DBConfig.objects()
for item in config:
setattr(self, item.key, item.value)
@classmethod
def from_yaml(cls, y):

14
jarvis/db/__init__.py
View file

@ -1,14 +0,0 @@
from pymongo import MongoClient
class DBManager(object):
def __new__(cls, *args, **kwargs):
it = cls.__dict__.get("it")
if it is not None:
return it
cls.__it__ = it = object.__new__(cls)
it.init(*args, **kwargs)
return it
def init(self, config: dict):
self.mongo = MongoClient(**config)

197
jarvis/db/models.py Normal file
View file

@ -0,0 +1,197 @@
from datetime import datetime
from mongoengine import Document
from mongoengine.fields import (
BooleanField,
DateTimeField,
DictField,
DynamicField,
IntField,
ListField,
LongField,
StringField,
)
class SnowflakeField(LongField):
pass
class Autopurge(Document):
guild = SnowflakeField(required=True)
channel = SnowflakeField(required=True)
delay = IntField(min_value=1, max_value=300, default=30)
admin = SnowflakeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Autoreact(Document):
guild = SnowflakeField(required=True)
channel = SnowflakeField(required=True)
reactions = ListField(field=StringField())
admin = SnowflakeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Ban(Document):
active = BooleanField(default=True)
admin = SnowflakeField(required=True)
user = SnowflakeField(required=True)
username = StringField(required=True)
discrim = IntField(min_value=1, max_value=9999, required=True)
duration = IntField(min_value=1, max_value=744, required=False)
guild = SnowflakeField(required=True)
type = StringField(default="perm", max_length=4, required=True)
reason = StringField(max_length=100, required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Config(Document):
key = StringField(required=True)
value = DynamicField(required=True)
meta = {"db_alias": "main"}
class Guess(Document):
correct = BooleanField(default=False)
guess = StringField(max_length=800, required=True)
user = SnowflakeField(required=True)
meta = {"db_alias": "ctc2"}
class Joke(Document):
rid = StringField()
body = StringField()
title = StringField()
created_utc = DateTimeField()
over_18 = BooleanField()
score = IntField()
meta = {"db_alias": "main"}
class Kick(Document):
admin = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
reason = StringField(max_length=100, required=True)
user = SnowflakeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Lock(Document):
active = BooleanField(default=True)
admin = SnowflakeField(required=True)
channel = SnowflakeField(required=True)
duration = IntField(min_value=1, max_value=300, default=10)
guild = SnowflakeField(required=True)
reason = StringField(max_length=100, required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Mute(Document):
active = BooleanField(default=True)
user = SnowflakeField(required=True)
admin = SnowflakeField(required=True)
duration = IntField(min_value=1, max_value=300, default=10)
guild = SnowflakeField(required=True)
reason = StringField(max_length=100, required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Purge(Document):
admin = SnowflakeField(required=True)
channel = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
count = IntField(min_value=1, default=10)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Reminder(Document):
active = BooleanField(default=True)
user = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
message = StringField(max_length=100, required=True)
remind_at = DateTimeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Roleping(Document):
active = BooleanField(default=True)
role = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
admin = SnowflakeField(required=True)
bypass = DictField()
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Setting(Document):
guild = SnowflakeField(required=True)
setting = StringField(required=True)
value = DynamicField()
meta = {"db_alias": "main"}
class Star(Document):
index = IntField(required=True)
message = SnowflakeField(required=True)
starboard = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
admin = SnowflakeField(required=True)
star = SnowflakeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Starboard(Document):
channel = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
admin = SnowflakeField(required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Unban(Document):
user = SnowflakeField(required=True)
username = StringField(required=True)
discrim = IntField(min_value=1, max_value=9999, required=True)
guild = SnowflakeField(required=True)
admin = SnowflakeField(required=True)
reason = StringField(max_length=100, required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}
class Warning(Document):
active = BooleanField(default=True)
admin = SnowflakeField(required=True)
user = SnowflakeField(required=True)
guild = SnowflakeField(required=True)
duration = IntField(min_value=1, max_value=120, default=24)
reason = StringField(max_length=100, required=True)
created_at = DateTimeField(default=datetime.utcnow)
meta = {"db_alias": "main"}

594
jarvis/db/types.py
View file

@ -1,594 +0,0 @@
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 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",
"Autopurge": "autopurges",
"Ban": "bans",
"Config": "config",
"Lock": "locks",
"Kick": "kicks",
"Mute": "mutes",
"Purge": "purges",
"Reminder": "reminders",
"Roleping": "rolepings",
"Setting": "settings",
"Starboard": "starboard",
"Star": "stars",
"Unban": "unbans",
"Warning": "warns",
}
m_config = get_config().mongo
db_instance = DBManager(m_config["connect"]).mongo[m_config["database"]]
#################
# 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()
if self._id:
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
else:
self.insert()
@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 Reminder(MongoObject, ActiveObject):
"""
User Reminder object
:param _id: MongoDB ID
:param user: User who wanted reminded
:param channel: ID of origin channel
:param guild: ID of origin guild
:param remind_at: Datetime delay to send message
:param active: If reminder is active
:param message: Reminder message
:param created_at: Time the reminder was created
"""
user: int
channel: int
guild: int
remind_at: datetime
message: str
created_at: datetime = field(default_factory=datetime.utcnow)
@dataclass
class Roleping(MongoObject, ActiveObject):
"""
Guild Roleping object
:param _id: MongoDB ID
:param role: Blocked role
:param guild: ID of origin guild
:param admin: Admin who added roleping
:param bypass: Roles/users who may bypass this roleping
:param active: If the roleping is disabled
:param created_at: Time the roleping was created
"""
role: int
guild: int
admin: int
bypass: dict
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, ActiveObject):
"""
Starboard Star object
:param _id: MongoDB ID
:param index: Starboard star index
:param message: Star Message ID
:param channel: Star Channel ID
:param starboard: Starboard Channel ID
:param guild: Starboard Guild ID
:param admin: ID of admin who created star
:param star: ID of the star message
:param created_at: Time created
"""
index: int
message: int
channel: int
starboard: int
guild: int
admin: int
star: 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)

5
jarvis/events/guild.py
View file

@ -2,7 +2,7 @@ import asyncio
from discord.utils import find
from jarvis.db.types import Setting
from jarvis.db.models import Setting
class GuildEventHandler(object):
@ -26,7 +26,6 @@ class GuildEventHandler(object):
)
# Set some default settings
setting = Setting(guild=guild.id, setting="massmention", value=5)
setting.insert()
_ = Setting(guild=guild.id, setting="massmention", value=5).save()
await general.send("Systems are now fully operational")

12
jarvis/events/member.py
View file

@ -1,6 +1,6 @@
from discord import Member
from jarvis.db.types import Mute, Setting
from jarvis.db.models import Mute, Setting
class MemberEventHandler(object):
@ -10,14 +10,16 @@ class MemberEventHandler(object):
async def on_member_join(self, user: Member):
guild = user.guild
mutes = Mute.get_active(guild=guild.id)
if mutes and len(mutes) >= 1:
mute_role = Setting.get(guild=guild.id, setting="mute")
mute = Mute.objects(guild=guild.id, user=user.id, active=True).first()
if mute:
mute_role = Setting.objects(guild=guild.id, setting="mute").first()
role = guild.get_role(mute_role.value)
await user.add_roles(
role, reason="User is muted still muted from prior mute"
)
unverified = Setting.get(guild=guild.id, setting="unverified")
unverified = Setting.objects(
guild=guild.id, setting="unverified"
).first()
if unverified:
role = guild.get_role(unverified.value)
if role not in user.roles:

31
jarvis/events/message.py
View file

@ -4,7 +4,7 @@ from discord import DMChannel, Message
from discord.utils import find
from jarvis.config import get_config
from jarvis.db.types import Autopurge, Autoreact, Roleping, Setting, Warning
from jarvis.db.models import Autopurge, Autoreact, Roleping, Setting, Warning
from jarvis.utils import build_embed
from jarvis.utils.field import Field
@ -21,17 +21,17 @@ class MessageEventHandler(object):
self.bot.add_listener(self.on_message_edit)
async def autopurge(self, message: Message):
autopurge = Autopurge.get(
autopurge = Autopurge.objects(
guild=message.guild.id, channel=message.channel.id
)
).first()
if autopurge:
await message.delete(delay=autopurge.delay)
async def autoreact(self, message: Message):
autoreact = Autoreact.get(
autoreact = Autoreact.objects(
guild=message.guild.id,
channel=message.channel.id,
)
).first()
if autoreact:
for reaction in autoreact.reactions:
await message.add_reaction(reaction)
@ -57,15 +57,14 @@ class MessageEventHandler(object):
]
if match.group(1) not in allowed:
await message.delete()
warning = 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()
).save()
fields = [
Field(
"Reason",
@ -92,10 +91,10 @@ class MessageEventHandler(object):
await message.channel.send(embed=embed)
async def massmention(self, message: Message):
massmention = Setting.get(
massmention = Setting.objects(
guild=message.guild.id,
setting="massmention",
)
).first()
if (
massmention
and massmention.value > 0
@ -103,15 +102,14 @@ class MessageEventHandler(object):
- (1 if message.author in message.mentions else 0)
> massmention.value
):
warning = Warning(
_ = Warning(
active=True,
admin=get_config().client_id,
duration=24,
guild=message.guild.id,
reason="Mass Mention",
user=message.author.id,
)
warning.insert()
).save()
fields = [Field("Reason", "Mass Mention", False)]
embed = build_embed(
title="Warning",
@ -131,7 +129,7 @@ class MessageEventHandler(object):
await message.channel.send(embed=embed)
async def roleping(self, message: Message):
rolepings = Roleping.get_active(guild=message.guild.id)
rolepings = Roleping.objects(guild=message.guild.id, active=True)
if not rolepings:
return
@ -180,15 +178,14 @@ class MessageEventHandler(object):
and not user_is_admin
and not user_has_bypass
):
warning = 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()
).save()
fields = [
Field(
"Reason",

33
jarvis/tasks/unban.py
View file

@ -1,16 +1,18 @@
from datetime import datetime, timedelta
import pymongo
from discord.ext.tasks import loop
import jarvis
from jarvis.db.types import Ban
from jarvis.config import get_config
from jarvis.db.models import Ban, Unban
jarvis_id = get_config().client_id
@loop(minutes=10)
async def unban():
bans = Ban.get_active(type="temp")
updates = []
bans = Ban.objects(type="temp", active=True)
unbans = []
for ban in bans:
if ban.created_at + timedelta(
hours=ban.duration
@ -19,16 +21,17 @@ async def unban():
user = await jarvis.jarvis.fetch_user(ban.user)
if user:
guild.unban(user)
updates.append(
pymongo.UpdateOne(
{
"user": user.id,
"guild": guild.id,
"created_at": ban.created_at,
"type": "temp",
},
{"$set": {"active": False}},
ban.active = False
ban.save()
unbans.append(
Unban(
user=user.id,
guild=guild.id,
username=user.name,
discrim=user.discriminator,
admin=jarvis_id,
reason="Ban expired",
)
)
if updates:
jarvis.jarvis_db.bans.bulk_write(updates)
if unbans:
Ban.objects().insert(unbans)

20
jarvis/tasks/unlock.py
View file

@ -1,16 +1,14 @@
from datetime import datetime, timedelta
import pymongo
from discord.ext.tasks import loop
import jarvis
from jarvis.db.types import Lock
from jarvis.db.models import Lock
@loop(minutes=1)
async def unlock():
locks = Lock.get_active()
updates = []
locks = Lock.objects(active=True)
for lock in locks:
if (
lock.created_at + timedelta(minutes=lock.duration)
@ -26,15 +24,5 @@ async def unlock():
await channel.set_permissions(
role, overwrite=overrides, reason="Lock expired"
)
updates.append(
pymongo.UpdateOne(
{
"channel": channel.id,
"guild": guild.id,
"created_at": lock.created_at,
},
{"$set": {"active": False}},
)
)
if updates:
jarvis.jarvis_db.locks.bulk_write(updates)
lock.active = False
lock.save()

22
jarvis/tasks/unmute.py
View file

@ -1,17 +1,15 @@
from datetime import datetime, timedelta
import pymongo
from discord.ext.tasks import loop
import jarvis
from jarvis.db.types import Mute, Setting
from jarvis.db.models import Mute, Setting
@loop(minutes=1)
async def unmute():
mutes = Mute.get_active(duration={"$gt": 0})
mute_roles = Setting.get_many(setting="mute")
updates = []
mutes = Mute.objects(duration__gt=0, active=True)
mute_roles = Setting.objects(setting="mute")
for mute in mutes:
if (
mute.created_at + timedelta(minutes=mute.duration)
@ -28,15 +26,5 @@ async def unmute():
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,
"created_at": mute.created_at,
},
{"$set": {"active": False}},
)
)
if updates:
jarvis.jarvis_db.mutes.bulk_write(updates)
mute.active = False
mute.save

15
jarvis/tasks/unwarn.py
View file

@ -1,25 +1,18 @@
from datetime import datetime, timedelta
import pymongo
from discord.ext.tasks import loop
import jarvis
from jarvis.db.types import Warning
from jarvis.db.models import Warning
@loop(hours=1)
async def unwarn():
warns = Warning.get_active()
updates = []
warns = Warning.objects(active=True)
for warn in warns:
if (
warn.created_at + timedelta(hours=warn.duration)
< datetime.utcnow()
):
updates.append(
pymongo.UpdateOne(
{"_id": warn._id}, {"$set": {"active": False}}
)
)
if updates:
jarvis.jarvis_db.warns.bulk_write(updates)
warn.active = False
warn.save()

1
requirements.txt
View file

@ -9,3 +9,4 @@ ButtonPaginator>=0.0.3
Pillow>=8.2.0, <9
python-gitlab>=2.9.0, <3
ulid-py>=1.1.0, <2
mongoengine>=0.23, <1