jarvis-bot/jarvis/cogs/verify.py

104 lines
3.6 KiB
Python

from random import randint
from discord import Role
from discord.ext import commands
from discord_slash import ComponentContext, SlashContext, cog_ext
from discord_slash.model import ButtonStyle
from discord_slash.utils import manage_commands, manage_components
from discord_slash.utils.manage_commands import create_option
import jarvis
from jarvis.config import get_config
from jarvis.utils.db import DBManager
def create_layout():
buttons = []
yes = randint(0, 2)
for i in range(3):
label = "YES" if i == yes else "NO"
id = f"no_{i}" if not i == yes else "yes"
color = ButtonStyle.green if i == yes else ButtonStyle.red
buttons.append(
manage_components.create_button(
style=color,
label=label,
custom_id=f"verify_button||{id}",
)
)
action_row = manage_components.spread_to_rows(*buttons, max_in_row=3)
return action_row
class VerifyCog(commands.Cog):
def __init__(self, bot):
self.bot = bot
self.db = DBManager(get_config().mongo).mongo
@cog_ext.cog_slash(
name="verify",
description="Verify that you've read the rules",
guild_ids=[418094694325813248, 578757004059738142],
)
async def _verify(self, ctx: SlashContext):
await ctx.defer()
role = self.db.jarvis.settings.find_one(
{"setting": "verified", "guild": ctx.guild.id}
)
if not role:
await ctx.send(
"This guild has not enabled verification", delete_after=5
)
return
if 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,
)
await message.delete(delay=15)
@cog_ext.cog_component(components=create_layout())
async def _process(self, ctx: ComponentContext):
await ctx.defer(edit_origin=True)
try:
if ctx.author.id != ctx.origin_message.mentions[0].id:
return
except Exception:
return
correct = ctx.custom_id.split("||")[-1] == "yes"
if correct:
components = ctx.origin_message.components
for c in components:
for c2 in c["components"]:
c2["disabled"] = True
setting = self.db.jarvis.settings.find_one(
{"setting": "verified", "guild": ctx.guild.id}
)
role = ctx.guild.get_role(setting["value"])
await ctx.author.add_roles(role, reason="Verified")
setting = self.db.jarvis.settings.find_one(
{"setting": "unverified", "guild": ctx.guild.id}
)
if setting:
role = ctx.guild.get_role(setting["value"])
await ctx.author.remove_roles(role, reason="Verified")
await ctx.edit_origin(
content=f"Welcome, {ctx.author.mention}. Please enjoy your stay.",
components=manage_components.spread_to_rows(
*components, max_in_row=5
),
)
await ctx.origin_message.delete(delay=5)
else:
await ctx.edit_origin(
content=f"{ctx.author.mention}, incorrect. "
+ "Please press the button that says `YES`",
)
def setup(bot):
bot.add_cog(VerifyCog(bot))