import weakref from datetime import datetime, timedelta, timezone from typing import Optional from discord import Guild, Member, Interaction from discord.app_commands import Group, guild_only, default_permissions, Transform from discord.ext.commands import Cog from config import CONFIG from rocketbot.bot import Rocketbot from rocketbot.cogs.basecog import BaseCog, BotMessage, BotMessageReaction, CogSetting from rocketbot.collections import AgeBoundList from rocketbot.storage import Storage from rocketbot.utils import TimeDeltaTransformer class JoinAgeQueryContext: """ Data about a join age query """ def __init__(self, join_members: list[Member], timespan: timedelta): self.join_members = list(join_members) self.timespan: timedelta = timespan self.kicked_members: set[Member] = set() self.banned_members: set[Member] = set() self.results_message_ref: Optional[weakref.ReferenceType[BotMessage]] = None class JoinAgeCog(BaseCog, name='Join Age'): """ Cog for finding users by when they joined. """ SETTING_ENABLED = CogSetting('enabled', bool, brief='join age', description='Whether this cog is enabled for a guild.') SETTING_JOIN_TIME = CogSetting('jointime', timedelta, brief='maximum length of time to track new joins', description='The number of seconds of join history to maintain.', usage='', min_value=1.0) STATE_KEY_RECENT_JOINS = "JoinAgeCog.recent_joins" def __init__(self, bot: Rocketbot): super().__init__( bot, config_prefix='joinage', name='join age', short_description='Tracks recently joined users with options to mass kick or ban.', ) self.add_setting(JoinAgeCog.SETTING_ENABLED) self.add_setting(JoinAgeCog.SETTING_JOIN_TIME) joinage = Group( name='joinage', description='Queries for users who joined in the past span of time', extras={ 'long_description': 'Searches for users who joined the server recently. ' + \ 'Can use time spans like 30s, 5m, 1h, 7d, etc.', 'usage': '', } ) @joinage.command( name='search', description='Searches for users who joined recently', ) @guild_only() @default_permissions(ban_members=True) async def search(self, interaction: Interaction, timespan: Transform[timedelta, TimeDeltaTransformer]) -> None: """Command handler""" guild: Guild = interaction.guild recent_joins: AgeBoundList[Member, datetime, timedelta] = Storage.get_state_value(guild, self.STATE_KEY_RECENT_JOINS) if recent_joins is None: max_age: timedelta = timedelta(seconds=self.get_guild_setting(guild, self.SETTING_JOIN_TIME)) recent_joins = AgeBoundList(max_age, lambda i, member0 : member0.joined_at) Storage.set_state_value(guild, self.STATE_KEY_RECENT_JOINS, recent_joins) results: list = [] cutoff: datetime = datetime.now(timezone.utc) - timespan for member in recent_joins: if member.joined_at > cutoff: results.append(member) ctx = JoinAgeQueryContext(results, timespan) msg = BotMessage(guild, text='', type=BotMessage.TYPE_INFO, context=ctx) ctx.results_message_ref = weakref.ref(msg) await self.__update_results_message(ctx) await self.post_message(msg) await interaction.response.send_message( "Search started", ephemeral=True, ) async def on_mod_react(self, bot_message: BotMessage, reaction: BotMessageReaction, reacted_by: Member) -> None: guild: Guild = bot_message.guild ctx: JoinAgeQueryContext = bot_message.context if reaction.emoji == CONFIG['kick_emoji']: to_kick = set(ctx.join_members) - ctx.kicked_members for member in to_kick: await member.kick( reason=f'Rocketbot: Mass kick based on join age, by {reacted_by.name}.') ctx.kicked_members |= to_kick await self.__update_results_message(ctx) self.log(guild, f'Users kicked by {reacted_by.name}.') elif reaction.emoji == CONFIG['ban_emoji']: to_ban = set(ctx.join_members) - ctx.banned_members for member in to_ban: await member.ban( reason=f'Rocketbot: Mass ban based on join age, by {reacted_by.name}.', delete_message_days=0) ctx.banned_members |= to_ban await self.__update_results_message(ctx) self.log(guild, f'Users banned by {reacted_by.name}') @Cog.listener() async def on_member_join(self, member: Member) -> None: """Event handler""" guild: Guild = member.guild if not self.get_guild_setting(guild, self.SETTING_ENABLED): return recent_joins: AgeBoundList[Member, datetime, timedelta] = Storage.get_state_value(guild, self.STATE_KEY_RECENT_JOINS) if recent_joins is None: max_age: timedelta = timedelta(seconds=self.get_guild_setting(guild, self.SETTING_JOIN_TIME)) recent_joins = AgeBoundList(max_age, lambda i, member : member.joined_at) Storage.set_state_value(guild, self.STATE_KEY_RECENT_JOINS, recent_joins) recent_joins.append(member) async def __update_results_message(self, context: JoinAgeQueryContext) -> None: if context.results_message_ref is None: return bot_message = context.results_message_ref() if bot_message is None: return text = f'The following members joined in the last {context.timespan}\n\n' if len(context.join_members) > 0: max_members = CONFIG['max_members_per_message'] for member in context.join_members[:max_members]: text += '\n• ' if member in context.banned_members: text += f'~~{member.mention} ({member.id})~~ - banned' elif member in context.kicked_members: text += f'~~{member.mention} ({member.id})~~ - kicked' else: text += f'{member.mention} ({member.id})' if len(context.join_members) > max_members: text += f'\n• {len(context.join_members) - max_members} more' else: text += 'No members found. If the bot was recently restarted ' + \ 'or JoinAgeCog just enabled, only new joins will be tracked.' await bot_message.set_text(text) if len(context.join_members) > 0: member_count = len(context.join_members) kick_count = len(context.kicked_members) ban_count = len(context.banned_members) await bot_message.set_reactions(BotMessageReaction.standard_set( did_kick=kick_count >= member_count, did_ban=ban_count >= member_count, user_count=member_count))