diff --git a/cogs/mod.py b/cogs/mod.py index 6d409a1..c294adb 100644 --- a/cogs/mod.py +++ b/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)) diff --git a/main.py b/main.py index 6a8d5d8..0b8765f 100644 --- a/main.py +++ b/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) diff --git a/utils/db.py b/utils/db.py index 6786934..89b93a7 100644 --- a/utils/db.py +++ b/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