Add mute, background unmute and unban tasks

jarvis/__init__.py

@@ -1,11 +1,15 @@
 from pathlib import Path
+import pymongo
+from datetime import datetime, timedelta
 from discord import Intents
 from discord.ext import commands
-from discord.utils import find
+from discord.ext.tasks import loop
+from discord.utils import find, get
 from discord_slash import SlashCommand
 from psutil import Process
 import asyncio
 from jarvis.config import get_config
+from jarvis.utils.db import DBManager
 from jarvis import utils, logo
 
 if asyncio.get_event_loop().is_closed():
@@ -61,6 +65,60 @@ async def on_guild_join(guild):
         await general.send("Systems are now fully operational")
 
 
+@loop(minutes=1)
+async def unmute():
+    db = DBManager(get_config().mongo).mongo
+    mutes = list(db.jarvis.mutes.find({"active": True}))
+    mute_roles = list(
+        db.jarvis.settings.find({"setting": "mute"}).sort[
+            ("guild", pymongo.ASCENDING)
+        ]
+    )
+    updates = []
+    for mute in mutes:
+        if mute["time"] + timedelta(minutes=mute["length"]) < datetime.now():
+            mute_role = [
+                x["value"] for x in mute_roles if x["guild"] == mute["guild"]
+            ]
+            guild = await jarvis.fetch_guild(mute["guild"])
+            role = guild.get_role(mute_role)
+            user = guild.get_member(mute["user"])
+            if user:
+                user.remove_roles(role, reason="No longer muted")
+            updates.append(
+                pymongo.UpdateOne(
+                    {"user": user.id, "guild": guild.id, "time": mute["time"]},
+                    {"$set": {"active": False}},
+                )
+            )
+    db.jarvis.mutes.update_many(updates)
+
+
+@loop(minutes=30)
+async def unban():
+    db = DBManager(get_config().mongo).mongo
+    bans = list(db.jarvis.bans.find({"active": True, "type": "temp"}))
+    updates = []
+    for ban in bans:
+        if ban["time"] + timedelta(days=ban["length"]) < datetime.now():
+            guild = await jarvis.fetch_guild(ban["guild"])
+            user = jarvis.fetch_user(ban["user"])
+            if user:
+                guild.unban(user)
+            updates.append(
+                pymongo.UpdateOne(
+                    {
+                        "user": user.id,
+                        "guild": guild.id,
+                        "time": ban["time"],
+                        "type": "temp",
+                    },
+                    {"$set": {"active": False}},
+                )
+            )
+    db.jarvis.bans.update_many(updates)
+
+
 def run(ctx=None):
     global restart_ctx
     if ctx:
@@ -74,6 +132,8 @@ def run(ctx=None):
             config.client_id
         )
     )
+    unmute.start()
+    unban.start()
     jarvis.run(config.token, bot=True, reconnect=True)
     for cog in jarvis.cogs:
         session = getattr(cog, "_session", None)

jarvis/cogs/admin.py

@@ -1,7 +1,8 @@
 import jarvis
 from datetime import datetime
-from discord import User
+from discord import User, Member
 from discord.ext import commands
+from discord.utils import get
 from discord_slash import cog_ext
 from discord_slash.utils.manage_commands import create_option, create_choice
 from jarvis.utils.db import DBManager
@@ -208,6 +209,57 @@ class AdminCog(commands.Cog):
         await ctx.defer()
         await self._purge(ctx, amount)
 
+    @cog_ext.cog_slash(
+        name="mute",
+        description="Mute a user",
+        guild_ids=[578757004059738142],
+        options=[
+            create_option(
+                name="user",
+                description="User to mute",
+                option_type=6,
+                required=True,
+            ),
+            create_option(
+                name="reason",
+                description="Reason for mute",
+                option_type=3,
+                required=True,
+            ),
+            create_option(
+                name="length",
+                description="Mute length",
+                option_type=4,
+                required=False,
+            ),
+        ],
+    )
+    async def _mute(self, ctx, user: Member, reason: str, length: int = 30):
+        mute_setting = self.db.jarvis.settings.find_one(
+            {"guild": ctx.guild.id, "setting": "mute"}
+        )
+        if not mute_setting:
+            await ctx.send(
+                "Please configure a mute role with /settings mute <role> first"
+            )
+            return
+        role = get(ctx.guild.roles, id=mute_setting["value"])
+        user.add_roles(role, reason=reason)
+        if length < 0:
+            length = -1
+        self.db.jarvis.mutes.insert_one(
+            {
+                "user": user.id,
+                "reason": reason,
+                "admin": ctx.author.id,
+                "time": datetime.now(),
+                "guild": ctx.guild.id,
+                "length": length,
+                "active": True,
+            }
+        )
+        await ctx.send(f"{user.mention} has been muted.\nReason: {reason}")
+
 
 def setup(bot):
     bot.add_cog(AdminCog(bot))

jarvis/cogs/settings.py

@@ -0,0 +1,45 @@
+import discord
+import jarvis
+from discord import Role
+from discord.ext import commands
+from discord_slash import cog_ext
+from discord_slash.utils.manage_commands import create_option
+from jarvis.config import get_config, reload_config
+from jarvis.utils import update, user_is_bot_admin, db
+
+
+class SettingsCog(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+        config = get_config()
+        self.db = db.DBManager(config.mongo).mongo
+        self.cache = {}
+
+    def update_settings(self, setting, value, guild):
+        settings = {"setting": setting, "value": value, "guild": guild}
+        updated = self.db.jarvis.settings.update_one(
+            {"settings": setting, "guild": guild},
+            {"$set": settings},
+            upsert=True,
+        )
+
+        return updated is not None
+
+    @cog_ext.cog_subcommand(
+        base="settings",
+        name="mute",
+        description="Set mute role",
+        options=[
+            create_option(
+                name="role",
+                description="Mute role",
+                option_type=8,
+                required=True,
+            )
+        ],
+    )
+    @commands.has_permissions(mute_members=True)
+    async def _mute(self, ctx, role: Role):
+        await ctx.defer()
+        self.update_settings("mute", role.id, ctx.guild.id)
+        await ctx.send(f"Settings applied. New mute role is {role.name}")

schema.yaml

@@ -28,7 +28,7 @@ jarvis:
     admin: User ID, admin who muted user
     time: datetime
     guild: Guild ID
-    type: String ("voice" | "text" | "all"), mute type, default "all"
+    #type: String ("voice" | "text" | "all"), mute type, default "all"
     length: int, mute length in minutes, default 30, -1 for permanent
     active: boolean, if mute is active