jarvis-bot/jarvis/cogs/owner.py

import discord
import jarvis
from discord import User
from discord.ext import commands
from jarvis.config import get_config, reload_config
from jarvis.utils import update, db
from jarvis.utils.permissions import 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.config = get_config()
        self.admins = self.config.admins
        self.db = db.DBManager(self.config.mongo)

    @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: `system <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.group(name="admin", hidden=True, pass_context=True)
    @commands.is_owner()
    async def _admin(self, ctx):
        if ctx.invoked_subcommand is None:
            await ctx.send(
                "Usage: `admin <subcommand>`\n"
                + "Subcommands: `add`, `remove`"
            )

    @_admin.command(name="add", hidden=True)
    @commands.is_owner()
    async def _add(self, ctx, user: User):
        if user.id in self.admins:
            await ctx.send(f"{user.mention} is already an admin.")
            return
        self.admins.append(user.id)
        self.db.mongo.jarvis.config.update_one(
            {"key": "admins"}, {"$set": {"value": self.admins}}
        )
        reload_config()
        await ctx.send(
            f"{user.mention} is now an admin. Use this power carefully."
        )

    @_admin.command(name="remove", hidden=True)
    @commands.is_owner()
    async def _remove(self, ctx, user: User):
        if user.id not in self.admins:
            await ctx.send(f"{user.mention} is not an admin.")
            return
        self.admins.remove(user.id)
        self.db.mongo.jarvis.config.update_one(
            {"key": "admins"}, {"$set": {"value": self.admins}}
        )
        reload_config()
        await ctx.send(f"{user.mention} is no longer an admin.")


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