Command groups and cogs

Pipfile

@@ -16,5 +16,8 @@ pytest-cov = "*"
 pytest = "*"
 exceptiongroup = "*"
 
+[requires]
+python_version = "3.11"
+
 [scripts]
 test = "python -m pytest --cov=app --cov-report=html"

app/bot.py

@@ -1,53 +1,12 @@
 """
 Discord bot that controls the power of a physical server
 """
-import os
-import sys
-import json
 import traceback
 from asyncio import TimeoutError as asyncTimeoutError
 from typing import List
 import requests
 import discord
-from discord.ext import commands
-
-def env_defined(key):
-    """
-    Checks if a given env var key is defined in the OS environment
-    """
-    return key in os.environ and len(os.environ[key]) > 0
-
-
-DISCORD_CHANNEL: List[str] = []
-
-# env variables are defaults, if no config file exists it'll be created.
-# If no env is set, stop the bot
-try:
-    DISCORD_TOKEN = os.environ["DISCORD_TOKEN"]
-    WOL_URL = os.environ["WOL_URL"]
-    SHUTDOWN_URL = os.environ["SHUTDOWN_URL"]
-    REBOOT_URL = os.environ["REBOOT_URL"]
-    LIVENESS_URL = os.environ["LIVENESS_URL"]
-except KeyError as e:
-    print(f"Missing {str(e)} token from .env.")
-    sys.exit()
-
-# Defaulting COOLDOWN to 300s unless set by the user
-if env_defined("COOLDOWN"):
-    COOLDOWN: int = int(os.environ["COOLDOWN"])
-else:
-    COOLDOWN: int = 300
-
-# Defaulting POWERBOT_ROLE to "@everyone" unless set by the user
-if env_defined("POWERBOT_ROLE"):
-    POWERBOT_ROLE = os.environ["POWERBOT_ROLE"].split(",")
-    # commands.has_any_role() takes either a role name as string,
-    # or a role ID as integer. Can't just map() a mixed list.
-    for i in range(0,len(POWERBOT_ROLE)):
-        if POWERBOT_ROLE[i].isdigit():
-            POWERBOT_ROLE[i] = int(POWERBOT_ROLE[i])
-else:
-    POWERBOT_ROLE = "@everyone"
+from cogs import config
 
 intents = discord.Intents.default()
 DESC = "Bot to control the power to physical game server"
@@ -106,106 +65,11 @@ async def on_ready():
     await bot.change_presence(status=discord.Status.idle, activity=game)
     print('Connected to API')
 
-
-@bot.slash_command(name="boot", description="Boots the game server")
-@commands.cooldown(rate=1, per=float(COOLDOWN), type=commands.BucketType.guild)
-@commands.check(check_cooldown)
-@commands.has_any_role(*POWERBOT_ROLE) # https://github.com/Pycord-Development/pycord/issues/974
-async def _boot(ctx):
-    try:
-        response = requests.get(WOL_URL, timeout=2)
-        jsonresponse = json.loads(response.content.decode())
-        if jsonresponse.get('success') is True:
-            game = discord.Activity(
-                name="Booting...", type=discord.ActivityType.playing)
-            await bot.change_presence(status=discord.Status.online, activity=game)
-            await ctx.respond('Server booted!')
-    except Exception:
-        await ctx.respond('Something went wrong, have an adult check the logs')
-        traceback.print_exc()
-
-
-@bot.slash_command(name="shutdown", description="Shuts down the game server")
-@commands.cooldown(rate=1, per=float(COOLDOWN), type=commands.BucketType.guild)
-@commands.check(check_cooldown)
-@commands.has_any_role(*POWERBOT_ROLE) # https://github.com/Pycord-Development/pycord/issues/974
-async def _shutdown(ctx):
-    try:
-        response = requests.get(SHUTDOWN_URL, timeout=2)
-        if response.status_code == 200:
-            game = discord.Activity(
-                name="Powering down...", type=discord.ActivityType.playing)
-            await bot.change_presence(status=discord.Status.do_not_disturb, activity=game)
-            await ctx.respond('Server shut down!')
-    except Exception:
-        await ctx.respond('Server is already offline')
-        traceback.print_exc()
-
-
-@bot.slash_command(name="reboot", description="Reboots the game server")
-@commands.cooldown(rate=1, per=float(COOLDOWN), type=commands.BucketType.guild)
-@commands.check(check_cooldown)
-@commands.has_any_role(*POWERBOT_ROLE) # https://github.com/Pycord-Development/pycord/issues/974
-async def _reboot(ctx):
-    try:
-        response = requests.get(REBOOT_URL, timeout=2)
-        if response.status_code == 200:
-            game = discord.Activity(
-                name="Rebooting...", type=discord.ActivityType.playing)
-            await bot.change_presence(status=discord.Status.streaming, activity=game)
-            await ctx.respond('Server rebooting!')
-    except Exception:
-        await ctx.respond('Server is already offline')
-        traceback.print_exc()
-
-
-@_boot.error
-@_shutdown.error
-@_reboot.error
-async def _error(ctx, error):
-    await on_application_command_error(ctx, error)
-
-
-@bot.slash_command(name="sudo", description="Use commands regardless of their cooldown")
-@commands.has_any_role(*POWERBOT_ROLE) # https://github.com/Pycord-Development/pycord/issues/974
-@discord.option(
-    "command",
-    description = "Command to be run ignoring any cooldown.",
-    choices = ["boot", "reboot", "shutdown"]
-)
-async def _sudo(ctx, command):
-    """
-    Allows to bypass cooldowns that are usually placed upon power functions
-    by resetting them all and then invoking the supplied command.
-    """
-    embed = discord.Embed(type="rich", colour=discord.Colour.red())
-    embed.title = '<:warning:1043511363441537046>' \
-                  ' WARNING <:warning:1043511363441537046>'
-    embed.description = f'Are you sure that you want to force `{command}`? ' \
-                        'If yes, react with <:sos:1043671788007211108>.'
-    res = await ctx.respond(embed=embed)
-    msg = await res.original_response()
-    await msg.add_reaction('\N{Squared SOS}') # default emojis need to be unicode
-
-    def check(reaction, user):
-        return user == ctx.author and str(reaction.emoji) == '\N{Squared SOS}'
-    try:
-        await bot.wait_for('reaction_add', timeout=5.0, check=check)
-    except asyncTimeoutError:
-        await msg.clear_reactions()
-    else:
-        bot.get_application_command(name="boot").reset_cooldown(ctx)
-        bot.get_application_command(name="reboot").reset_cooldown(ctx)
-        bot.get_application_command(name="shutdown").reset_cooldown(ctx)
-        await bot.get_application_command(command).invoke(ctx)
-        await ctx.respond(f'`{command}` executed.')
-
-
 @bot.slash_command(name="status",
                    description="Checks current power status of game server")
 async def _status(ctx):
     try:
-        response = requests.get(LIVENESS_URL, timeout=2)
+        response = requests.get(config.LIVENESS_URL, timeout=2)
         if response.status_code == 200:
             game = discord.Activity(
                 name="Server online", type=discord.ActivityType.playing)
@@ -219,4 +83,6 @@ async def _status(ctx):
         print("Server host is offline")
         traceback.print_exc()
 
-bot.run(DISCORD_TOKEN)
+bot.load_extension('cogs.poweractions')
+
+bot.run(config.DISCORD_TOKEN)

app/cogs/config.py

@@ -0,0 +1,36 @@
+import os
+import sys
+
+
+def env_defined(key):
+    """
+    Checks if a given env var key is defined in the OS environment
+    """
+    return key in os.environ and len(os.environ[key]) > 0
+
+# env variables are defaults, if no config file exists it'll be created.
+# If no env is set, stop the bot
+if not env_defined("DISCORD_TOKEN"):
+    print("Missing bot token from .env")
+    sys.exit()
+DISCORD_TOKEN = os.environ["DISCORD_TOKEN"]
+
+if not env_defined("WOL_URL"):
+    print("Missing wake on lan URL from .env")
+    sys.exit()
+WOL_URL = os.environ["WOL_URL"]
+
+if not env_defined("SHUTDOWN_URL"):
+    print("Missing shutdown URL from .env")
+    sys.exit()
+SHUTDOWN_URL = os.environ["SHUTDOWN_URL"]
+
+if not env_defined("REBOOT_URL"):
+    print("Missing liveness URL from .env")
+    sys.exit()
+REBOOT_URL = os.environ["REBOOT_URL"]
+
+if not env_defined("LIVENESS_URL"):
+    print("Missing liveness URL from .env")
+    sys.exit()
+LIVENESS_URL = os.environ["LIVENESS_URL"]

app/cogs/poweractions.py

@@ -0,0 +1,59 @@
+import json
+import traceback
+import requests
+import discord
+from discord.ext import commands
+import config
+
+
+class PowerActions(commands.Cog):
+    def __init__(self, bot):
+        self.bot = bot
+
+    @commands.group()
+    async def actions(self, ext):
+        pass
+
+    @actions.slash_command(name="boot", description="Boots the game server")
+    async def _boot(self, ctx):
+        try:
+            response = requests.get(config.WOL_URL, timeout=2)
+            jsonresponse = json.loads(response.content.decode())
+            if jsonresponse.get('success') is True:
+                game = discord.Activity(
+                    name="Booting...", type=discord.ActivityType.playing)
+                await self.bot.change_presence(status=discord.Status.online, activity=game)
+                await ctx.respond('Server booted!')
+        except Exception:
+            await ctx.respond('Something went wrong, have an adult check the logs')
+            traceback.print_exc()
+
+    @actions.slash_command(name="shutdown", description="Shuts down the game server")
+    async def _shutdown(self, ctx):
+        try:
+            response = requests.get(config.SHUTDOWN_URL, timeout=2)
+            if response.status_code == 200:
+                game = discord.Activity(
+                    name="Powering down...", type=discord.ActivityType.playing)
+                await self.bot.change_presence(status=discord.Status.do_not_disturb, activity=game)
+                await ctx.respond('Server shut down!')
+        except Exception:
+            await ctx.respond('Server is already offline')
+            traceback.print_exc()
+
+    @actions.slash_command(name="reboot", description="Reboots the game server")
+    async def _reboot(self, ctx):
+        try:
+            response = requests.get(config.REBOOT_URL, timeout=2)
+            if response.status_code == 200:
+                game = discord.Activity(
+                    name="Rebooting...", type=discord.ActivityType.playing)
+                await self.bot.change_presence(status=discord.Status.streaming, activity=game)
+                await ctx.respond('Server rebooting!')
+        except Exception:
+            await ctx.respond('Server is already offline')
+            traceback.print_exc()
+
+
+def setup(bot):
+    bot.add_cog(PowerActions(bot))