From 05237dbe7b4502fadeda17c123413e9c817e5ed2 Mon Sep 17 00:00:00 2001 From: flifloo Date: Thu, 5 Nov 2020 13:56:00 +0100 Subject: [PATCH] Add list, enable and disable extension commands --- config_exemple.json | 1 - db/Extension.py | 26 +++++++++++ db/__init__.py | 1 + extensions/extension.py | 96 ++++++++++++++++++++++++++++++++++++++--- 4 files changed, 117 insertions(+), 7 deletions(-) create mode 100644 db/Extension.py diff --git a/config_exemple.json b/config_exemple.json index 8fc5dd5..a6ff4dd 100644 --- a/config_exemple.json +++ b/config_exemple.json @@ -1,5 +1,4 @@ {"prefix": "!", "token": "GOOD_BOT_TOKEN", - "admin_id": 1234567890, "db": "postgresql://usr:pass@localhost:5432/sqlalchemy" } \ No newline at end of file diff --git a/db/Extension.py b/db/Extension.py new file mode 100644 index 0000000..f3fc96b --- /dev/null +++ b/db/Extension.py @@ -0,0 +1,26 @@ +from db import Base +from sqlalchemy import Column, String, Boolean, ForeignKey, BigInteger +from sqlalchemy.orm import relationship + + +class Extension(Base): + __tablename__ = "extension" + name = Column(String, primary_key=True) + default_state = Column(Boolean, nullable=False, default=True) + extension_state = relationship("ExtensionState", backref="extension") + + def __init__(self, name: int, default_state: bool = True): + self.name = name + self.default_state = default_state + + +class ExtensionState(Base): + __tablename__ = "extension_state" + extension_name = Column(String, ForeignKey("extension.name"), primary_key=True) + guild_id = Column(BigInteger, nullable=False, primary_key=True) + state = Column(Boolean, nullable=False, default=True) + + def __init__(self, extension_name: str, guild_id: int, state: bool = True): + self.extension_name = extension_name + self.guild_id = guild_id + self.state = state diff --git a/db/__init__.py b/db/__init__.py index 3d2878c..1741487 100644 --- a/db/__init__.py +++ b/db/__init__.py @@ -15,4 +15,5 @@ from db.WarnAction import WarnAction from db.InviteRole import InviteRole from db.Tomuss import Tomuss from db.PCP import PCP +from db.Extension import Extension, ExtensionState Base.metadata.create_all(engine) diff --git a/extensions/extension.py b/extensions/extension.py index 42be296..afc97ac 100644 --- a/extensions/extension.py +++ b/extensions/extension.py @@ -1,7 +1,10 @@ from traceback import format_exc from discord.ext import commands -from discord import Embed +from discord import Embed, Guild +from discord.ext.commands import MissingPermissions, BadArgument + +import db from administrator.logger import logger @@ -17,7 +20,6 @@ class Extension(commands.Cog): return "Manage bot's extensions" @commands.group("extension", pass_context=True) - @commands.is_owner() async def extension(self, ctx: commands.Context): if ctx.invoked_subcommand is None: await ctx.invoke(self.extension_help) @@ -25,20 +27,69 @@ class Extension(commands.Cog): @extension.group("help", pass_context=True) async def extension_help(self, ctx: commands.Context): embed = Embed(title="Extension help") - embed.add_field(name="extension list", value="List all loaded extensions", inline=False) - embed.add_field(name="extension load ", value="Load an extension", inline=False) - embed.add_field(name="extension unload ", value="Unload an extension", inline=False) - embed.add_field(name="extension reload ", value="Reload an extension", inline=False) + if await self.extension_list.can_run(ctx): + embed.add_field(name="extension list", value="List all enabled extensions", inline=False) + if await self.extension_enable.can_run(ctx): + embed.add_field(name="extension enable", value="Enable an extensions", inline=False) + if await self.extension_disable.can_run(ctx): + embed.add_field(name="extension disable", value="Disable an extensions", inline=False) + if await self.extension_load.can_run(ctx): + embed.add_field(name="extension loaded", value="List all loaded extensions", inline=False) + if await self.extension_load.can_run(ctx): + embed.add_field(name="extension load ", value="Load an extension", inline=False) + if await self.extension_unload.can_run(ctx): + embed.add_field(name="extension unload ", value="Unload an extension", inline=False) + if await self.extension_reload.can_run(ctx): + embed.add_field(name="extension reload ", value="Reload an extension", inline=False) + if not embed.fields: + raise MissingPermissions(None) await ctx.send(embed=embed) @extension.group("list", pass_context=True) + @commands.has_guild_permissions(administrator=True) async def extension_list(self, ctx: commands.Context): + s = db.Session() embed = Embed(title="Extensions list") + for es in s.query(db.ExtensionState).filter(db.ExtensionState.guild_id == ctx.guild.id): + embed.add_field(name=es.extension_name, value="Enable" if es.state else "Disable") + await ctx.send(embed=embed) + + @extension.group("enable", pass_context=True) + @commands.has_guild_permissions(administrator=True) + async def extension_enable(self, ctx: commands.Context, name: str): + s = db.Session() + es = s.query(db.ExtensionState).get((name, ctx.guild.id)) + if not es or es.state: + raise BadArgument() + es.state = True + s.add(es) + s.commit() + s.close() + await ctx.message.add_reaction("\U0001f44d") + + @extension.group("disable", pass_context=True) + @commands.has_guild_permissions(administrator=True) + async def extension_disable(self, ctx: commands.Context, name: str): + s = db.Session() + es = s.query(db.ExtensionState).get((name, ctx.guild.id)) + if not es or not es.state: + raise BadArgument() + es.state = False + s.add(es) + s.commit() + s.close() + await ctx.message.add_reaction("\U0001f44d") + + @extension.group("loaded", pass_context=True) + @commands.is_owner() + async def extension_loaded(self, ctx: commands.Context): + embed = Embed(title="Extensions loaded") for extension in self.bot.extensions: embed.add_field(name=extension, value="Loaded", inline=False) await ctx.send(embed=embed) @extension.group("load", pass_context=True) + @commands.is_owner() async def extension_load(self, ctx: commands.Context, name: str): try: self.bot.load_extension(name) @@ -49,6 +100,7 @@ class Extension(commands.Cog): await ctx.message.add_reaction("\U0001f44d") @extension.group("unload", pass_context=True) + @commands.is_owner() async def extension_unload(self, ctx: commands.Context, name: str): try: self.bot.unload_extension(name) @@ -59,6 +111,7 @@ class Extension(commands.Cog): await ctx.message.add_reaction("\U0001f44d") @extension.group("reload", pass_context=True) + @commands.is_owner() async def extension_reload(self, ctx: commands.Context, name: str): try: self.bot.unload_extension(name) @@ -69,6 +122,37 @@ class Extension(commands.Cog): else: await ctx.message.add_reaction("\U0001f44d") + @commands.Cog.listener() + async def on_ready(self): + s = db.Session() + for guild in self.bot.guilds: + for extension in filter(lambda x: x not in ["Extension", "Help"], self.bot.cogs): + e = s.query(db.Extension).get(extension) + if not e: + s.add(db.Extension(extension)) + s.commit() + es = s.query(db.ExtensionState).get((extension, guild.id)) + if not es: + s.add(db.ExtensionState(extension, guild.id)) + s.commit() + s.close() + + @commands.Cog.listener() + async def on_guild_join(self, guild: Guild): + s = db.Session() + for extension in s.query(db.Extension).all(): + s.add(db.ExtensionState(extension.name, guild.id)) + s.commit() + s.close() + + @commands.Cog.listener() + async def on_guild_remove(self, guild: Guild): + s = db.Session() + for es in s.query(db.ExtensionState).filter(db.ExtensionState.guild_id == guild.id): + s.delete(es) + s.commit() + s.close() + def setup(bot): logger.info(f"Loading...")