From 959a7609fbf0003486ef570003be5e34994f5ec4 Mon Sep 17 00:00:00 2001
From: 0x5c <0x5c.dev@gmail.com>
Date: Mon, 7 Oct 2019 09:11:03 -0400
Subject: [PATCH] main.py: added checks to restart and shutdown (#27)

- Added a check for ownership of the bot
- Added a helper function to add reactions
  (and handle lack of permissions)

Fixes #7
---
 main.py | 46 ++++++++++++++++++++++++++++------------------
 1 file changed, 28 insertions(+), 18 deletions(-)

diff --git a/main.py b/main.py
index 4601d51..d4ac98c 100644
--- a/main.py
+++ b/main.py
@@ -43,35 +43,45 @@ class GlobalSettings(commands.Cog):
 bot = commands.Bot(command_prefix=opt.prefix, description=info.description, help_command=commands.MinimalHelpCommand())
 
 
+# --- Helper functions ---
+
+async def add_react(msg: discord.Message, react: str):
+    try:
+        await msg.add_reaction(react)
+    except discord.Forbidden:
+        print(f"!! Missing permissions to add reaction in '{msg.guild.id}/{msg.channel.id}'!")
+
+
+# --- Checks ---
+
+async def check_if_owner(ctx: commands.Context):
+    if ctx.author.id in opt.owners_uids:
+        return True
+    else:
+        await add_react(ctx.message, "❌")
+        return False
+
+
 # --- Commands ---
 
 @bot.command(name="restart")
+@commands.check(check_if_owner)
 async def _restart_bot(ctx):
     """Restarts the bot."""
     global exit_code
-    if ctx.author.id in opt.owners_uids:
-        await ctx.message.add_reaction("✅")
-        exit_code = 42  # Signals to the wrapper script that the bot needs to be restarted.
-        await bot.logout()
-    else:
-        try:
-            await ctx.message.add_reaction("❌")
-        except:
-            return
+    await add_react(ctx.message, "✅")
+    exit_code = 42  # Signals to the wrapper script that the bot needs to be restarted.
+    await bot.logout()
+
 
 @bot.command(name="shutdown")
+@commands.check(check_if_owner)
 async def _shutdown_bot(ctx):
     """Shuts down the bot."""
     global exit_code
-    if ctx.author.id in opt.owners_uids:
-        await ctx.message.add_reaction("✅")
-        exit_code = 0  # Signals to the wrapper script that the bot should not be restarted.
-        await bot.logout()
-    else:
-        try:
-            await ctx.message.add_reaction("❌")
-        except:
-            return
+    await add_react(ctx.message, "✅")
+    exit_code = 0  # Signals to the wrapper script that the bot should not be restarted.
+    await bot.logout()
 
 
 # --- Events ---