jarvis-bot/jarvis/cogs/admin/warning.py

"""J.A.R.V.I.S. WarningCog."""
from datetime import datetime, timedelta

from dis_snek import InteractionContext, Permissions, Snake
from dis_snek.ext.paginators import Paginator
from dis_snek.models.discord.user import User
from dis_snek.models.snek.application_commands import (
    OptionTypes,
    SlashCommandChoice,
    slash_command,
    slash_option,
)

from jarvis.db.models import Warning
from jarvis.utils import build_embed
from jarvis.utils.cachecog import CacheCog
from jarvis.utils.field import Field
from jarvis.utils.permissions import admin_or_permissions


class WarningCog(CacheCog):
    """J.A.R.V.I.S. WarningCog."""

    def __init__(self, bot: Snake):
        super().__init__(bot)

    @slash_command(name="warn", description="Warn a user")
    @slash_option(
        name="user", description="User to warn", option_type=OptionTypes.USER, required=True
    )
    @slash_option(
        name="reason",
        description="Reason for warning",
        option_type=OptionTypes.STRING,
        required=True,
    )
    @slash_option(
        name="duration",
        description="Duration of warning in hours, default 24",
        option_type=OptionTypes.INTEGER,
        required=False,
    )
    @admin_or_permissions(Permissions.MANAGE_GUILD)
    async def _warn(
        self, ctx: InteractionContext, user: User, reason: str, duration: int = 24
    ) -> None:
        if len(reason) > 100:
            await ctx.send("Reason must be < 100 characters", hidden=True)
            return
        if duration <= 0:
            await ctx.send("Duration must be > 0", hidden=True)
            return
        elif duration >= 120:
            await ctx.send("Duration must be < 5 days", hidden=True)
            return
        await ctx.defer()
        _ = Warning(
            user=user.id,
            reason=reason,
            admin=ctx.author.id,
            guild=ctx.guild.id,
            duration=duration,
            active=True,
        ).save()
        fields = [Field("Reason", reason, False)]
        embed = build_embed(
            title="Warning",
            description=f"{user.mention} has been warned",
            fields=fields,
        )
        embed.set_author(
            name=user.display_name,
            icon_url=user.display_avatar,
        )
        embed.set_footer(text=f"{user.name}#{user.discriminator} | {user.id}")

        await ctx.send(embed=embed)

    @slash_command(name="warnings", description="Get count of user warnings")
    @slash_option(
        name="user", description="User to view", option_type=OptionTypes.USER, required=True
    )
    @slash_option(
        name="active",
        description="View active only",
        option_type=OptionTypes.INTEGER,
        required=False,
        choices=[
            SlashCommandChoice(name="Yes", value=1),
            SlashCommandChoice(name="No", value=0),
        ],
    )
    @admin_or_permissions(Permissions.MANAGE_GUILD)
    async def _warnings(self, ctx: InteractionContext, user: User, active: bool = 1) -> None:
        active = bool(active)
        exists = self.check_cache(ctx, user_id=user.id, active=active)
        if exists:
            await ctx.defer(hidden=True)
            await ctx.send(
                f"Please use existing interaction: {exists['paginator']._message.jump_url}",
                hidden=True,
            )
            return
        warnings = Warning.objects(
            user=user.id,
            guild=ctx.guild.id,
        ).order_by("-created_at")
        active_warns = Warning.objects(user=user.id, guild=ctx.guild.id, active=True).order_by(
            "-created_at"
        )

        pages = []
        if active:
            if active_warns.count() == 0:
                embed = build_embed(
                    title="Warnings",
                    description=f"{warnings.count()} total | 0 currently active",
                    fields=[],
                )
                embed.set_author(name=user.username, icon_url=user.display_avatar)
                embed.set_thumbnail(url=ctx.guild.icon_url)
                pages.append(embed)
            else:
                fields = []
                for warn in active_warns:
                    admin = await ctx.guild.get_member(warn.admin)
                    admin_name = "||`[redacted]`||"
                    if admin:
                        admin_name = f"{admin.username}#{admin.discriminator}"
                    fields.append(
                        Field(
                            name=warn.created_at.strftime("%Y-%m-%d %H:%M:%S UTC"),
                            value=f"{warn.reason}\nAdmin: {admin_name}\n\u200b",
                            inline=False,
                        )
                    )
                for i in range(0, len(fields), 5):
                    embed = build_embed(
                        title="Warnings",
                        description=(
                            f"{warnings.count()} total | {active_warns.count()} currently active"
                        ),
                        fields=fields[i : i + 5],
                    )
                    embed.set_author(
                        name=user.username + "#" + user.discriminator,
                        icon_url=user.display_avatar,
                    )
                    embed.set_thumbnail(url=ctx.guild.icon_url)
                    embed.set_footer(text=f"{user.username}#{user.discriminator} | {user.id}")
                    pages.append(embed)
        else:
            fields = []
            for warn in warnings:
                title = "[A] " if warn.active else "[I] "
                title += warn.created_at.strftime("%Y-%m-%d %H:%M:%S UTC")
                fields.append(
                    Field(
                        name=title,
                        value=warn.reason + "\n\u200b",
                        inline=False,
                    )
                )
            for i in range(0, len(fields), 5):
                embed = build_embed(
                    title="Warnings",
                    description=(
                        f"{warnings.count()} total | {active_warns.count()} currently active"
                    ),
                    fields=fields[i : i + 5],
                )
                embed.set_author(
                    name=user.username + "#" + user.discriminator,
                    icon_url=user.display_avatar,
                )
                embed.set_thumbnail(url=ctx.guild.icon_url)
                pages.append(embed)

        paginator = Paginator(bot=self.bot, *pages, timeout=300)

        self.cache[hash(paginator)] = {
            "guild": ctx.guild.id,
            "user": ctx.author.id,
            "timeout": datetime.utcnow() + timedelta(minutes=5),
            "command": ctx.subcommand_name,
            "user_id": user.id,
            "active": active,
            "paginator": paginator,
        }

        await paginator.send(ctx)