Cleanup, message on reboot, global context state for reboots, banning

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

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):

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):
+        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...")
-        try:
+            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:
+        else:
-            pass
+            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):

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)
 

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

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()