python-app-rocketbot/rocketbot/cogs/joinagecog.py

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='<seconds:float>',
		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': '<time_period>',
		}
	)

	@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))