diff --git a/config.example.yaml b/config.example.yaml index fcbb2f4..7cf7ec1 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1,3 +1,8 @@ --- token: api key here client_id: 123456789012345678 + admins: + - list + - of + - user + - ids diff --git a/jarvis/cogs/admin.py b/jarvis/cogs/admin.py new file mode 100644 index 0000000..8c3fbfc --- /dev/null +++ b/jarvis/cogs/admin.py @@ -0,0 +1,17 @@ +from discord import User +from discord.ext import commands + + +class AdminCog(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @commands.command(name="ban") + @commands.has_permissions(administrator=True) + async def _ban(self, ctx, user: User = None, reason=None): + if not user or user == ctx.message.author: + await ctx.channel.send("You cannot ban yourself") + + +def setup(bot): + bot.add_cog(AdminCog(bot)) diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py new file mode 100644 index 0000000..47f4bf9 --- /dev/null +++ b/jarvis/cogs/owner.py @@ -0,0 +1,67 @@ +from discord.ext import commands +from jarvis.config import get_config + + +class OwnerCog(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.admins = get_config().admins + + @commands.command(name="load", hidden=True) + 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) + async def _unload_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.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) + async def _cog_reload(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.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") + + +def setup(bot): + bot.add_cog(OwnerCog(bot)) diff --git a/jarvis/cogs/util.py b/jarvis/cogs/util.py index e8d18af..156ba51 100644 --- a/jarvis/cogs/util.py +++ b/jarvis/cogs/util.py @@ -1,5 +1,6 @@ from jarvis import jarvis_self -from jarvis.utils import convert_bytesize +from jarvis.utils import convert_bytesize, build_embed +from jarvis.utils.field import Field from discord import Embed, Color from discord.ext import commands @@ -11,16 +12,19 @@ class UtilCog(commands.Cog): @commands.command(name="status", help="Retrieve J.A.R.V.I.S. status") async def _status(self, ctx): title = "J.A.R.V.I.S. Status" - description = "All systems online" - color = Color.from_rgb(152, 204, 218) - embed = Embed(title=title, description=description, color=color) + desc = "All systems online" + color = "#98CCDA" + fields = [] with jarvis_self.oneshot(): - embed.add_field(name="CPU Usage", value=jarvis_self.cpu_percent()) - embed.add_field( - name="RAM Usage", - value=convert_bytesize(jarvis_self.memory_info().rss), + fields.append(Field("CPU Usage", jarvis_self.cpu_percent())) + fields.append( + Field( + "RAM Usage", + convert_bytesize(jarvis_self.memory_info().rss), + ) ) - embed.add_field(name="PID", value=jarvis_self.pid) + fields.append(Field("PID", jarvis_self.pid)) + embed = build_embed(title, desc, fields, color) await ctx.send(embed=embed) diff --git a/jarvis/config.py b/jarvis/config.py index 85f93f0..76a438d 100644 --- a/jarvis/config.py +++ b/jarvis/config.py @@ -11,10 +11,13 @@ except ImportError: class Config: token: str client_id: str + admins: list @classmethod def from_yaml(cls, y): - instance = cls(token=y["token"], client_id=y["client_id"]) + instance = cls( + token=y["token"], client_id=y["client_id"], admins=y["admins"] + ) return instance diff --git a/jarvis/utils/__init__.py b/jarvis/utils/__init__.py index 8e6fdf7..639a04a 100644 --- a/jarvis/utils/__init__.py +++ b/jarvis/utils/__init__.py @@ -1,4 +1,5 @@ from discord.ext import commands +from discord import Embed, Color from pkgutil import iter_modules import jarvis.cogs @@ -23,3 +24,20 @@ def get_prefix(bot, message): def get_extensions(path=jarvis.cogs.__path__) -> list: return ["jarvis.cogs.{}".format(x.name) for x in iter_modules(path)] + + +def parse_color_hex(hex: str) -> Color: + hex = hex.lstrip("#") + rgb = tuple(int(hex[i : i + 2], 16) for i in (0, 2, 4)) # noqa: E203 + return Color.from_rgb(*rgb) + + +def build_embed( + title: str, desc: str, fields: list, color: str = "#FF0000", **kwargs +) -> Embed: + embed = Embed( + title=title, desc=desc, color=parse_color_hex(color), **kwargs + ) + for field in fields: + embed.add_field(**field.to_dict()) + return embed diff --git a/jarvis/utils/field.py b/jarvis/utils/field.py new file mode 100644 index 0000000..8fb683f --- /dev/null +++ b/jarvis/utils/field.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass +from typing import Any + + +@dataclass +class Field: + name: Any + value: Any + inline: bool = True + + def to_dict(self): + return {"name": self.name, "value": self.value, "inline": self.inline}