diff --git a/jarvis/__init__.py b/jarvis/__init__.py index cb45306..f9cf227 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -13,18 +13,34 @@ if asyncio.get_event_loop().is_closed(): intents = Intents.default() intents.members = True +restart_ctx = None jarvis = commands.Bot(command_prefix=utils.get_prefix, intents=intents) jarvis_self = Process() +__version__ = "0.1.0" @jarvis.event async def on_ready(): + global restart_ctx print(" Logged in as {0.user}".format(jarvis)) print(" Connected to {} guild(s)".format(len(jarvis.guilds))) with jarvis_self.oneshot(): print(f" Current PID: {jarvis_self.pid}") Path(f"jarvis.{jarvis_self.pid}.pid").touch() + if restart_ctx: + channel = None + if "guild" in restart_ctx: + guild = find(lambda x: x.id == restart_ctx["guild"], jarvis.guilds) + if guild: + channel = find( + lambda x: x.id == restart_ctx["channel"], guild.channels + ) + elif "user" in restart_ctx: + channel = jarvis.get_user(restart_ctx["user"]) + if channel: + await channel.send("Core systems restarted and back online.") + restart_ctx = None @jarvis.event @@ -43,7 +59,10 @@ async def on_guild_join(guild): await general.send("Systems are now fully operational") -def run(): +def run(ctx=None): + global restart_ctx + if ctx: + restart_ctx = ctx for extension in utils.get_extensions(): jarvis.load_extension(extension) config = get_config() @@ -52,3 +71,5 @@ def run(): + "{}&permissions=8&scope=bot".format(config.client_id) ) jarvis.run(config.token, bot=True, reconnect=True) + if restart_ctx: + return restart_ctx diff --git a/jarvis/cogs/admin.py b/jarvis/cogs/admin.py index 8c3fbfc..cba225b 100644 --- a/jarvis/cogs/admin.py +++ b/jarvis/cogs/admin.py @@ -10,7 +10,24 @@ class AdminCog(commands.Cog): @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") + await ctx.send("You cannot ban yourself.") + if not reason: + reason = ( + "Mr. Stark is displeased with your presence. Please leave." + ) + await user.send( + f"You have been banned from Stark Industries. Reason:\n{reason}" + ) + await ctx.guild.ban(user, reason=reason) + await ctx.send( + f"{user.name} has been banned from Stark Industries." + + f"Reason:\n{reason}" + ) + + @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): diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py index fe3ec3e..d5b06eb 100644 --- a/jarvis/cogs/owner.py +++ b/jarvis/cogs/owner.py @@ -1,6 +1,8 @@ +import jarvis +import discord from discord.ext import commands from jarvis.config import get_config -from jarvis.utils import get_prefix +from jarvis.utils import update class OwnerCog(commands.Cog): @@ -28,6 +30,9 @@ class OwnerCog(commands.Cog): @commands.command(name="unload", hidden=True) 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 @@ -46,6 +51,9 @@ class OwnerCog(commands.Cog): @commands.command(name="reload", hidden=True) 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 @@ -73,11 +81,57 @@ class OwnerCog(commands.Cog): @_system.command(name="restart", hidden=True) async def _restart(self, ctx): - await ctx.send("Restarting core systems...") - try: + 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): + print("IsDM") + 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() - except RuntimeError: - pass + else: + await ctx.send("I'm afraid I can't let you do that") + + @_system.command(name="update", hidden=True) + 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") def setup(bot): diff --git a/jarvis/cogs/util.py b/jarvis/cogs/util.py index 156ba51..2a3ff96 100644 --- a/jarvis/cogs/util.py +++ b/jarvis/cogs/util.py @@ -1,7 +1,7 @@ +import jarvis from jarvis import jarvis_self from jarvis.utils import convert_bytesize, build_embed from jarvis.utils.field import Field -from discord import Embed, Color from discord.ext import commands @@ -24,6 +24,7 @@ class UtilCog(commands.Cog): ) ) fields.append(Field("PID", jarvis_self.pid)) + fields.append(Field("Version", jarvis.__version__, False)) embed = build_embed(title, desc, fields, color) await ctx.send(embed=embed) diff --git a/jarvis/utils/__init__.py b/jarvis/utils/__init__.py index 639a04a..4fb9dfa 100644 --- a/jarvis/utils/__init__.py +++ b/jarvis/utils/__init__.py @@ -2,6 +2,7 @@ from discord.ext import commands from discord import Embed, Color from pkgutil import iter_modules import jarvis.cogs +import git def convert_bytesize(bytes: int) -> str: @@ -41,3 +42,22 @@ def build_embed( for field in fields: embed.add_field(**field.to_dict()) return embed + + +def update(): + repo = git.Repo(".") + dirty = repo.is_dirty() + current_hash = repo.head.object.hexsha + origin = repo.remotes.origin + origin.fetch() + if current_hash != origin.refs["main"].object.hexsha: + if dirty: + return 2 + origin.pull() + return 0 + return 1 + + +def get_repo_hash(): + repo = git.Repo(".") + return repo.head.object.hexsha diff --git a/run.py b/run.py index 72f9a22..0d042f9 100644 --- a/run.py +++ b/run.py @@ -8,9 +8,10 @@ from time import sleep def run(): + ctx = None while True: ireload(jarvis) - jarvis.run() + ctx = jarvis.run(ctx) def restart(): @@ -78,7 +79,7 @@ Command List: elif status == 1: print(" No core updates available.") elif status == 2: - print(" Core updates available, but not applied.") + print(" Core system update available, but core is dirty.") if cmd.lower() in ["r", "reload"]: print(" Reloading core systems...") restart()