Add temprole, closes #128

jarvis/cogs/temprole.py

@@ -0,0 +1,120 @@
+"""JARVIS temporary role handler."""
+import logging
+from datetime import datetime, timezone
+
+from dateparser import parse
+from dateparser_data.settings import default_parsers
+from dis_snek import InteractionContext, Permissions, Scale, Snake
+from dis_snek.models.discord.embed import EmbedField
+from dis_snek.models.discord.role import Role
+from dis_snek.models.discord.user import Member
+from dis_snek.models.snek.application_commands import (
+    OptionTypes,
+    slash_command,
+    slash_option,
+)
+from dis_snek.models.snek.command import check
+from jarvis_core.db.models import Temprole
+
+from jarvis.utils import build_embed
+from jarvis.utils.permissions import admin_or_permissions
+
+
+class TemproleCog(Scale):
+    """JARVIS Temporary Role Cog."""
+
+    def __init__(self, bot: Snake):
+        self.bot = bot
+        self.logger = logging.getLogger(__name__)
+
+    @slash_command(name="temprole", description="Give a user a temporary role")
+    @slash_option(
+        name="user", description="User to grant role", opt_type=OptionTypes.USER, required=True
+    )
+    @slash_option(
+        name="role", description="Role to grant", opt_type=OptionTypes.ROLE, required=True
+    )
+    @slash_option(
+        name="duration",
+        description="Duration of temp role (i.e. 2 hours)",
+        opt_type=OptionTypes.STRING,
+        required=True,
+    )
+    @slash_option(
+        name="reason",
+        description="Reason for temporary role",
+        opt_type=OptionTypes.STRING,
+        required=False,
+    )
+    @check(admin_or_permissions(Permissions.MANAGE_ROLES))
+    async def _temprole(
+        self, ctx: InteractionContext, user: Member, role: Role, duration: str, reason: str = None
+    ) -> None:
+        if not isinstance(user, Member):
+            await ctx.send("User not in guild", ephemeral=True)
+            return
+
+        if role.id == ctx.guild.id:
+            await ctx.send("Cannot add `@everyone` to users", ephemeral=True)
+            return
+
+        if role.bot_managed or not role.is_assignable:
+            await ctx.send(
+                "Cannot assign this role, try lowering it below my role or using a different role",
+                ephemeral=True,
+            )
+            return
+
+        base_settings = {
+            "PREFER_DATES_FROM": "future",
+            "TIMEZONE": "UTC",
+            "RETURN_AS_TIMEZONE_AWARE": True,
+        }
+        rt_settings = base_settings.copy()
+        rt_settings["PARSERS"] = [
+            x for x in default_parsers if x not in ["absolute-time", "timestamp"]
+        ]
+
+        rt_duration = parse(duration, settings=rt_settings)
+
+        at_settings = base_settings.copy()
+        at_settings["PARSERS"] = [x for x in default_parsers if x != "relative-time"]
+        at_duration = parse(duration, settings=at_settings)
+
+        if rt_duration:
+            duration = rt_duration
+        elif at_duration:
+            duration = at_duration
+        else:
+            self.logger.debug(f"Failed to parse duration: {duration}")
+            await ctx.send(f"`{duration}` is not a parsable date, please try again", ephemeral=True)
+            return
+
+        if duration < datetime.now(tz=timezone.utc):
+            await ctx.send(
+                f"`{duration}` is in the past. Past durations aren't allowed", ephemeral=True
+            )
+            return
+
+        await user.add_role(role, reason=reason)
+        await Temprole(
+            guild=ctx.guild.id, user=user.id, role=role.id, admin=ctx.author.id, expires_at=duration
+        ).commit()
+
+        ts = int(duration.timestamp())
+
+        fields = (
+            EmbedField(name="Role", value=role.mention),
+            EmbedField(name="Valid Until", value=f"<t:{ts}:F> (<t:{ts}:R>)"),
+        )
+
+        embed = build_embed(
+            title="Role granted",
+            description=f"Role temporarily granted to {user.mention}",
+            fields=fields,
+        )
+        embed.set_author(
+            name=f"{user.username}#{user.discriminator}", icon_url=user.display_avatar.url
+        )
+
+        await ctx.send(embed=embed)