Heimdallr/commands/infractions.py

import logging
from typing import List
from naff import (
    Extension,
    slash_command,
    slash_option,
    OptionTypes,
    InteractionContext,
    Embed,
    EmbedField,
    Permissions,
    Member,
    Client,
    ActionRow,
    Button,
    ButtonStyles,
)
from database import (
    Infractions as InfractionsModel,
)


class Infractions(Extension):
    def __init__(self, client):
        self.client: Client = client

    @slash_command(
        name="infractions", description="View your infractions", dm_permission=False
    )
    async def infractions_command(self, ctx: InteractionContext):
        await ctx.defer(ephemeral=True)
        infractions: List[InfractionsModel] = InfractionsModel.select().where(
            InfractionsModel.guild_id == int(ctx.guild_id),
            InfractionsModel.user_id == int(ctx.author.id),
        )

        if len(infractions) == 0:
            await ctx.send("You have no infractions.", ephemeral=True)
            return

        embeds: List[Embed] = []
        for infraction in infractions:  # pylint: disable=not-an-iterable
            embed = Embed(
                title="Infraction",
                description=f"{infraction.reason}",
                color=infraction_colour(infraction.weight),
                fields=[
                    EmbedField(
                        name="**Received**",
                        value=f"<t:{int(infraction.at_time.timestamp())}:F>",
                    ),
                    EmbedField(name="**Severity**", value=f"{infraction.weight}"),
                ],
            )
            embeds.append(embed)

        await ctx.send(embed=embeds, ephemeral=True)

    @slash_command(
        name="user-infractions",
        description="View a user's infractions",
        dm_permission=False,
        default_member_permissions=Permissions.KICK_MEMBERS,
    )
    @slash_option(
        name="user",
        description="User to view",
        required=True,
        opt_type=OptionTypes.USER,
    )
    async def user_infractions(self, ctx: InteractionContext, user: Member):
        await ctx.defer(ephemeral=False)
        infractions: List[InfractionsModel] = InfractionsModel.select().where(
            InfractionsModel.guild_id == int(ctx.guild_id),
            InfractionsModel.user_id == int(ctx.author.id),
        )

        if len(infractions) == 0:
            await ctx.send(f"{user.mention} has no infractions.")
            return

        embeds: List[Embed] = []
        for infraction in infractions:  # pylint: disable=not-an-iterable
            issuer = await self.client.fetch_member(
                guild_id=ctx.guild, user_id=infraction.given_by
            )
            embed = Embed(
                title=f"Infraction for user {user.display_name} ({user.username}#{user.discriminator}, {user.id})", # pylint: disable=line-too-long
                description=f"{infraction.reason}",
                color=infraction_colour(infraction.weight),
                fields=[
                    EmbedField(name="**ID**", value=f"{infraction.id}"),
                    EmbedField(
                        name="**Received**",
                        value=f"<t:{int(infraction.at_time.timestamp())}:F>",
                    ),
                    EmbedField(name="**Severity**", value=f"{infraction.weight}"),
                    EmbedField(name="**Issued by**", value=f"{issuer.display_name}"),
                ],
            )
            embeds.append(embed)

        await ctx.send(embed=embeds, ephemeral=False)

    @slash_command(
        name="warn",
        description="Warn a user",
        dm_permission=False,
        default_member_permissions=Permissions.KICK_MEMBERS,
    )
    @slash_option(
        name="user",
        description="User to warn",
        required=True,
        opt_type=OptionTypes.USER,
    )
    @slash_option(
        name="reason",
        description="Reason for warning",
        required=False,
        opt_type=OptionTypes.STRING,
    )
    @slash_option(
        name="weight",
        description="Severity of warning",
        required=False,
        opt_type=OptionTypes.NUMBER,
        min_value=0.0,
        max_value=10.0,
    )
    @slash_option(
        name="silent",
        description="Silent warning (will not notify user)",
        required=False,
        opt_type=OptionTypes.BOOLEAN,
    )
    async def warn_user(
        self,
        ctx: InteractionContext,
        user: Member,
        reason: str = None,
        weight: float = None,
        silent: bool = None,
    ):
        await ctx.defer(ephemeral=False)
        if weight is None:
            weight = 1.0
        if silent is None:
            silent = False
        if reason is None:
            reason = ""
        infraction = InfractionsModel.create(
            guild_id=int(ctx.guild_id),
            user_id=int(user.id),
            given_by=int(ctx.author.id),
            reason=reason,
            weight=weight,
            silent=silent,
        )
        warning_msg = ...
        if not silent:
            try:
                warning_msg = await user.send(
                    embed=Embed(
                        title=f"You have received a warning in {ctx.guild.name}",
                        description=f"{reason}",
                        color=infraction_colour(weight),
                        fields=[
                            EmbedField(name="**Severity**", value=f"{weight}"),
                            EmbedField(
                                name="**Issued at**",
                                value=f"<t:{int(infraction.at_time.timestamp())}:F>",  # pylint: disable=no-member
                            ),
                        ],
                    )
                )
            except Exception:  # pylint: disable=broad-except
                warning_msg = None

        await ctx.send(
            f'Warned {user.display_name} ({user.username}#{user.discriminator}, {user.id}) for "{reason}" with severity {weight}',  # pylint: disable=line-too-long
            ephemeral=True,
        )
        if not silent and warning_msg is None:
            await ctx.send(
                f"{user.mention} has been warned, but I couldn't DM them.",
                ephemeral=True,
            )

    @slash_command(
        name="remove-infraction",
        description="Remove an infraction",
        dm_permission=False,
        default_member_permissions=Permissions.KICK_MEMBERS,
    )
    @slash_option(
        name="infraction-id",
        description="ID of infraction to remove",
        required=True,
        opt_type=OptionTypes.INTEGER,
    )
    async def remove_infraction(self, ctx: InteractionContext, infraction_id: int):
        await ctx.defer(ephemeral=True)
        infraction: InfractionsModel = InfractionsModel.get_or_none(
            InfractionsModel.id == infraction_id
        )
        if infraction is None:
            await ctx.send("That infraction doesn't exist.", ephemeral=True)
            return

        if infraction.guild_id != int(ctx.guild_id):
            await ctx.send("That infraction doesn't exist.", ephemeral=True)
            return

        # pylint: disable=unexpected-keyword-arg
        components = [
            ActionRow(
                Button(
                    custom_id=f"remove-infraction:remove:{infraction.id}",
                    style=ButtonStyles.DANGER,
                    label="Remove infraction",
                ),
                Button(
                    custom_id=f"remove-infraction:cancel:{infraction.id}",
                    style=ButtonStyles.SECONDARY,
                    label="Cancel",
                ),
            )
        ]

        await ctx.send(
            content="Remove this infraction? (times out in 60 seconds)",
            embed=Embed(
                title="Infraction",
                description=f"{infraction.reason}",
                color=infraction_colour(infraction.weight),
                fields=[
                    EmbedField(name="**ID**", value=f"{infraction.id}"),
                    EmbedField(
                        name="**Received**",
                        value=f"<t:{int(infraction.at_time.timestamp())}:F>",
                    ),
                    EmbedField(name="**Severity**", value=f"{infraction.weight}"),
                    EmbedField(name="**Issued by**", value=f"{infraction.given_by}"),
                ],
            ),
            ephemeral=True,
            components=components,
        )

        try:
            used_comp = await self.client.wait_for_component(
                components=components, timeout=60
            )
        except TimeoutError:
            await ctx.send("Timed out.", ephemeral=True)
            return
        else:
            (group, action, inf_id) = used_comp.context.custom_id.split(":")
            if group != "remove-infraction" or action == "cancel":
                await used_comp.context.send("Cancelled.", ephemeral=True)
                return

            if action == "remove" and inf_id == str(infraction.id):
                infraction.delete_instance()
                await used_comp.context.send("Removed.", ephemeral=True)
                return


def infraction_colour(w: float) -> int:
    if w < 0.5:
        return 0xBCBCBC
    if w < 1.0:
        return 0xFF7711

    return 0xFF2211


def setup(client):
    InfractionsModel.create_table()
    Infractions(client)
    logging.info("Infractions extension loaded.")