From b23437c4005378c462db497265ea13f706dbecdc Mon Sep 17 00:00:00 2001 From: Jareth Gomes Date: Thu, 28 Aug 2025 07:15:10 +0000 Subject: [PATCH 1/3] feat: Add member role save and restoration on member leave and join --- bot.py | 2 + src/discord/rolerestore.py | 95 ++++++++++++++++++++++++++++++++++++++ src/mongo/models.py | 10 ++++ 3 files changed, 107 insertions(+) create mode 100644 src/discord/rolerestore.py diff --git a/bot.py b/bot.py index e1076d3..f1a51d0 100644 --- a/bot.py +++ b/bot.py @@ -190,6 +190,7 @@ async def setup_hook(self) -> None: src.mongo.models.Event, src.mongo.models.Censor, src.mongo.models.Settings, + src.mongo.models.UserRoles, # TODO ], ) @@ -210,6 +211,7 @@ async def setup_hook(self) -> None: "src.discord.spam", "src.discord.reporter", "src.discord.logger", + "src.discord.rolerestore", ) for i, extension in enumerate(extensions): try: diff --git a/src/discord/rolerestore.py b/src/discord/rolerestore.py new file mode 100644 index 0000000..41af9c4 --- /dev/null +++ b/src/discord/rolerestore.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import discord +from beanie.odm.operators.update.general import Set +from discord import Member, RawMemberRemoveEvent +from discord.ext import commands + +from src.discord.globals import ( + ROLE_EM, + ROLE_LH, + ROLE_MUTED, + ROLE_QUARANTINE, + ROLE_SELFMUTE, + ROLE_STAFF, + ROLE_VIP, +) +from src.mongo.models import UserRoles + +if TYPE_CHECKING: + from bot import PiBot + +NON_PUBLIC_ROLES = [ + # Note that any role above `Bot` or `Pi-Bot` will not be assignable by the bot. This essentially + # means any Staff role is not assignable and existing staff must assign the user those roles manually. + ROLE_LH, + ROLE_EM, + ROLE_MUTED, + ROLE_SELFMUTE, + ROLE_QUARANTINE, +] + + +class RoleRestore(commands.Cog): + """ + Cog responsible for maintaining members roles when they leave and rejoin. This is in particular + to *awarded* roles and not roles that are publicly assignable. + """ + + def __init__(self, bot: PiBot): + self.bot = bot + + @commands.Cog.listener() + async def on_raw_member_remove(self, payload: RawMemberRemoveEvent): + if isinstance(payload.user, discord.User): + # figure out if we can find member here + # the payload type .user can be User | Member, but this should only be called when the member leaves a Guild?? + return + + roles_to_save = [ + role.name for role in payload.user.roles if role.name in NON_PUBLIC_ROLES + ] + + await UserRoles.find_one( + UserRoles.user_id == payload.user.id, + UserRoles.guild_id == payload.guild_id, + ).upsert( + Set({UserRoles.roles: roles_to_save}), + on_insert=UserRoles( + user_id=payload.user.id, + guild_id=payload.guild_id, + roles=roles_to_save, + ), + ) + + logging.info( + "%d roles were saved for `%s` (id: %d)", + len(roles_to_save), + payload.user.name, + payload.user.id, + ) + + @commands.Cog.listener() + async def on_member_join(self, member: Member): + user_roles = await UserRoles.find_one( + UserRoles.user_id == member.id, + UserRoles.guild_id == member.guild.id, + ) + + if user_roles: + roles_to_add = [] + for role_name in user_roles.roles: + role = discord.utils.get(member.guild.roles, name=role_name) + roles_to_add.append(role) + + await member.add_roles( + *roles_to_add, + reason="Existing user rejoined. Restoring non-public roles.", + ) + + +async def setup(bot: PiBot): + await bot.add_cog(RoleRestore(bot)) diff --git a/src/mongo/models.py b/src/mongo/models.py index 000f099..de05fc5 100644 --- a/src/mongo/models.py +++ b/src/mongo/models.py @@ -94,3 +94,13 @@ class Settings(Document): class Settings: name = "settings" use_cache = True + + +class UserRoles(Document): + user_id: int + guild_id: int + roles: list[str] + + class Settings: + name = "user_roles" + use_cache = False From a93030a5fb38245c636ace663c29bdf32c01c073 Mon Sep 17 00:00:00 2001 From: Jareth Gomes Date: Mon, 15 Sep 2025 00:38:42 +0000 Subject: [PATCH 2/3] feat: Add role sync to trigger on member update --- src/discord/rolerestore.py | 58 +++++++++++++++++++++++++------------- 1 file changed, 39 insertions(+), 19 deletions(-) diff --git a/src/discord/rolerestore.py b/src/discord/rolerestore.py index 41af9c4..0cc26d6 100644 --- a/src/discord/rolerestore.py +++ b/src/discord/rolerestore.py @@ -7,6 +7,7 @@ from beanie.odm.operators.update.general import Set from discord import Member, RawMemberRemoveEvent from discord.ext import commands +from motor.core import ClientSession from src.discord.globals import ( ROLE_EM, @@ -49,31 +50,15 @@ async def on_raw_member_remove(self, payload: RawMemberRemoveEvent): # the payload type .user can be User | Member, but this should only be called when the member leaves a Guild?? return - roles_to_save = [ - role.name for role in payload.user.roles if role.name in NON_PUBLIC_ROLES - ] - - await UserRoles.find_one( - UserRoles.user_id == payload.user.id, - UserRoles.guild_id == payload.guild_id, - ).upsert( - Set({UserRoles.roles: roles_to_save}), - on_insert=UserRoles( - user_id=payload.user.id, - guild_id=payload.guild_id, - roles=roles_to_save, - ), - ) - logging.info( - "%d roles were saved for `%s` (id: %d)", - len(roles_to_save), + "Syncing roles for `%s` (triggered by on_raw_member_remove)", payload.user.name, - payload.user.id, ) + await sync_roles(payload.user, None) @commands.Cog.listener() async def on_member_join(self, member: Member): + logging.info("Restoring roles for user `%s`", member.name) user_roles = await UserRoles.find_one( UserRoles.user_id == member.id, UserRoles.guild_id == member.guild.id, @@ -89,6 +74,41 @@ async def on_member_join(self, member: Member): *roles_to_add, reason="Existing user rejoined. Restoring non-public roles.", ) + logging.info("Finish restoring roles for `%s`", member.name) + else: + logging.info( + "RoleRestore could not find existing entry for `%s` for restoration", + member.name, + ) + + @commands.Cog.listener() + async def on_member_update(self, _: Member, after: Member): + logging.info("Updating roles for user `%s`", after.name) + # TODO: add logic to prevent unnecessary syncs + await sync_roles(after, None) + + +async def sync_roles(member: Member, session: ClientSession | None): + roles = [role.name for role in member.roles if role.name in NON_PUBLIC_ROLES] + + await UserRoles.find_one( + UserRoles.user_id == member.id, + UserRoles.guild_id == member.guild.id, + session=session, + ).upsert( + Set({UserRoles.roles: roles}), + on_insert=UserRoles(user_id=member.id, guild_id=member.guild.id, roles=roles), + session=session, + ) + + logging.info( + "%d roles were saved for `%s` (id: %d)", + len(roles), + member.name, + member.id, + ) + + return roles async def setup(bot: PiBot): From f3ac17dd566729f54bf175935e05e0bb419ae072 Mon Sep 17 00:00:00 2001 From: Jareth Gomes Date: Mon, 15 Sep 2025 00:43:56 +0000 Subject: [PATCH 3/3] feat: Add `/rolerestore sync` command to allow for manual syncing --- src/discord/rolerestore.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/src/discord/rolerestore.py b/src/discord/rolerestore.py index 0cc26d6..e28c568 100644 --- a/src/discord/rolerestore.py +++ b/src/discord/rolerestore.py @@ -5,10 +5,11 @@ import discord from beanie.odm.operators.update.general import Set -from discord import Member, RawMemberRemoveEvent +from discord import AllowedMentions, Member, RawMemberRemoveEvent, app_commands from discord.ext import commands from motor.core import ClientSession +from env import env from src.discord.globals import ( ROLE_EM, ROLE_LH, @@ -43,6 +44,37 @@ class RoleRestore(commands.Cog): def __init__(self, bot: PiBot): self.bot = bot + rolerestore_group = app_commands.Group( + name="rolerestore", + description="Commands pertaining to RoleRestore.", + guild_ids=env.slash_command_guilds, + default_permissions=discord.Permissions(manage_roles=True), + ) + + @rolerestore_group.command( + name="sync", + description="Staff command. Syncs Pi-Bot's backend with the current state of a user's non-public roles.", + ) + @app_commands.checks.has_any_role(ROLE_STAFF, ROLE_VIP) + @app_commands.describe( + member="Sync a particular user rather than the entire server.", + ) + async def sync(self, interaction: discord.Interaction, member: Member): + logging.info("Manually syncing roles for user `%s`", member.name) + await interaction.response.send_message( + f"Syncing roles for {member.mention}", + allowed_mentions=AllowedMentions.none(), + ) + roles = await sync_roles(member, None) + + roles_markdown = [f"`{discord.utils.escape_markdown(role)}`" for role in roles] + roles_str = ", ".join(roles_markdown) + await interaction.edit_original_response( + content=f"The following roles were synced: {roles_str}", + allowed_mentions=AllowedMentions.none(), + ) + return + @commands.Cog.listener() async def on_raw_member_remove(self, payload: RawMemberRemoveEvent): if isinstance(payload.user, discord.User):