diff --git a/roleutils/autorole.py b/roleutils/autorole.py index cc8462d..08bb171 100644 --- a/roleutils/autorole.py +++ b/roleutils/autorole.py @@ -26,10 +26,11 @@ import asyncio import logging from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Final, List, NoReturn, Optional, Tuple, Union +from typing import Any, Dict, Final, List, NoReturn, Optional, Tuple, Union, Literal import discord from redbot.core import commands +from redbot.core.utils.mod import get_audit_reason from redbot.core.modlog import Case, create_case from redbot.core.utils.antispam import AntiSpam @@ -109,6 +110,62 @@ async def _create_case( case = None return case + async def _give_sticky_role( + self, + member: discord.Member, + role: Union[discord.Role, StrictRole], + reason: Optional[str] = None, + ) -> None: + if not member.guild.get_member(member.id): + return + guild: discord.Guild = member.guild + if not guild.me.guild_permissions.manage_roles: + return + await member.add_roles(role, reason=reason) + + async def _remove_sticky_role( + self, + member: discord.Member, + role: Union[discord.Role, StrictRole], + reason: Optional[str] = None, + ) -> None: + if not member.guild.get_member(member.id): + return + guild: discord.Guild = member.guild + if not guild.me.guild_permissions.manage_roles: + return + await member.remove_roles(role, reason=reason) + + async def _sticky_join(self, member: discord.Member) -> None: + guild: discord.Guild = member.guild + if await self.bot.cog_disabled_in_guild(self, guild): + return + to_reapply: List[int] = await self.config.member(member).sticky_roles() + if not to_reapply: + return + to_add: List[discord.Role] = [] + for role_id in to_reapply: + role: Optional[discord.Role] = guild.get_role(int(role_id)) + if role and role.position < guild.me.top_role.position: + to_add.append(role) + if to_add: + await member.add_roles(*to_add, reason="Applied configured sticky roles.") + + async def _sticky_remove(self, member: discord.Member) -> None: + guild: discord.Guild = member.guild + if await self.bot.cog_disabled_in_guild(self, guild): + return + to_reapply: List[int] = await self.config.member(member).sticky_roles() + reapply: bool = False + for role in member.roles: + if not await self.config.role(role).sticky(): + continue + if role.id not in to_reapply: + to_reapply.append(role.id) + reapply: bool = True + if reapply: + await self.config.member(member).sticky_roles.set(to_reapply) + async def _handle_member_join(self, member: discord.Member): role_ids: List[int] = [] settings: Dict[str, Any] = await self.config.guild(member.guild).autoroles.all() @@ -263,8 +320,101 @@ async def bots_remove(self, ctx: commands.Context, *, role: FuzzyRole): ) await ctx.send(f"Removed {role.name} ({role.id}) from the humans autorole list.") + @_autorole.group(name="sticky", aliases=["stickyrole"]) + async def _sticky(self, _: commands.GuildContext): + """""" + + @_sticky.command(name="set", aliases=["role"]) + async def _autorole_sticky_set( + self, + ctx: commands.GuildContext, + add_or_remove: Literal["add", "remove"], + *, + role: StrictRole, + ): + """""" + if add_or_remove.lower() == "add": + if await self.config.role(role).sticky(): + await ctx.send( + "{} is already a sticky role.".format(role.name), + ) + return + await self.config.role(role).sticky.set(True) + await ctx.send("Successfully configured {} as a sticky role.".format(role.name)) + elif add_or_remove.lower() == "remove": + if not await self.config.role(role).sticky(): + await ctx.send( + "{} is not a sticky role.".format(role.name), + ) + return + await self.config.role(role).sticky.set(False) + await ctx.send("Successfully removed {} from sticky roles.".format(role.name)) + else: + await ctx.send_help(ctx.command) + + @_sticky.command(name="add") + async def _autorole_sticky_add( + self, + ctx: commands.GuildContext, + users: commands.Greedy[discord.Member], + *, + role: StrictRole, + ): + """""" + failed: List[str] = [] + for user in users: + async with self.config.member(user).sticky_roles() as settings: + if role.id not in settings: + settings.append(role.id) + try: + await self._give_sticky_role( + user, role, reason=get_audit_reason(ctx.author, "Sticky role applied.") + ) + except discord.HTTPException: + failed.append( + "There was an error applying the sticky role to {} ({}).\n".format( + user.display_name, user.id + ) + ) + await ctx.send( + "{} is configured as a sticky role for {} users.".format(role.name, len(users)) + ) + if failed: + await ctx.send("".join([f for f in failed])) + + @_sticky.command(name="remove") + async def _autorole_sticky_remove( + self, + ctx: commands.GuildContext, + users: commands.Greedy[discord.Member], + *, + role: StrictRole, + ): + """""" + failed: List[str] = [] + for user in users: + async with self.config.member(user).sticky_roles() as settings: + if role.id in settings: + settings.remove(role.id) + try: + await self._remove_sticky_role( + user, role, reason=get_audit_reason(ctx.author, "Sticky role removed.") + ) + except discord.HTTPException: + failed.append( + "There was an error applying the sticky role to {} ({}).\n".format( + user.display_name, user.id + ) + ) + await ctx.send("Removed the sticky role {} from {} users.".format(role.name, len(users))) + if failed: + await ctx.send("".join([f for f in failed])) + @commands.Cog.listener() async def on_member_join(self, member: discord.Member): + if await self.bot.cog_disabled_in_guild(self, member.guild): + return + await self._sticky_join(member) if member.pending: return await self.queue.put(member) @@ -273,3 +423,9 @@ async def on_member_join(self, member: discord.Member): async def on_member_update(self, before: discord.Member, after: discord.Member): if before.pending and not after.pending: await self.queue.put(after) + + @commands.Cog.listener() + async def on_member_remove(self, member: discord.Member): + if await self.bot.cog_disabled_in_guild(self, member.guild): + return + await self._sticky_remove(member) diff --git a/roleutils/roleutils.py b/roleutils/roleutils.py index aadbadf..a498651 100644 --- a/roleutils/roleutils.py +++ b/roleutils/roleutils.py @@ -93,7 +93,15 @@ def __init__(self, bot: Red, *_args: Any) -> None: }, }, } + default_role: Dict[str, bool] = { + "sticky": False, + } + default_member: Dict[str, List[int]] = { + "sticky_roles": [], + } self.config.register_guild(**default_guild) + self.config.register_role(**default_role) + self.config.register_member(**default_member) default_guildmessage: Dict[str, Dict[str, Any]] = {"reactroles": {"react_to_roleid": {}}} self.config.init_custom("GuildMessage", 2)