Adding JoinAgeCog for searching for recent users

bot.py

@@ -13,6 +13,7 @@ from config import CONFIG
 from rocketbot.cogs.configcog import ConfigCog
 from rocketbot.cogs.crosspostcog import CrossPostCog
 from rocketbot.cogs.generalcog import GeneralCog
+from rocketbot.cogs.joinagecog import JoinAgeCog
 from rocketbot.cogs.joinraidcog import JoinRaidCog
 from rocketbot.cogs.patterncog import PatternCog
 from rocketbot.cogs.urlspamcog import URLSpamCog
@@ -63,6 +64,7 @@ bot.add_cog(ConfigCog(bot))
 
 # Optional
 bot.add_cog(CrossPostCog(bot))
+bot.add_cog(JoinAgeCog(bot))
 bot.add_cog(JoinRaidCog(bot))
 bot.add_cog(PatternCog(bot))
 bot.add_cog(URLSpamCog(bot))

config.py.sample

@@ -20,6 +20,9 @@ CONFIG = {
 			'minlength': 1,
 			'timespan': 60,
 		},
+		'JoinAgeCog': {
+			'jointime': 3600,
+		},
 		'JoinRaidCog': {
 			'enabled': False,
 			'warning_count': 5,

rocketbot/cogs/joinagecog.py

@@ -0,0 +1,147 @@
+import weakref
+
+from datetime import datetime, timedelta
+from discord import Guild, Member
+from discord.ext import commands
+
+from config import CONFIG
+from rocketbot.cogs.basecog import BaseCog, BotMessage, BotMessageReaction, CogSetting
+from rocketbot.collections import AgeBoundList
+from rocketbot.storage import Storage
+from rocketbot.utils import timedelta_from_str
+
+class JoinAgeQueryContext:
+	"""
+	Data about a join age query
+	"""
+	def __init__(self, join_members: list, timespan: str):
+		self.join_members = list(join_members)
+		self.timespan = timespan
+		self.kicked_members = set()
+		self.banned_members = set()
+		self.results_message_ref = 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', float,
+		brief='maximum length of time to track new joins',
+		description='The number of seconds of join history to maintain.',
+		usage='<seconds:float>',
+		min_value=1.0)
+
+	STATE_KEY_RECENT_JOINS = "JoinAgeCog.recent_joins"
+
+	def __init__(self, bot):
+		super().__init__(bot)
+		self.add_setting(JoinAgeCog.SETTING_ENABLED)
+		self.add_setting(JoinAgeCog.SETTING_JOIN_TIME)
+
+	@commands.group(
+		brief='Tracks recently joined users with options to mass kick or ban',
+	)
+	@commands.has_permissions(ban_members=True)
+	@commands.guild_only()
+	async def joinage(self, context: commands.Context):
+		'Join age tracking'
+		if context.invoked_subcommand is None:
+			await context.send_help()
+
+	@joinage.command(
+		brief='Queries for users who joined in the past span of time',
+		description='Searches for users who joined the server recently. ' + \
+			'Can use time spans like 30s, 5m, 1h, 7d, etc.',
+		usage='<time_period>'
+	)
+	async def search(self, context: commands.Context, timespan: str):
+		'Command handler'
+		guild: Guild = context.guild
+		recent_joins: AgeBoundList = 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)
+		results: list = []
+		ts: timedelta = timedelta_from_str(timespan)
+		cutoff: datetime = datetime.utcnow() - ts
+		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)
+
+	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}')
+
+	@commands.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 = 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:
+			for member in context.join_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})'
+		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))