Add unban command (thanks @cpixl, @GlitterButts). Needs testing, references #25

jarvis/cogs/admin.py

@@ -1,9 +1,11 @@
 import jarvis
+import pymongo
+import re
 from datetime import datetime, timedelta
-from discord import User, Member
+from discord import User, Member, BanEntry
 from discord.ext import commands
-from discord.utils import get
+from discord.utils import get, find
-from discord_slash import cog_ext
+from discord_slash import cog_ext, SlashContext
 from discord_slash.utils.manage_commands import create_option, create_choice
 from jarvis.utils.db import DBManager
 from jarvis.utils.permissions import admin_or_permissions
@@ -16,22 +18,62 @@ class AdminCog(commands.Cog):
     Used to manage guilds
     """
 
-    def __init__(self, bot):
+    def __init__(self, bot: commands.Bot):
         self.bot = bot
         config = jarvis.config.get_config()
         self.db = DBManager(config.mongo).mongo
 
+    @cog_ext.cog_slash(
+        name="ban",
+        description="Ban a user",
+        guild_ids=[418094694325813248, 578757004059738142],
+        options=[
+            create_option(
+                name="user",
+                description="User to ban",
+                option_type=6,
+                required=True,
+            ),
+            create_option(
+                name="reason",
+                description="Ban reason",
+                required=True,
+                option_type=3,
+            ),
+            create_option(
+                name="type",
+                description="Ban type",
+                option_type=3,
+                required=False,
+                choices=[
+                    create_choice(value="perm", name="Permanent"),
+                    create_choice(value="temp", name="Temporary"),
+                    create_choice(value="soft", name="Soft"),
+                ],
+            ),
+            create_option(
+                name="duration",
+                description="Ban duration in hours if temporary",
+                required=False,
+                option_type=4,
+            ),
+        ],
+    )
+    @admin_or_permissions(ban_members=True)
     async def _ban(
         self,
-        ctx,
+        ctx: SlashContext,
         user: User = None,
-        type: str = "perm",
         reason: str = None,
+        type: str = "perm",
         length: int = 4,
     ):
         if not user or user == ctx.author:
             await ctx.send("You cannot ban yourself.")
             return
+        if user == self.bot.user:
+            await ctx.send("I'm afraid I can't let you do that")
+            return
         if type == "temp" and length < 0:
             await ctx.send("You cannot set a temp ban to < 0 hours.")
             return
@@ -71,6 +113,8 @@ class AdminCog(commands.Cog):
         self.db.jarvis.bans.insert_one(
             {
                 "user": user.id,
+                "username": user.name,
+                "discrim": user.discriminator,
                 "reason": reason,
                 "admin": ctx.author.id,
                 "time": datetime.now(),
@@ -82,57 +126,224 @@ class AdminCog(commands.Cog):
             }
         )
 
+    async def discord_apply_unban(self, ctx, user, reason):
+        await ctx.guild.unban(user, reason=reason)
+        user: User = user
+        self.db.jarvis.unbans.insert_one(
+            {
+                "user": user.id,
+                "username": user.name,
+                "discrim": user.discriminator,
+                "guild": ctx.guild.id,
+                "admin": ctx.author,
+                "reason": reason,
+                "time": datetime.now(),
+            }
+        )
+        _ = self.db.jarvis.bans.update(
+            {"user": user.id, "guild": ctx.guild.id},
+            {"$set": {"active": False}},
+        )
+        await ctx.send("User successfully unbanned.\nReason: " + reason)
+
     @cog_ext.cog_slash(
-        name="ban",
+        name="unban",
-        description="Ban a user",
+        description="Unban a user",
         guild_ids=[418094694325813248, 578757004059738142],
         options=[
             create_option(
-                name="user",
+                name="id",
-                description="User to ban",
+                description="User ID to unban",
-                option_type=6,
+                option_type=3,
                 required=True,
             ),
-            create_option(
-                name="type",
-                description="Ban type",
-                option_type=3,
-                required=False,
-                choices=[
-                    create_choice("Permanent", "perm"),
-                    create_choice("Temporary", "temp"),
-                    create_choice("Soft", "soft"),
-                ],
-            ),
             create_option(
                 name="reason",
-                description="Ban reason",
+                description="Unban reason",
-                required=False,
+                required=True,
                 option_type=3,
             ),
-            create_option(
-                name="duration",
-                description="Ban duration in hours if temporary",
-                required=False,
-                option_type=4,
-            ),
         ],
     )
     @admin_or_permissions(ban_members=True)
-    async def _ban_slash(
+    async def _unban(
         self,
-        ctx,
+        ctx: SlashContext,
-        user: User = None,
+        user: str,
-        type: str = "perm",
+        reason: str,
-        reason: str = None,
-        length: int = 12,
     ):
-        await self._ban(ctx, user, type, reason, length)
+        ctx.defer()
 
-    async def _kick(self, ctx, user: User, reason=None):
+        orig_user = user
+        discrim = None
+        discord_ban_info: BanEntry = None
+        database_ban_info = None
+
+        bans = await ctx.guild.bans()
+
+        # Try to get ban information out of Discord
+        if re.match("^[0-9]{1,}$"):  # User ID
+            user = int(user)
+            discord_ban_info = find(lambda x: x.id == user, bans)
+        else:  # User name
+            if re.match("#[0-9]{4}$", user):  # User name has discrim
+                user, discrim = user.split("#")
+            if discrim:
+                discord_ban_info = find(
+                    lambda x: x.name == user and x.discriminator == discrim,
+                    bans,
+                )
+            else:
+                results = filter(lambda x: x.name == user, bans)
+                if results:
+                    if len(results) > 1:
+                        # TODO: send list of bans that matched
+                        # for admins to use on next attempt
+                        pass
+                    else:
+                        discord_ban_info = results[0]
+
+        # If we don't have the ban information in Discord,
+        # try to find the relevant information in the database.
+        # We take advantage of the previous checks to save CPU cycles
+        if not discord_ban_info:
+            if isinstance(user, int):
+                database_ban_info = self.db.jarvis.bans.find_one(
+                    {"guild": ctx.guild.id, "user": user, "active": True}
+                )
+            else:
+                search = {
+                    "guild": ctx.guild.id,
+                    "username": user,
+                    "active": True,
+                }
+                if discrim:
+                    search["discrim"] = discrim
+                database_ban_info = self.db.jarvis.bans.find_one(search)
+
+        if not discord_ban_info and not database_ban_info:
+            await ctx.send(f"Unable to find user {orig_user}")
+
+        elif discord_ban_info:
+            await self.discord_apply_unban(ctx, discord_ban_info.user, reason)
+        else:
+            discord_ban_info = find(
+                lambda x: x.id == database_ban_info["id"], bans
+            )
+            if discord_ban_info:
+                await self.discord_apply_unban(
+                    ctx, discord_ban_info.user, reason
+                )
+            else:
+                self.db.jarvis.bans.update_many(
+                    {"user": database_ban_info["id"], "guild": ctx.guild.id},
+                    {"$set": {"active": False}},
+                )
+                self.db.jarvis.unbans.insert_one(
+                    {
+                        "user": discord_ban_info.id,
+                        "username": discord_ban_info.name,
+                        "discrim": discord_ban_info.discriminator,
+                        "guild": ctx.guild.id,
+                        "admin": ctx.author,
+                        "reason": reason,
+                        "time": datetime.now(),
+                    }
+                )
+                await ctx.send(
+                    "Unable to find user in Discord, "
+                    + "but removed entry from database."
+                )
+
+    @cog_ext.cog_subcommand(
+        base="bans",
+        name="list",
+        description="List bans",
+        guild_ids=[418094694325813248, 578757004059738142],
+        options=[
+            create_option(
+                name="type",
+                description="Ban type",
+                option_type=4,
+                required=False,
+                choices=[
+                    create_choice(value=0, name="All"),
+                    create_choice(value=1, name="Permanent"),
+                    create_choice(value=2, name="Temporary"),
+                    create_choice(value=3, name="Soft"),
+                ],
+            ),
+            create_option(
+                name="active",
+                description="Active bans",
+                option_type=5,
+                required=False,
+                choices=[
+                    create_choice(value=True, name="Yes"),
+                    create_choice(value=False, name="No"),
+                ],
+            ),
+        ],
+    )
+    @admin_or_permissions(ban_members=True)
+    async def _bans_list(
+        self, ctx: SlashContext, type: int = 0, active: bool = True
+    ):
+        ctx.defer()
+        types = [0, "perm", "temp", "soft"]
+        search = {"guild": ctx.guild.id, "active": active}
+        if type > 0:
+            search["type"] = types[type]
+        bans = self.db.jarvis.bans.find(search).sort(
+            [("time", pymongo.DESCENDING)]
+        )
+        message = "Bans:\n```\n"
+        db_bans = []
+        for ban in bans:
+            if "username" not in ban:
+                user = await self.bot.fetch_user(ban["user"])
+                ban["username"] = user.name if user else "[deleted user]"
+            message += "[{0}] {1} ({2}): {3}\n".format(
+                ban["time"].strftime("%d-%m-%Y"),
+                ban["username"],
+                ban["user"],
+                ban["reason"],
+            )
+        async for ban in ctx.guild.bans():
+            if ban.user.id not in db_bans:
+                message += "[unknown] {0} ({1}): {2}".format(
+                    ban.user.name, ban.user.id, ban.reason
+                )
+        message += "```"
+        await ctx.send(message)
+
+    @cog_ext.cog_slash(
+        name="kick",
+        description="Kick a user",
+        guild_ids=[418094694325813248, 578757004059738142],
+        options=[
+            create_option(
+                name="user",
+                description="User to kick",
+                option_type=6,
+                required=True,
+            ),
+            create_option(
+                name="reason",
+                description="Kick reason",
+                required=False,
+                option_type=3,
+            ),
+        ],
+    )
+    @admin_or_permissions(kick_members=True)
+    async def _kick(self, ctx: SlashContext, user: User, reason=None):
         if not user or user == ctx.author:
             await ctx.send("You cannot kick yourself.")
             return
+        if user == self.bot.user:
+            await ctx.send("I'm afraid I can't let you do that")
+            return
         if not reason:
             reason = (
                 "Mr. Stark is displeased with your presence. Please leave."
@@ -160,29 +371,24 @@ class AdminCog(commands.Cog):
         )
 
     @cog_ext.cog_slash(
-        name="kick",
+        name="purge",
-        description="Kick a user",
+        description="Purge messages from channel",
-        guild_ids=[418094694325813248, 578757004059738142],
+        guild_ids=[578757004059738142],
         options=[
             create_option(
-                name="user",
+                name="amount",
-                description="User to kick",
+                description="Amount of messages to purge",
-                option_type=6,
-                required=True,
-            ),
-            create_option(
-                name="reason",
-                description="Kick reason",
                 required=False,
-                option_type=3,
+                option_type=4,
-            ),
+            )
         ],
     )
-    @admin_or_permissions(kick_members=True)
+    @admin_or_permissions(manage_messages=True)
-    async def _kick_slash(self, ctx, user: User, reason=None):
+    async def _purge(self, ctx: SlashContext, amount: int = 10):
-        await self._kick(ctx, user, reason)
+        if amount < 1:
-
+            await ctx.send("Amount must be >= 1")
-    async def _purge(self, ctx, amount: int = 10):
+            return
+        await ctx.defer()
         channel = ctx.channel
         messages = []
         async for message in channel.history(limit=amount + 1):
@@ -198,27 +404,6 @@ class AdminCog(commands.Cog):
             }
         )
 
-    @cog_ext.cog_slash(
-        name="purge",
-        description="Purge messages from channel",
-        guild_ids=[578757004059738142],
-        options=[
-            create_option(
-                name="amount",
-                description="Amount of messages to purge",
-                required=False,
-                option_type=4,
-            )
-        ],
-    )
-    @admin_or_permissions(manage_messages=True)
-    async def _purge_slash(self, ctx, amount: int = 10):
-        if amount < 1:
-            await ctx.send("Amount must be >= 1")
-            return
-        await ctx.defer()
-        await self._purge(ctx, amount)
-
     @cog_ext.cog_slash(
         name="mute",
         description="Mute a user",
@@ -245,11 +430,16 @@ class AdminCog(commands.Cog):
         ],
     )
     @admin_or_permissions(mute_members=True)
-    async def _mute(self, ctx, user: Member, reason: str, length: int = 30):
+    async def _mute(
+        self, ctx: SlashContext, user: Member, reason: str, length: int = 30
+    ):
         ctx.defer()
         if user == ctx.author:
             await ctx.send("You cannot mute yourself.")
             return
+        if user == self.bot.user:
+            await ctx.send("I'm afraid I can't let you do that")
+            return
         mute_setting = self.db.jarvis.settings.find_one(
             {"guild": ctx.guild.id, "setting": "mute"}
         )
@@ -302,14 +492,15 @@ class AdminCog(commands.Cog):
         ],
     )
     @admin_or_permissions(mute_members=True)
-    async def _unmute(self, ctx, user: Member):
+    async def _unmute(self, ctx: SlashContext, user: Member):
         ctx.defer()
         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."
+                "Please configure a mute role with "
+                + "/settings mute <role> first."
             )
             return