diff --git a/timerole/timerole.py b/timerole/timerole.py index 7484267..1b56b69 100644 --- a/timerole/timerole.py +++ b/timerole/timerole.py @@ -1,6 +1,7 @@ import asyncio import logging from datetime import datetime, timedelta +from typing import Optional import discord from redbot.core import Config, checks, commands @@ -19,6 +20,15 @@ async def sleep_till_next_hour(): await asyncio.sleep((next_hour - datetime.utcnow()).seconds) +async def announce_to_channel(channel, results, title): + if channel is not None and results: + await channel.send(title) + for page in pagify(results, shorten_by=50): + await channel.send(page) + elif results: # Channel is None, log the results + log.info(results) + + class Timerole(Cog): """Add roles to users based on time on server""" @@ -27,10 +37,15 @@ class Timerole(Cog): self.bot = bot self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True) default_global = {} - default_guild = {"announce": None, "roles": {}} + default_guild = {"announce": None, "reapply": True, "roles": {}} + default_rolemember = {"had_role": False, "check_again_time": None} self.config.register_global(**default_global) self.config.register_guild(**default_guild) + + self.config.init_custom("RoleMember", 2) + self.config.register_custom("RoleMember", **default_rolemember) + self.updating = asyncio.create_task(self.check_hour()) async def red_delete_data_for_user(self, **kwargs): @@ -49,11 +64,14 @@ class Timerole(Cog): Useful for troubleshooting the initial setup """ - async with ctx.typing(): + pre_run = datetime.utcnow() await self.timerole_update() + after_run = datetime.utcnow() await ctx.tick() + await ctx.maybe_send_embed(f"Took {after_run-pre_run} seconds") + @commands.group() @checks.mod_or_permissions(administrator=True) @commands.guild_only() @@ -84,9 +102,7 @@ class Timerole(Cog): await self.config.guild(guild).roles.set_raw(role.id, value=to_set) await ctx.maybe_send_embed( - "Time Role for {0} set to {1} days and {2} hours until added".format( - role.name, days, hours - ) + f"Time Role for {role.name} set to {days} days and {hours} hours until added" ) @timerole.command() @@ -114,18 +130,27 @@ class Timerole(Cog): await self.config.guild(guild).roles.set_raw(role.id, value=to_set) await ctx.maybe_send_embed( - "Time Role for {0} set to {1} days and {2} hours until removed".format( - role.name, days, hours - ) + f"Time Role for {role.name} set to {days} days and {hours} hours until removed" ) @timerole.command() - async def channel(self, ctx: commands.Context, channel: discord.TextChannel): + async def channel(self, ctx: commands.Context, channel: Optional[discord.TextChannel] = None): """Sets the announce channel for role adds""" guild = ctx.guild + if channel is None: + await self.config.guild(guild).announce.clear() + await ctx.maybe_send_embed(f"Announce channel has been cleared") + else: + await self.config.guild(guild).announce.set(channel.id) + await ctx.send(f"Announce channel set to {channel.mention}") - await self.config.guild(guild).announce.set(channel.id) - await ctx.send("Announce channel set to {0}".format(channel.mention)) + @timerole.command() + async def reapply(self, ctx: commands.Context): + """Toggle reapplying roles if the member loses it somehow. Defaults to True""" + guild = ctx.guild + current_setting = await self.config.guild(guild).reapply() + await self.config.guild(guild).reapply.set(not current_setting) + await ctx.maybe_send_embed(f"Reapplying roles is now set to: {not current_setting}") @timerole.command() async def delrole(self, ctx: commands.Context, role: discord.Role): @@ -133,7 +158,8 @@ class Timerole(Cog): guild = ctx.guild await self.config.guild(guild).roles.set_raw(role.id, value=None) - await ctx.send("{0} will no longer be applied".format(role.name)) + await self.config.custom("RoleMember", role.id).clear() + await ctx.maybe_send_embed(f"{role.name} will no longer be applied") @timerole.command() async def list(self, ctx: commands.Context): @@ -153,89 +179,197 @@ class Timerole(Cog): str(discord.utils.get(guild.roles, id=int(new_id))) for new_id in r_data["required"] ] - out += "{} | {} days | requires: {}\n".format(str(role), r_data["days"], r_roles) + out += f"{role} | {r_data['days']} days | requires: {r_roles}\n" await ctx.maybe_send_embed(out) async def timerole_update(self): - async for guild in AsyncIter(self.bot.guilds): - addlist = [] - removelist = [] + utcnow = datetime.utcnow() + all_guilds = await self.config.all_guilds() + + # all_mrs = await self.config.custom("RoleMember").all() + + # log.debug(f"Begin timerole update") + + for guild in self.bot.guilds: + guild_id = guild.id + if guild_id not in all_guilds: + log.debug(f"Guild has no configured settings: {guild}") + continue + + add_results = "" + remove_results = "" + reapply = all_guilds[guild_id]["reapply"] + role_dict = all_guilds[guild_id]["roles"] - role_dict = await self.config.guild(guild).roles() if not any(role_data for role_data in role_dict.values()): # No roles + log.debug(f"No roles are configured for guild: {guild}") continue - async for member in AsyncIter(guild.members): - has_roles = [r.id for r in member.roles] + # all_mr = await self.config.all_custom("RoleMember") + # log.debug(f"{all_mr=}") + + async for member in AsyncIter(guild.members, steps=10): + addlist = [] + removelist = [] + + for role_id, role_data in role_dict.items(): + # Skip non-configured roles + if not role_data: + continue + + mr_dict = await self.config.custom("RoleMember", role_id, member.id).all() + + # Stop if they've had the role and reapplying is disabled + if not reapply and mr_dict["had_role"]: + log.debug(f"{member.display_name} - Not reapplying") + continue + + # Stop if the check_again_time hasn't passed yet + if ( + mr_dict["check_again_time"] is not None + and datetime.fromisoformat(mr_dict["check_again_time"]) >= utcnow + ): + log.debug(f"{member.display_name} - Not time to check again yet") + continue + member: discord.Member + has_roles = set(r.id for r in member.roles) + + # Stop if they currently have or don't have the role, and mark had_role + if (role_id in has_roles and not role_data["remove"]) or ( + role_id not in has_roles and role_data["remove"] + ): + if not mr_dict["had_role"]: + await self.config.custom( + "RoleMember", role_id, member.id + ).had_role.set(True) + log.debug(f"{member.display_name} - applying had_role") + continue + + # Stop if they don't have all the required roles + if role_data is None or ( + "required" in role_data and not set(role_data["required"]) & has_roles + ): + continue + + check_time = member.joined_at + timedelta( + days=role_data["days"], + hours=role_data.get("hours", 0), + ) + + # Check if enough time has passed to get the role and save the check_again_time + if check_time >= utcnow: + await self.config.custom( + "RoleMember", role_id, member.id + ).check_again_time.set(check_time.isoformat()) + log.debug( + f"{member.display_name} - Not enough time has passed to qualify for the role\n" + f"Waiting until {check_time}" + ) + continue + + if role_data["remove"]: + removelist.append(role_id) + else: + addlist.append(role_id) + + # Done iterating through roles, now add or remove the roles + if not addlist and not removelist: + continue + # log.debug(f"{addlist=}\n{removelist=}") add_roles = [ - int(rID) - for rID, r_data in role_dict.items() - if r_data is not None and not r_data["remove"] + discord.utils.get(guild.roles, id=int(role_id)) for role_id in addlist ] remove_roles = [ - int(rID) - for rID, r_data in role_dict.items() - if r_data is not None and r_data["remove"] + discord.utils.get(guild.roles, id=int(role_id)) for role_id in removelist ] - check_add_roles = set(add_roles) - set(has_roles) - check_remove_roles = set(remove_roles) & set(has_roles) - - await self.check_required_and_date( - addlist, check_add_roles, has_roles, member, role_dict - ) - await self.check_required_and_date( - removelist, check_remove_roles, has_roles, member, role_dict - ) - + if None in add_roles or None in remove_roles: + log.info( + f"Timerole ran into an error with the roles in: {add_roles + remove_roles}" + ) + + if addlist: + try: + await member.add_roles(*add_roles, reason="Timerole", atomic=False) + except (discord.Forbidden, discord.NotFound) as e: + log.exception("Failed Adding Roles") + add_results += f"{member.display_name} : **(Failed Adding Roles)**\n" + else: + add_results += " \n".join( + f"{member.display_name} : {role.name}" for role in add_roles + ) + for role_id in addlist: + await self.config.custom( + "RoleMember", role_id, member.id + ).had_role.set(True) + + if removelist: + try: + await member.remove_roles(*remove_roles, reason="Timerole", atomic=False) + except (discord.Forbidden, discord.NotFound) as e: + log.exception("Failed Removing Roles") + remove_results += f"{member.display_name} : **(Failed Removing Roles)**\n" + else: + remove_results += " \n".join( + f"{member.display_name} : {role.name}" for role in remove_roles + ) + for role_id in removelist: + await self.config.custom( + "RoleMember", role_id, member.id + ).had_role.set(True) + + # Done iterating through members, now maybe announce to the guild channel = await self.config.guild(guild).announce() if channel is not None: channel = guild.get_channel(channel) - title = "**These members have received the following roles**\n" - await self.announce_roles(title, addlist, channel, guild, to_add=True) - title = "**These members have lost the following roles**\n" - await self.announce_roles(title, removelist, channel, guild, to_add=False) - - async def announce_roles(self, title, role_list, channel, guild, to_add: True): - results = "" - async for member, role_id in AsyncIter(role_list): - role = discord.utils.get(guild.roles, id=role_id) - try: - if to_add: - await member.add_roles(role, reason="Timerole") - else: - await member.remove_roles(role, reason="Timerole") - except (discord.Forbidden, discord.NotFound) as e: - results += "{} : {} **(Failed)**\n".format(member.display_name, role.name) - else: - results += "{} : {}\n".format(member.display_name, role.name) - if channel is not None and results: - await channel.send(title) - for page in pagify(results, shorten_by=50): - await channel.send(page) - elif results: # Channel is None, log the results - log.info(results) - - async def check_required_and_date(self, role_list, check_roles, has_roles, member, role_dict): - async for role_id in AsyncIter(check_roles): - # Check for required role - if "required" in role_dict[str(role_id)]: - if not set(role_dict[str(role_id)]["required"]) & set(has_roles): - # Doesn't have required role - continue - - if ( - member.joined_at - + timedelta( - days=role_dict[str(role_id)]["days"], - hours=role_dict[str(role_id)].get("hours", 0), - ) - <= datetime.today() - ): - # Qualifies - role_list.append((member, role_id)) + if add_results: + title = "**These members have received the following roles**\n" + await announce_to_channel(channel, add_results, title) + if remove_results: + title = "**These members have lost the following roles**\n" + await announce_to_channel(channel, remove_results, title) + # End + + # async def announce_roles(self, title, role_list, channel, guild, to_add: True): + # results = "" + # async for member, role_id in AsyncIter(role_list): + # role = discord.utils.get(guild.roles, id=role_id) + # try: + # if to_add: + # await member.add_roles(role, reason="Timerole") + # else: + # await member.remove_roles(role, reason="Timerole") + # except (discord.Forbidden, discord.NotFound) as e: + # results += f"{member.display_name} : {role.name} **(Failed)**\n" + # else: + # results += f"{member.display_name} : {role.name}\n" + # if channel is not None and results: + # await channel.send(title) + # for page in pagify(results, shorten_by=50): + # await channel.send(page) + # elif results: # Channel is None, log the results + # log.info(results) + + # async def check_required_and_date(self, role_list, check_roles, has_roles, member, role_dict): + # async for role_id in AsyncIter(check_roles): + # # Check for required role + # if "required" in role_dict[str(role_id)]: + # if not set(role_dict[str(role_id)]["required"]) & set(has_roles): + # # Doesn't have required role + # continue + # + # if ( + # member.joined_at + # + timedelta( + # days=role_dict[str(role_id)]["days"], + # hours=role_dict[str(role_id)].get("hours", 0), + # ) + # <= datetime.utcnow() + # ): + # # Qualifies + # role_list.append((member, role_id)) async def check_hour(self): await sleep_till_next_hour()