Add API commands to get level/xp, dynamic command restrictions and cooldowns based on level.

brandons209 · brandons209 · commit 2becd56823f2 · 2025-05-04T22:11:04.000-04:00
diff --git a/levelup/commands/admin.py b/levelup/commands/admin.py
@@ -4,14 +4,10 @@
 from time import perf_counter
 
 import discord
+from discord.ext.commands import cooldowns
 from redbot.core import commands
 from redbot.core.i18n import Translator, cog_i18n
-from redbot.core.utils.chat_formatting import (
-    box,
-    humanize_number,
-    humanize_timedelta,
-    pagify,
-)
+from redbot.core.utils.chat_formatting import box, humanize_number, humanize_timedelta, pagify, error, warning, info
 
 from ..abc import MixinMeta
 from ..common import const, utils
@@ -29,6 +25,200 @@ async def levelset(self, ctx: commands.Context):
         """Configure LevelUp Settings"""
         pass
 
+    @levelset.group(name="bypass")
+    async def levelset_bypass(self, ctx: commands.Context):
+        """Set roles/members that bypass level requirement and cooldowns set"""
+        pass
+
+    @levelset_bypass.command(name="member")
+    async def levelset_bypass_member(self, ctx: commands.Context, member: t.Optional[discord.Member] = None):
+        """
+        Add a member to the bypass list.
+
+        Run with no arguments to see all the bypass members
+        Run with a member already in the list to remove them from the list
+        """
+        conf = self.db.get_conf(ctx.guild)
+        if member is None:
+            if not conf.cmd_bypass_member:
+                await ctx.reply(
+                    info("No members configured for bypassing command restrictions."),
+                    delete_after=30,
+                    mention_author=False,
+                )
+                return
+            text = "# Bypass Members:\n"
+            for member_id in conf.cmd_bypass_member.copy():
+                member = ctx.guild.get_member(member_id)
+                if member is None:  # member left, remove from conf
+                    conf.cmd_bypass_member.remove(member_id)
+                    continue
+                text += f"- {member.mention}\n"
+            pages = list(pagify(text))
+            await ctx.send_interactive(pages)
+            return
+        if member.id not in conf.cmd_bypass_member:
+            conf.cmd_bypass_member.append(member.id)
+        else:
+            conf.cmd_bypass_member.remove(member.id)
+            await ctx.reply(
+                info(f"Member {member.mention} removed from command bypasses."), delete_after=30, mention_author=False
+            )
+        await ctx.tick()
+
+    @levelset_bypass.command(name="role")
+    async def levelset_bypass_role(self, ctx: commands.Context, role: t.Optional[discord.Role] = None):
+        """
+        Add a role to the bypass list.
+
+        Run with no arguments to see all the bypass roles
+        Run with a role already in the list to remove it from the list
+        """
+        conf = self.db.get_conf(ctx.guild)
+        if role is None:
+            if not conf.cmd_bypass_roles:
+                await ctx.reply(
+                    info("No roles configured for bypassing command restrictions."),
+                    delete_after=30,
+                    mention_author=False,
+                )
+                return
+            text = "# Bypass Roles:\n"
+            for role_id in conf.cmd_bypass_roles.copy():
+                role = ctx.guild.get_role(role_id)
+                if role is None:  # role gone, remove from conf
+                    conf.cmd_bypass_roles.remove(role_id)
+                    continue
+                text += f"- {role.mention}\n"
+            pages = list(pagify(text))
+            await ctx.send_interactive(pages)
+            return
+        if role.id not in conf.cmd_bypass_roles:
+            conf.cmd_bypass_roles.append(role.id)
+        else:
+            conf.cmd_bypass_roles.remove(role.id)
+            await ctx.reply(
+                info(f"Role {role.mention} removed from command bypasses."),
+                delete_after=30,
+                mention_author=False,
+            )
+        await ctx.tick()
+
+    @levelset.group(name="cooldowns")
+    async def levelset_cooldowns(self, ctx: commands.Context):
+        """Manage per level command cooldowns"""
+        pass
+
+    @levelset_cooldowns.command(name="add")
+    async def levelset_cooldowns_add(self, ctx: commands.Context, level: int, cooldown: int, *, command: str):
+        """
+        Add a cooldown for a command based on level
+        Multiple cooldown levels can be set, the cooldown will be applied to members at the specified level and under
+
+        **Warning:** This will override any default cooldowns for the command
+
+        Example:
+        [p]lset cooldowns add 5 15 mycommand
+        [p]lset cooldowns add 10 5 mycommand
+        Members who are level [0, 5] will have a cooldown of 15 seconds for mycommand (including members at level 5)
+        Members who are level (5, 10] will have a cooldown of 5 seconds
+        Members above level 10 will have no cooldown
+        """
+        if self.bot.get_command(command) is None:
+            return await ctx.reply(error(f"Invalid command: `{command}`"), delete_after=30, mention_author=False)
+        conf = self.db.get_conf(ctx.guild)
+        command_cooldowns = conf.cmd_cooldowns.get(command, {})
+        command_cooldowns[level] = cooldown
+        conf.cmd_cooldowns[command] = command_cooldowns
+        self.save()
+        await ctx.tick()
+
+    @levelset_cooldowns.command(name="del")
+    async def levelset_cooldowns_del(self, ctx: commands.Context, level: int, *, command: str):
+        """Delete a cooldown for a specific command and level"""
+        if self.bot.get_command(command) is None:
+            return await ctx.reply(error(f"Invalid command: `{command}`"), delete_after=30, mention_author=False)
+        conf = self.db.get_conf(ctx.guild)
+        command_cooldowns = conf.cmd_cooldowns.get(command, {})
+        if not command_cooldowns:
+            return await ctx.reply(
+                warning(f"No cooldowns are set for `{command}`"), delete_after=30, mention_author=False
+            )
+        if level not in command_cooldowns:
+            return await ctx.reply(
+                warning(f"There is no cooldown for level {level}"), delete_after=30, mention_author=False
+            )
+        del command_cooldowns[level]
+        conf.cmd_cooldowns[command] = command_cooldowns
+        if command_cooldowns == {}:
+            del conf.cmd_cooldowns[command]
+
+        self.save()
+        await ctx.tick()
+
+    @levelset_cooldowns.command(name="list")
+    async def levelset_cooldowns_list(self, ctx: commands.Context):
+        """List cooldowns for all commands"""
+        conf = self.db.get_conf(ctx.guild)
+        cmds = conf.cmd_cooldowns
+        if not cmds:
+            await ctx.send(info("No commands configured."))
+            return
+
+        msg = f"# Cooldowns for {ctx.guild.name}\n"
+        for cmd, cooldowns in cmds.items():
+            msg += f"- `{cmd}`:\n"
+            for level, cooldown in cooldowns.items():
+                msg += f"   - Level `{level}`: `{humanize_timedelta(seconds=cooldown)}`\n"
+
+        for page in pagify(msg):
+            await ctx.send(page)
+
+    @levelset.group(name="lvlreq")
+    async def levelset_lvlreq(self, ctx: commands.Context):
+        """Manage level requirement for commands"""
+        pass
+
+    @levelset_lvlreq.command(name="add")
+    async def levelset_lvlreq_add(self, ctx: commands.Context, level: int, *, command: str):
+        """Add a level requirement to a command."""
+        if self.bot.get_command(command) is None:
+            return await ctx.reply(error(f"Invalid command: `{command}`"), delete_after=30, mention_author=False)
+        conf = self.db.get_conf(ctx.guild)
+        conf.cmd_requirements[command] = level
+        self.save()
+        await ctx.tick()
+
+    @levelset_lvlreq.command(name="del")
+    async def levelset_lvlreq_del(self, ctx: commands.Context, *, command: str):
+        """Delete a level requirement for a command."""
+        conf = self.db.get_conf(ctx.guild)
+        if command not in conf.cmd_requirements:
+            return await ctx.reply(
+                warning(f"No level requirement was set for `{command}`"),
+                delete_after=30,
+                mention_author=False,
+            )
+        del conf.cmd_requirements[command]
+        self.save()
+        await ctx.tick()
+
+    @levelset_lvlreq.command(name="list")
+    async def levelset_lvlreq_list(self, ctx: commands.Context):
+        """List all command level requirements"""
+        conf = self.db.get_conf(ctx.guild)
+        cmds = conf.cmd_requirements
+        if not cmds:
+            await ctx.send(info("No commands configured."))
+            return
+
+        msg = f"# Command Level Requirements for {ctx.guild.name}\n"
+        for cmd, level in cmds.items():
+            msg += f"- `{cmd}`: `{level}`\n"
+
+        for page in pagify(msg):
+            await ctx.send(page)
+
     @levelset.command(name="view")
     @commands.bot_has_permissions(embed_links=True)
     async def view_settings(self, ctx: commands.Context):
diff --git a/levelup/common/models.py b/levelup/common/models.py
@@ -281,6 +281,13 @@ class GuildSettings(Base):
     cooldown: int = 60  # Only gives XP every 60 seconds
     min_length: int = 0  # Minimum length of message to be considered eligible for XP gain
 
+    # command checks and cooldown based on level
+    cmd_requirements: t.Dict[str, int] = {}  # command name -> level requirement
+    # command name -> Dict[level, cooldown], levels at specified level and under will be under that cooldown
+    cmd_cooldowns: t.Dict[str, t.Dict[int, int]] = {}
+    cmd_bypass_roles: t.List[int] = []  # allow roles to bypass req level and cooldown
+    cmd_bypass_member: t.List[int] = []  # allow members to bypass req level and cooldown
+
     # Voice
     voicexp: int = 2  # XP per minute in voice
     ignore_muted: bool = True  # Ignore XP while being muted in voice
diff --git a/levelup/main.py b/levelup/main.py
@@ -131,6 +131,8 @@ async def cog_load(self) -> None:
 
     async def cog_unload(self) -> None:
         self.bot.tree.remove_command(view_profile_context)
+        self.bot.remove_before_invoke_hook(self.level_check)
+        self.bot.remove_before_invoke_hook(self.cooldown_check)
         self.stop_levelup_tasks()
 
     async def start_api(self) -> bool:
@@ -245,6 +247,10 @@ async def initialize(self) -> None:
         if voice_initialized := await self.initialize_voice_states():
             log.info(f"Initialized {voice_initialized} voice states")
 
+        # add checks
+        self.bot.before_invoke(self.level_check)
+        self.bot.before_invoke(self.cooldown_check)
+
         self.start_levelup_tasks()
         self.custom_fonts.mkdir(exist_ok=True)
         self.custom_backgrounds.mkdir(exist_ok=True)
diff --git a/levelup/shared/__init__.py b/levelup/shared/__init__.py
@@ -2,9 +2,10 @@
 from .levelups import LevelUps
 from .profile import ProfileFormatting
 from .weeklyreset import WeeklyReset
+from .checks import Checks
 
 
-class SharedFunctions(LevelUps, ProfileFormatting, WeeklyReset, metaclass=CompositeMetaClass):
+class SharedFunctions(LevelUps, ProfileFormatting, WeeklyReset, Checks, metaclass=CompositeMetaClass):
     """
     Subclass all shared metaclassed parts of the cog
 
diff --git a/levelup/shared/checks.py b/levelup/shared/checks.py
@@ -0,0 +1,92 @@
+import logging
+import bisect
+import typing as t
+
+import discord
+from redbot.core.i18n import Translator
+from redbot.core.utils.chat_formatting import warning
+from redbot.core import commands
+
+from ..abc import MixinMeta
+
+log = logging.getLogger("red.vrt.levelup.shared.checks")
+_ = Translator("LevelUp", __file__)
+
+
+class Checks(MixinMeta):
+    def __init__(self):
+        self.last_invoked: t.Dict[t.Tuple[int, str], float] = {}  # (member id, command) -> last usage time
+
+    async def level_check(self, ctx: commands.Context) -> bool:
+        if not ctx.guild:
+            return True
+        member = ctx.author
+        if isinstance(member, discord.User):
+            return True
+        conf = self.db.get_conf(member.guild)
+        command = ctx.command.qualified_name
+        if command not in conf.cmd_requirements:
+            return True
+
+        # allow help command
+        if ctx.command.name == "help" or ctx.invoked_with == "help":
+            return True
+
+        # check bypasses
+        bypass_roles = set([r.id for r in member.roles]) & set(conf.cmd_bypass_roles)
+        if member.id in conf.cmd_bypass_member or len(bypass_roles) >= 1:
+            return True
+
+        profile = conf.get_profile(member)
+        level = profile.level
+        req_level = conf.cmd_requirements[command]
+        if level < req_level:
+            await ctx.reply(
+                warning(f"You need level `{req_level}` to use `{ctx.command}`. (current level: {level})"),
+                delete_after=30,
+                mention_author=False,
+            )
+            raise commands.CheckFailure()
+        return True
+
+    async def cooldown_check(self, ctx: commands.Context) -> bool:
+        if not ctx.guild:
+            return True
+        member = ctx.author
+        if isinstance(member, discord.User):
+            return True
+        conf = self.db.get_conf(member.guild)
+        command = ctx.command.qualified_name
+        if command not in conf.cmd_cooldowns:
+            return True
+
+        # check bypasses
+        bypass_roles = set([r.id for r in member.roles]) & set(conf.cmd_bypass_roles)
+        if member.id in conf.cmd_bypass_member or len(bypass_roles) >= 1:
+            return True
+
+        cooldowns = conf.cmd_cooldowns[command]
+        levels = sorted(cooldowns.keys())
+        profile = conf.get_profile(member)
+        level = profile.level
+
+        cooldown_level = bisect.bisect_left(levels, level)
+        if cooldown_level == len(
+            levels
+        ):  # user has a higher level then all of the specified cooldowns, so no cooldown is applied
+            return True
+        cooldown = cooldowns[levels[cooldown_level]]
+
+        key = (member.id, command)
+        now = ctx.message.created_at.timestamp()
+        last = self.last_invoked.get(key, 0)
+        retry_after = last + cooldown - now
+
+        if retry_after > 0:
+            bucket_cooldown = commands.Cooldown(1, cooldown)
+            raise commands.CommandOnCooldown(bucket_cooldown, retry_after, commands.BucketType.member)
+        self.last_invoked[key] = now
+        # override any built in cooldowns for the command
+        ctx.command.reset_cooldown(ctx)
+
+        return True
diff --git a/levelup/shared/profile.py b/levelup/shared/profile.py
@@ -52,6 +52,22 @@ async def remove_xp(self, member: discord.Member, xp: int) -> int:
         self.save()
         return int(profile.xp)
 
+    async def get_xp(self, member: discord.Member) -> int:
+        """Get the XP for a member"""
+        if not isinstance(member, discord.Member):
+            raise TypeError("member must be a discord.Member")
+        conf = self.db.get_conf(member.guild)
+        profile = conf.get_profile(member)
+        return int(profile.xp)
+
+    async def get_level(self, member: discord.Member) -> int:
+        """Get the level for a member"""
+        if not isinstance(member, discord.Member):
+            raise TypeError("member must be a discord.Member")
+        conf = self.db.get_conf(member.guild)
+        profile = conf.get_profile(member)
+        return profile.level
+
     async def get_profile_background(
         self, user_id: int, profile: Profile, try_return_url: bool = False
     ) -> t.Union[bytes, str]:
@@ -108,7 +124,9 @@ async def get_banner(self, user_id: int) -> t.Optional[str]:
             return f"https://cdn.discordapp.com/banners/{user_id}/{banner_id}?size=1024"
 
     async def get_user_profile(
-        self, member: discord.Member, reraise: bool = False
+        self,
+        member: discord.Member,
+        reraise: bool = False,
     ) -> t.Union[discord.Embed, discord.File]:
         """
         Get a user's profile as an embed or file
