From dd1c304522ec9eeddfbbbdcf3de4328795782d5a Mon Sep 17 00:00:00 2001
From: imcompiling <develop@imcompiling.com>
Date: Wed, 8 Nov 2023 16:17:16 -0500
Subject: [PATCH] Add role check to nominate button.

---
 classes/election_cycle.py |  2 +-
 classes/nominate_views.py | 44 ++++++++++++++++++++++++++++++++-------
 dakiny_bot.py             |  6 ++++--
 3 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/classes/election_cycle.py b/classes/election_cycle.py
index 3197891..373db04 100644
--- a/classes/election_cycle.py
+++ b/classes/election_cycle.py
@@ -210,7 +210,7 @@ class ElectionCycle:
         for index, message in enumerate(messages):
             # if last item, add the nominate button
             if index + 1 == len(messages):
-                await thread.send(message, view=NominateView(self.config))
+                await thread.send(message, view=NominateView(self.config, self.guild))
             else:
                 await thread.send(message)
 
diff --git a/classes/nominate_views.py b/classes/nominate_views.py
index d51a752..23dda9e 100644
--- a/classes/nominate_views.py
+++ b/classes/nominate_views.py
@@ -19,7 +19,13 @@ class NominateModal(nextcord.ui.Modal):
     Modal for nominating oneself. User writes a speech and then the callback sends the speech into the correct channel.
     """
 
-    def __init__(self, config: DakinyConfig, usr_id: int, usr_nominated: bool, speech_content: str):
+    def __init__(
+        self,
+        config: DakinyConfig,
+        usr_id: int,
+        usr_nominated: bool,
+        speech_content: str,
+    ):
         super().__init__(
             "Nomination",
             timeout=5 * 60,  # 5 minutes
@@ -70,8 +76,12 @@ class NominateModal(nextcord.ui.Modal):
                 )
                 obj = session.scalar(stmt)
 
-                msg: nextcord.Message = await interaction.channel.fetch_message(obj.SpeechMsgId)
-                await msg.edit(content=decorate_nomination(interaction.user.id, self.speech.value))
+                msg: nextcord.Message = await interaction.channel.fetch_message(
+                    obj.SpeechMsgId
+                )
+                await msg.edit(
+                    content=decorate_nomination(interaction.user.id, self.speech.value)
+                )
 
 
 class NominateView(nextcord.ui.View):
@@ -79,12 +89,19 @@ class NominateView(nextcord.ui.View):
     The button to press to start your nomination.
     """
 
-    def __init__(self, config):
+    def __init__(self, config: DakinyConfig, guild: nextcord.Guild):
         super().__init__(timeout=None)
         self.config: DakinyConfig = config
+        self.guild: nextcord.Guild = guild
 
-    @nextcord.ui.button(label="Nominate Yourself", style=nextcord.ButtonStyle.primary, custom_id="nominate")
-    async def nominate(self, button: nextcord.ui.Button, interaction: nextcord.Interaction) -> None:
+    @nextcord.ui.button(
+        label="Nominate Yourself",
+        style=nextcord.ButtonStyle.primary,
+        custom_id="nominate",
+    )
+    async def nominate(
+        self, button: nextcord.ui.Button, interaction: nextcord.Interaction
+    ) -> None:
         """
         Gathers previous information about the user if they have already nominated themselves, then launches
         the nomination modal to get their speech.
@@ -94,6 +111,17 @@ class NominateView(nextcord.ui.View):
             interaction (nextcord.Interaction): Interaction of user who pressed the button
         """
 
+        # enforce members only
+        if (
+            self.guild.get_role(self.config.params.clan_member_role_id)
+            not in interaction.user.roles
+        ):
+            await interaction.response.send_message(
+                f"You must have the <@&{self.config.params.clan_member_role_id}> role to vote.",
+                ephemeral=True,
+            )
+            return
+
         cycle = get_election_cycle_name(self.config)
         usr_nominated = False
         speech_content = ""
@@ -110,7 +138,9 @@ class NominateView(nextcord.ui.View):
                 speech_content = msg.content.split(NOMINEE_SPLITTER)[1]
 
         await interaction.response.send_modal(
-            NominateModal(self.config, interaction.user.id, usr_nominated, speech_content)
+            NominateModal(
+                self.config, interaction.user.id, usr_nominated, speech_content
+            )
         )
 
 
diff --git a/dakiny_bot.py b/dakiny_bot.py
index 0c0cc53..ab9bcf6 100644
--- a/dakiny_bot.py
+++ b/dakiny_bot.py
@@ -20,7 +20,9 @@ birthday_list: list[int] = []
 logger = logging.getLogger("nextcord")
 logger.setLevel(logging.DEBUG)
 handler = logging.FileHandler(filename="nextcord.log", encoding="utf-8", mode="w")
-handler.setFormatter(logging.Formatter("%(asctime)s:%(levelname)s:%(name)s: %(message)s"))
+handler.setFormatter(
+    logging.Formatter("%(asctime)s:%(levelname)s:%(name)s: %(message)s")
+)
 logger.addHandler(handler)
 
 
@@ -51,7 +53,7 @@ def main(cfg: DakinyConfig):
         """When Discord is connected"""
         assert bot.user is not None
         create_database  # pylint: disable=pointless-statement
-        bot.add_view(NominateView(cfg))
+        bot.add_view(NominateView(cfg, bot.get_guild(cfg.params.guild_id)))
         bot.add_view(VoteView(cfg, bot.get_guild(cfg.params.guild_id)))
         print(f"{bot.user.name} has connected to Discord!")