jarvis-bot/jarvis/cogs/owner.py

import jarvis
import discord
from discord.ext import commands
from jarvis.config import get_config, reload_config
from jarvis.utils import update, user_is_bot_admin


class OwnerCog(commands.Cog):
    """
    J.A.R.V.I.S. management cog

    Used by admins to control core J.A.R.V.I.S. systems
    """

    def __init__(self, bot):
        self.bot = bot
        self.admins = get_config().admins

    @commands.command(name="load", hidden=True)
    @user_is_bot_admin()
    async def _load_cog(self, ctx, *, cog: str):
        info = await self.bot.application_info()
        if (
            ctx.message.author == info.owner
            or ctx.message.author.id in self.admins
        ):
            try:
                self.bot.load_extension(cog)
            except Exception as e:
                await ctx.send(
                    f"Failed to load new cog `{cog}`: {type(e).name} - {e}"
                )
            else:
                await ctx.send(f"Successfully loaded new cog `{cog}`")
        else:
            await ctx.send("I'm afraid I can't let you do that")

    @commands.command(name="unload", hidden=True)
    @user_is_bot_admin()
    async def _unload_cog(self, ctx, *, cog: str):
        if cog == "jarvis.cogs.owner":
            await ctx.send("Cannot unload `owner` cog")
            return
        info = await self.bot.application_info()
        if (
            ctx.message.author == info.owner
            or ctx.message.author.id in self.admins
        ):
            try:
                self.bot.unload_extension(cog)
            except Exception as e:
                await ctx.send(
                    f"Failed to unload cog `{cog}` {type(e).__name__} - {e}"
                )
            else:
                await ctx.send(f"Successfully unloaded cog `{cog}`")
        else:
            await ctx.send("I'm afraid I can't let you do that")

    @commands.command(name="reload", hidden=True)
    @user_is_bot_admin()
    async def _cog_reload(self, ctx, *, cog: str):
        if cog == "jarvis.cogs.owner":
            await ctx.send("Cannot reload `owner` cog")
            return
        info = await self.bot.application_info()
        if (
            ctx.message.author == info.owner
            or ctx.message.author.id in self.admins
        ):
            try:
                self.bot.unload_extension(cog)
                self.bot.load_extension(cog)
            except Exception as e:
                await ctx.send(
                    f"Failed to reload cog `{cog}` {type(e).__name__} - {e}"
                )
            else:
                await ctx.send(f"Successfully reloaded cog `{cog}`")
        else:
            await ctx.send("I'm afraid I can't let you do that")

    @commands.group(name="system", hidden=True, pass_context=True)
    @user_is_bot_admin()
    async def _system(self, ctx):
        if ctx.invoked_subcommand is None:
            await ctx.send(
                f"Usage: `{ctx.message.content} <subcommand>`\n"
                + "Subcommands: `restart`, `update`"
            )

    @_system.command(name="restart", hidden=True)
    @user_is_bot_admin()
    async def _restart(self, ctx):
        info = await self.bot.application_info()
        if (
            ctx.message.author == info.owner
            or ctx.message.author.id in self.admins
        ):
            await ctx.send("Restarting core systems...")
            if isinstance(ctx.channel, discord.channel.DMChannel):
                jarvis.restart_ctx = {
                    "user": ctx.message.author.id,
                    "channel": ctx.channel.id,
                }
            else:
                jarvis.restart_ctx = {
                    "guild": ctx.message.guild.id,
                    "channel": ctx.channel.id,
                }
            await self.bot.close()
        else:
            await ctx.send("I'm afraid I can't let you do that")

    @_system.command(name="update", hidden=True)
    @user_is_bot_admin()
    async def _update(self, ctx):
        info = await self.bot.application_info()
        if (
            ctx.message.author == info.owner
            or ctx.message.author.id in self.admins
        ):
            await ctx.send("Updating core systems...")
            status = update()
            if status == 0:
                await ctx.send("Core systems updated. Restarting...")
                if isinstance(ctx.channel, discord.channel.DMChannel):
                    jarvis.restart_ctx = {
                        "user": ctx.message.author.id,
                        "channel": ctx.channel.id,
                    }
                else:
                    jarvis.restart_ctx = {
                        "guild": ctx.message.guild.id,
                        "channel": ctx.channel.id,
                    }
                await self.bot.close()
            elif status == 1:
                await ctx.send("Core systems already up to date.")
            elif status == 2:
                await ctx.send(
                    "Core system update available, but core is dirty."
                )
        else:
            await ctx.send("I'm afraid I can't let you do that")

    @_system.command(name="refresh", hidden=True)
    @user_is_bot_admin()
    async def _refresh(self, ctx):
        reload_config()
        await ctx.send("System refreshed")

    @commands.Cog.listener()
    async def on_command_error(self, ctx, error):
        if isinstance(error, commands.errors.MissingPermissions):
            await ctx.send("I'm afraid I can't let you do that.")


def setup(bot):
    bot.add_cog(OwnerCog(bot))