add blocking

cogs/mod.py

@@ -1,6 +1,8 @@
+from datetime import datetime
+
 import discord
 from discord.ext import commands
-from utils import Student, get_or_none, BannedStudentID
+from utils import Student, get_or_none, BannedStudentID, owner_or_admin, JimmyBans
 
 
 class Mod(commands.Cog):
@@ -51,6 +53,36 @@ class Mod(commands.Cog):
             else:
                 return await ctx.respond(f"\N{white heavy check mark} Unbanned {student_id}. Unbanned {user_id}.")
 
+    @commands.slash_command(name="block")
+    @owner_or_admin()
+    async def block_user(self, ctx: discord.ApplicationContext, user: discord.Member, reason: str, *, until: str):
+        """Blocks a user from using the bot."""
+        await ctx.defer()
+        try:
+            hour, minute = map(int, until.split(":"))
+        except ValueError:
+            return await ctx.respond("\N{cross mark} Invalid time format. Use HH:MM.")
+        end = datetime.utcnow().replace(hour=hour, minute=minute)
+
+        # get an entry for the user's ID, and if it doesn't exist, create it. Otherwise, alert the user.
+        entry = await get_or_none(JimmyBans, user_id=user.id)
+        if entry is None:
+            await JimmyBans.objects.create(user_id=user.id, reason=reason, until=end.timestamp())
+        else:
+            return await ctx.respond("\N{cross mark} That user is already blocked.")
+        await ctx.respond(f"\N{white heavy check mark} Blocked {user.mention} until {discord.utils.format_dt(end)}.")
+
+    @commands.slash_command(name="unblock")
+    @owner_or_admin()
+    async def unblock_user(self, ctx: discord.ApplicationContext, user: discord.Member):
+        """Unblocks a user from using the bot."""
+        await ctx.defer()
+        entry = await get_or_none(JimmyBans, user_id=user.id)
+        if entry is None:
+            return await ctx.respond("\N{cross mark} That user isn't blocked.")
+        await entry.delete()
+        await ctx.respond(f"\N{white heavy check mark} Unblocked {user.mention}.")
+
 
 def setup(bot):
     bot.add_cog(Mod(bot))

main.py

@@ -2,7 +2,8 @@ import discord
 from discord.ext import commands
 from asyncio import Lock
 import config
-from utils import registry, console
+from datetime import datetime, timezone
+from utils import registry, console, get_or_none, JimmyBans
 
 
 intents = discord.Intents.default()
@@ -81,6 +82,27 @@ async def ping(ctx: discord.ApplicationContext):
     return await ctx.respond(f"\N{white heavy check mark} Pong! `{gateway}ms`.")
 
 
+@bot.check_once
+async def check_not_banned(ctx: discord.ApplicationContext | commands.Context):
+    if await bot.is_owner(ctx.author):
+        return True
+    user = ctx.author
+    ban: JimmyBans = await get_or_none(JimmyBans, user_id=user.id)
+    if ban:
+        dt = datetime.fromtimestamp(ban.until, timezone.utc)
+        if dt < discord.utils.utcnow():
+            await ban.delete()
+        else:
+            reply = ctx.reply if isinstance(ctx, commands.Context) else ctx.respond
+            try:
+                await reply(content=f":x: You can use commands {discord.utils.format_dt(dt, 'R')}")
+            except discord.HTTPException:
+                pass
+            finally:
+                return False
+    return True
+
+
 if __name__ == "__main__":
-    print("Starting...")
+    console.log("Starting...")
     bot.run(config.token)

utils/db.py

@@ -30,6 +30,7 @@ __all__ = [
     "Assignments",
     "Tutors",
     "UptimeEntry",
+    "JimmyBans"
 ]
 
 T = TypeVar("T")
@@ -168,3 +169,21 @@ class UptimeEntry(orm.Model):
         timestamp: float
         response_time: int | None
         notes: str | None
+
+
+class JimmyBans(orm.Model):
+    tablename = "jimmy_bans"
+    registry = registry
+    fields = {
+        "entry_id": orm.UUID(primary_key=True, default=uuid.uuid4),
+        "user_id": orm.BigInteger(),
+        "reason": orm.Text(allow_null=True, default=None),
+        "timestamp": orm.Float(default=lambda: datetime.datetime.utcnow().timestamp()),
+        "until": orm.Float(),
+    }
+    if TYPE_CHECKING:
+        entry_id: uuid.UUID
+        user_id: int
+        reason: str | None
+        timestamp: float
+        until: float | None