Merge pull request #789 from allthingslinux/statusroles

feat(status-roles): add cog to assign roles based on user status

Author: electron271

config/settings.yml.example

@@ -49,6 +49,11 @@ TEMPVC_CATEGORY_ID: 123456789012345679
 # Set this to the channel ID where you want the temporary voice channels to be created.
 TEMPVC_CHANNEL_ID: 123456789012345679
 
+# This will automatically give people with a status regex a role.
+STATUS_ROLES:
+  #- server_id: 123456789012345679
+  #  status_regex: ".*"
+  #  role_id: 123456789012345679
 
 SNIPPETS:
   LIMIT_TO_ROLE_IDS: false # Only allow users with the specified role IDs to use the snippet command

tux/cogs/services/status_roles.py

@@ -0,0 +1,113 @@
+import re
+
+import discord
+from discord.ext import commands
+from loguru import logger
+
+from tux.utils.config import CONFIG
+
+
+class StatusRoles(commands.Cog):
+    """Assign roles to users based on their status."""
+
+    def __init__(self, bot: commands.Bot):
+        self.bot = bot
+        self.status_roles = CONFIG.STATUS_ROLES
+        logger.info("StatusRoles cog initialized with %d role configurations", len(self.status_roles))
+
+    @commands.Cog.listener()
+    async def on_ready(self):
+        """Check all users' statuses when the bot starts up."""
+        logger.info("StatusRoles cog ready, checking all users' statuses")
+        for guild in self.bot.guilds:
+            for member in guild.members:
+                await self.check_and_update_roles(member)
+
+    @commands.Cog.listener()
+    async def on_presence_update(self, before: discord.Member, after: discord.Member):
+        """Event triggered when a user's presence changes."""
+        logger.debug(f"Presence update for {after.display_name}: {before.status} -> {after.status}")
+        # Only process if the custom status changed
+        before_status = self.get_custom_status(before)
+        after_status = self.get_custom_status(after)
+
+        if before_status != after_status or self.has_activity_changed(before, after):
+            logger.debug(f"Status change detected for {after.display_name}: '{before_status}' -> '{after_status}'")
+            await self.check_and_update_roles(after)
+
+    def has_activity_changed(self, before: discord.Member, after: discord.Member) -> bool:
+        """Check if there was a relevant change in activities."""
+        before_has_custom = (
+            any(isinstance(a, discord.CustomActivity) for a in before.activities) if before.activities else False
+        )
+        after_has_custom = (
+            any(isinstance(a, discord.CustomActivity) for a in after.activities) if after.activities else False
+        )
+        return before_has_custom != after_has_custom
+
+    def get_custom_status(self, member: discord.Member) -> str | None:
+        """Extract the custom status text from a member's activities."""
+        if not member.activities:
+            return None
+
+        return next(
+            (
+                activity.name
+                for activity in member.activities
+                if isinstance(activity, discord.CustomActivity) and activity.name
+            ),
+            None,
+        )
+
+    async def check_and_update_roles(self, member: discord.Member):
+        """Check a member's status against configured patterns and update roles accordingly."""
+        if member.bot:
+            return
+
+        status_text = self.get_custom_status(member)
+        if status_text is None:
+            status_text = ""  # Use empty string for regex matching if no status
+
+        for config in self.status_roles:
+            # Skip if the config is for a different server
+            if int(config.get("server_id", 0)) != member.guild.id:
+                continue
+
+            role_id = int(config.get("role_id", 0))
+            pattern = str(config.get("status_regex", ".*"))
+
+            role = member.guild.get_role(role_id)
+            if not role:
+                logger.warning(f"Role {role_id} configured in STATUS_ROLES not found in guild {member.guild.name}")
+                continue
+
+            try:
+                matches = bool(re.search(pattern, status_text, re.IGNORECASE))
+
+                has_role = role in member.roles
+
+                if matches and not has_role:
+                    # Add role if status matches and member doesn't have the role
+                    logger.info(
+                        f"Adding role {role.name} to {member.display_name} (status: '{status_text}' matched '{pattern}')",
+                    )
+                    await member.add_roles(role, reason="Status role automation")
+
+                elif not matches and has_role:
+                    # Remove role if status doesn't match and member has the role
+                    logger.info(f"Removing role {role.name} from {member.display_name} (status no longer matches)")
+                    await member.remove_roles(role, reason="Status role automation")
+
+            except re.error:
+                logger.exception(f"Invalid regex pattern '{pattern}' in STATUS_ROLES config")
+            except discord.Forbidden:
+                logger.exception(
+                    f"Bot lacks permission to modify roles for {member.display_name} in {member.guild.name}",
+                )
+            except Exception:
+                logger.exception(f"Error updating roles for {member.display_name}")
+
+
+async def setup(bot: commands.Bot):
+    await bot.add_cog(StatusRoles(bot))
+    logger.info("Loaded StatusRoles cog")

tux/utils/config.py

@@ -55,6 +55,9 @@ class Config:
     ACTIVITIES: Final[str] = config["BOT_INFO"]["ACTIVITIES"]
     HIDE_BOT_OWNER: Final[bool] = config["BOT_INFO"]["HIDE_BOT_OWNER"]
 
+    # Status Roles
+    STATUS_ROLES: Final[list[dict[str, int]]] = config["STATUS_ROLES"]
+
     # Debug env
     DEBUG: Final[bool] = bool(os.getenv("DEBUG", "True"))