python-app-rocketbot/cogs/joinraidcog.py

from datetime import datetime, timedelta
from discord import Guild, Member
from discord.ext import commands
import weakref

from cogs.basecog import BaseCog, BotMessage, BotMessageReaction, CogSetting
from config import CONFIG
from rscollections import AgeBoundList
from storage import Storage

class JoinRaidContext:
	def __init__(self, join_members: list):
		self.join_members = list(join_members)
		self.kicked_members = set()
		self.banned_members = set()
		self.warning_message_ref = None

	def last_join_time(self) -> datetime:
		return self.join_members[-1].joined_at

class JoinRaidCog(BaseCog):
	"""
	Cog for monitoring member joins and detecting potential bot raids.
	"""
	SETTING_ENABLED = CogSetting('enabled', bool,
			brief='join raid detection',
			description='Whether this cog is enabled for a guild.')
	SETTING_JOIN_COUNT = CogSetting('joincount', int,
			brief='number of joins to trigger a warning',
			description='The number of joins occuring within the time ' + \
				'window to trigger a mod warning.',
			usage='<count:int>',
			min_value=2)
	SETTING_JOIN_TIME = CogSetting('jointime', float,
			brief='time window length to look for joins',
			description='The number of seconds of join history to look ' + \
				'at when counting recent joins. If joincount or more ' + \
				'joins occur within jointime seconds a mod warning is issued.',
			usage='<seconds:float>',
			min_value=1.0,
			max_value=900.0)

	STATE_KEY_RECENT_JOINS = "JoinRaidCog.recent_joins"	
	STATE_KEY_LAST_RAID = "JoinRaidCog.last_raid"

	def __init__(self, bot):
		super().__init__(bot)
		self.add_setting(JoinRaidCog.SETTING_ENABLED)
		self.add_setting(JoinRaidCog.SETTING_JOIN_COUNT)
		self.add_setting(JoinRaidCog.SETTING_JOIN_TIME)

	@commands.group(
		brief='Manages join raid detection and handling',
	)
	@commands.has_permissions(ban_members=True)
	@commands.guild_only()
	async def joinraid(self, context: commands.Context):
		'Command group'
		if context.invoked_subcommand is None:
			await context.send_help()

	async def on_mod_react(self,
			bot_message: BotMessage,
			reaction: BotMessageReaction,
			reacted_by: Member) -> None:
		guild: Guild = bot_message.guild
		raid: JoinRaidRecord = bot_message.context
		if reaction.emoji == CONFIG['kick_emoji']:
			to_kick = set(raid.join_members) - raid.kicked_members
			for member in to_kick:
				await member.kick(
					reason=f'Rocketbot: Part of join raid. Kicked by {reacted_by.name}.')
			raid.kicked_members |= to_kick
			await self.__update_warning_message(raid)
			self.log(guild, f'Join raid users kicked by {reacted_by.name}.')
		elif reaction.emoji == CONFIG['ban_emoji']:
			to_ban = set(raid.join_members) - raid.banned_members
			for member in to_ban:
				await member.ban(
					reason=f'Rocketbot: Part of join raid. Banned by {reacted_by.name}.',
					delete_message_days=0)
			raid.banned_members |= to_ban
			await self.__update_warning_message(raid)
			self.log(guild, f'Join raid 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
		min_count = self.get_guild_setting(guild, self.SETTING_JOIN_COUNT)
		seconds = self.get_guild_setting(guild, self.SETTING_JOIN_TIME)
		timespan: timedelta = timedelta(seconds=seconds)

		last_raid: JoinRaidContext = Storage.get_state_value(guild, self.STATE_KEY_LAST_RAID)
		recent_joins: AgeBoundList = Storage.get_state_value(guild, self.STATE_KEY_RECENT_JOINS)
		if recent_joins is None:
			recent_joins = AgeBoundList(timespan, lambda i, member : member.joined_at)
			Storage.set_state_value(guild, self.STATE_KEY_RECENT_JOINS, recent_joins)
		if last_raid:
			if member.joined_at - last_raid.last_join_time() > timespan:
				# Last raid is over
				Storage.set_state_value(guild, self.STATE_KEY_LAST_RAID, None)
				recent_joins.append(member)
				return
			# Add join to existing raid
			last_raid.join_members.append(member)
			await self.__update_warning_message(last_raid)
		else:
			# Add join to the general, non-raid recent join list
			recent_joins.append(member)
			if len(recent_joins) >= min_count:
				self.log(guild, '\u0007Join raid detected')
				last_raid = JoinRaidContext(recent_joins)
				Storage.set_state_value(guild, self.STATE_KEY_LAST_RAID, last_raid)
				recent_joins.clear()
				msg = BotMessage(guild,
						text='',
						type=BotMessage.TYPE_MOD_WARNING,
						context=last_raid)
				last_raid.warning_message_ref = weakref.ref(msg)
				await self.__update_warning_message(last_raid)
				await self.post_message(msg)

	async def on_setting_updated(self, guild: Guild, setting: CogSetting) -> None:
		if setting is self.SETTING_JOIN_TIME:
			seconds = self.get_guild_setting(guild, self.SETTING_JOIN_TIME)
			timespan: timedelta = timedelta(seconds=seconds)
			recent_joins: AgeBoundList = Storage.get_state_value(guild,
				self.STATE_KEY_RECENT_JOINS)
			if recent_joins:
				recent_joins.max_age = timespan
				recent_joins.purge_old_elements()

	async def __update_warning_message(self, context: JoinRaidContext) -> None:
		if context.warning_message_ref is None:
			return
		bot_message = context.warning_message_ref()
		if bot_message is None:
			return
		text = 'JOIN RAID DETECTED\n\n' + \
			'The following members joined in close succession:\n'
		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})'
		text += '\n_(list updates automatically)_'
		await bot_message.set_text(text)
		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))