Add warning system, closes #12. Add /unlock command, /lockdown {start,end}

jarvis/__init__.py

@@ -212,6 +212,22 @@ async def unlock():
         db.jarvis.locks.bulk_write(updates)
 
 
+@loop(hours=1)
+async def unwarn():
+    db = DBManager(get_config().mongo).mongo
+    warns = list(db.jarvis.warns.find({"active": True}))
+    updates = []
+    for warn in warns:
+        if warn["time"] + timedelta(hours=warn["duration"]) < datetime.now():
+            updates.append(
+                pymongo.UpdateOne(
+                    {{"_id": warn["_id"]}, {"$set": {"active": False}}}
+                )
+            )
+    if updates:
+        db.jarvis.warns.bulk_write(updates)
+
+
 def run(ctx=None):
     global restart_ctx
     if ctx:
@@ -228,6 +244,7 @@ def run(ctx=None):
     unmute.start()
     unban.start()
     unlock.start()
+    unwarn.start()
     jarvis.max_messages = config.max_messages
     jarvis.run(config.token, bot=True, reconnect=True)
     for cog in jarvis.cogs:

jarvis/cogs/admin.py

@@ -9,7 +9,9 @@ from discord_slash import SlashContext, cog_ext
 from discord_slash.utils.manage_commands import create_choice, create_option
 
 import jarvis
+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
 
 
@@ -564,6 +566,16 @@ class AdminCog(commands.Cog):
         overrides.send_messages = allow_send
         await channel.set_permissions(role, overwrite=overrides, reason=reason)
 
+    async def _unlock_channel(
+        self,
+        channel: TextChannel,
+        role: Role,
+        admin: User,
+    ):
+        overrides = channel.overwrites_for(role)
+        overrides.send_messages = None
+        await channel.set_permissions(role, overwrite=overrides)
+
     @cog_ext.cog_slash(
         name="lock",
         description="Locks a channel",
@@ -589,6 +601,7 @@ class AdminCog(commands.Cog):
             ),
         ],
     )
+    @commands.has_permissions(administrator=True)
     async def _lock(
         self,
         ctx: SlashContext,
@@ -618,7 +631,50 @@ class AdminCog(commands.Cog):
         await ctx.send(f"{channel.mention} locked for {duration} minute(s)")
 
     @cog_ext.cog_slash(
-        name="lockdown",
+        name="unlock",
+        description="Unlocks a channel",
+        guild_ids=[418094694325813248, 578757004059738142, 862402786116763668],
+        options=[
+            create_option(
+                name="channel",
+                description="Channel to lock",
+                option_type=7,
+                required=False,
+            ),
+        ],
+    )
+    @commands.has_permissions(administrator=True)
+    async def _unlock(
+        self,
+        ctx: SlashContext,
+        channel: TextChannel = None,
+    ):
+        await ctx.defer()
+        if not channel:
+            channel = ctx.channel
+        lock = self.db.jarvis.locks.find_one(
+            {"guild": ctx.guild.id, "channel": channel.id, "active": True}
+        )
+        if not lock:
+            await ctx.send(f"{channel.mention} not locked.")
+            return
+        for role in ctx.guild.roles:
+            try:
+                await self._unlock_channel(channel, role, ctx.author)
+            except Exception:
+                continue  # Just continue on error
+        self.db.jarvis.locks.update_one(
+            {
+                "channel": channel.id,
+                "guild": ctx.guild.id,
+            },
+            {"$set": {"active": False}},
+        )
+        await ctx.send(f"{channel.mention} unlocked")
+
+    @cog_ext.cog_subcommand(
+        base="lockdown",
+        name="start",
         description="Locks a server",
         guild_ids=[418094694325813248, 578757004059738142, 862402786116763668],
         options=[
@@ -636,7 +692,8 @@ class AdminCog(commands.Cog):
             ),
         ],
     )
-    async def _lockdown(
+    @commands.has_permissions(administrator=True)
+    async def _lockdown_start(
         self,
         ctx: SlashContext,
         reason: str,
@@ -669,6 +726,109 @@ class AdminCog(commands.Cog):
             self.db.jarvis.locks.bulk_write(updates)
         await ctx.send(f"Server locked for {duration} minute(s)")
 
+    @cog_ext.cog_subcommand(
+        base="lockdown",
+        name="end",
+        description="Unlocks a server",
+        guild_ids=[418094694325813248, 578757004059738142, 862402786116763668],
+    )
+    @commands.has_permissions(administrator=True)
+    async def _lockdown_end(
+        self,
+        ctx: SlashContext,
+    ):
+        await ctx.defer()
+        channels = ctx.guild.channels
+        roles = ctx.guild.roles
+        updates = []
+        locks = list(
+            self.db.jarvis.locks.find({"guild": ctx.guild.id, "active": True})
+        )
+        if not locks:
+            await ctx.send("No lockdown detected.")
+            return
+        for channel in channels:
+            for role in roles:
+                try:
+                    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}},
+                )
+            )
+        if updates:
+            self.db.jarvis.locks.bulk_write(updates)
+        await ctx.send(f"Server unlocked")
+
+    @cog_ext.cog_slash(
+        name="warn",
+        description="Warn a user",
+        guild_ids=[418094694325813248, 578757004059738142, 862402786116763668],
+        options=[
+            create_option(
+                name="user",
+                description="User to warn",
+                option_type=6,
+                required=True,
+            ),
+            create_option(
+                name="reason",
+                description="Reason for warning",
+                option_type=3,
+                required=True,
+            ),
+            create_option(
+                name="duration",
+                description="Duration of warning in hours, default 24",
+                option_type=4,
+                required=False,
+            ),
+        ],
+    )
+    @commands.has_permissions(administrator=True)
+    async def _warn(
+        self, ctx: SlashContext, user: User, reason: str, duration: int = 24
+    ):
+        await ctx.defer()
+        self.db.jarvis.warns.insert_one(
+            {
+                "user": user.id,
+                "reason": reason,
+                "admin": ctx.author.id,
+                "time": datetime.now(),
+                "guild": ctx.guild.id,
+                "duration": duration,
+                "active": True,
+            }
+        )
+        count = len(
+            list(
+                self.db.jarvis.warns.find(
+                    {"user": user.id, "guild": ctx.guild.id, "active": True}
+                )
+            )
+        )
+        fields = [Field("Reason", reason, False)]
+        embed = build_embed(
+            title="Warning",
+            description=f"{user.mention} has been warned",
+            fields=fields,
+        )
+        embed.set_author(
+            name=user.nick if user.nick else user.name,
+            icon_url=user.avatar_url,
+        )
+        embed.set_footer(text=f"{user.name}#{user.discriminator} | {user.id}")
+
+        await ctx.send(embed=embed)
+
 
 def setup(bot):
     bot.add_cog(AdminCog(bot))