jarvis-bot/jarvis/cogs/verify.py

91 lines
3.4 KiB
Python

"""J.A.R.V.I.S. Verify Cog."""
import asyncio
from random import randint
from dis_snek import InteractionContext, Scale, Snake
from dis_snek.models.application_commands import slash_command
from dis_snek.models.discord.components import Button, ButtonStyles, spread_to_rows
from dis_snek.models.snek.command import cooldown
from dis_snek.models.snek.cooldowns import Buckets
from jarvis_core.db import q
from jarvis_core.db.models import Setting
def create_layout() -> list:
"""Create verify component layout."""
buttons = []
yes = randint(0, 2) # noqa: S311
for i in range(3):
label = "YES" if i == yes else "NO"
id = f"no_{i}" if not i == yes else "yes"
color = ButtonStyles.GREEN if i == yes else ButtonStyles.RED
buttons.append(
Button(
style=color,
label=label,
custom_id=f"verify_button||{id}",
)
)
return spread_to_rows(*buttons, max_in_row=3)
class VerifyCog(Scale):
"""J.A.R.V.I.S. Verify Cog."""
def __init__(self, bot: Snake):
self.bot = bot
@slash_command(name="verify", description="Verify that you've read the rules")
@cooldown(bucket=Buckets.USER, rate=1, interval=15)
async def _verify(self, ctx: InteractionContext) -> None:
await ctx.defer()
role = await Setting.find_one(q(guild=ctx.guild.id, setting="verified"))
if not role:
await ctx.send("This guild has not enabled verification", delete_after=5)
return
if await ctx.guild.get_role(role.value) in ctx.author.roles:
await ctx.send("You are already verified.", delete_after=5)
return
components = create_layout()
message = await ctx.send(
content=f"{ctx.author.mention}, please press the button that says `YES`.",
components=components,
)
try:
context = await self.bot.wait_for_component(
messages=message, check=lambda x: ctx.author.id == x.author.id, timeout=30
)
correct = context.context.custom_id.split("||")[-1] == "yes"
if correct:
for row in components:
for component in row.components:
component.disabled = True
setting = await Setting.find_one(guild=ctx.guild.id, setting="verified")
role = await ctx.guild.get_role(setting.value)
await ctx.author.add_roles(role, reason="Verification passed")
setting = await Setting.find_one(guild=ctx.guild.id, setting="unverified")
if setting:
role = await ctx.guild.get_role(setting.value)
await ctx.author.remove_roles(role, reason="Verification passed")
await context.context.edit_origin(
content=f"Welcome, {ctx.author.mention}. Please enjoy your stay.",
components=components,
)
await context.context.message.delete(delay=5)
else:
await context.context.edit_origin(
content=(
f"{ctx.author.mention}, incorrect. "
"Please press the button that says `YES`"
)
)
except asyncio.TimeoutError:
await message.delete(delay=30)
def setup(bot: Snake) -> None:
"""Add VerifyCog to J.A.R.V.I.S."""
VerifyCog(bot)