python-app-rocketbot/cogs/crosspostcog.py

from discord import Guild, Member, Message, PartialEmoji
from discord.ext import commands
from datetime import datetime, timedelta
import math

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

class SpamContext:
	def __init__(self, member, message_hash):
		self.member = member
		self.message_hash = message_hash
		self.age = datetime.now()
		self.bot_message = None  # BotMessage
		self.is_kicked = False
		self.is_banned = False
		self.is_autobanned = False
		self.spam_messages = set()  # of Message
		self.deleted_messages = set()  # of Message

class CrossPostCog(BaseCog):
	SETTING_WARN_COUNT = CogSetting('warncount',
		brief='number of messages to trigger a warning',
		description='The number of identical messages to trigger a mod warning.',
		usage='<count:int>',
		min_value=2)
	SETTING_BAN_COUNT = CogSetting('bancount',
		brief='number of messages to trigger a ban',
		description='The number of identical messages to trigger an ' + \
			'automatic ban. Set to a large value to effectively disable, e.g. 9999.',
		usage='<count:int>',
		min_value=2)
	SETTING_MIN_LENGTH = CogSetting('minlength',
		brief='minimum message length',
		description='The minimum number of characters in a message to be ' + \
			'checked for duplicates. This can help ignore common short ' + \
			'messages like "lol" or a single emoji. Set to 0 to count all ' + \
			'message lengths.',
		usage='<character_count:int>',
		min_value=0)
	SETTING_TIMESPAN = CogSetting('timespan',
		brief='time window to look for dupe messages',
		description='The number of seconds of message history to look at ' + \
			'when looking for duplicates. Shorter values are preferred, ' + \
			'both to detect bots and avoid excessive memory usage.',
		usage='<seconds:int>',
		min_value=1)

	STATE_KEY_RECENT_MESSAGES = "crosspost_recent_messages"
	STATE_KEY_SPAM_CONTEXT = "crosspost_spam_context"

	CONFIG_KEY_WARN_COUNT = "crosspost_warn_count"
	CONFIG_KEY_BAN_COUNT = "crosspost_ban_count"
	CONFIG_KEY_MIN_MESSAGE_LENGTH = "crosspost_min_message_length"
	CONFIG_KEY_MESSAGE_AGE = "crosspost_message_age"

	def __init__(self, bot):
		super().__init__(bot)
		self.add_setting(CrossPostCog.SETTING_WARN_COUNT)
		self.add_setting(CrossPostCog.SETTING_BAN_COUNT)
		self.add_setting(CrossPostCog.SETTING_MIN_LENGTH)
		self.add_setting(CrossPostCog.SETTING_TIMESPAN)
		self.max_spam_contexts = 12

	# Config

	async def __record_message(self, message: Message) -> None:
		if message.author.permissions_in(message.channel).ban_members:
			# User exempt from spam detection
			return
		if len(message.content) < self.get_guild_setting(message.guild, self.SETTING_MIN_LENGTH):
			# Message too short to count towards spam total
			return
		max_age = timedelta(seconds=self.get_guild_setting(message.guild, self.SETTING_TIMESPAN))
		recent_messages = Storage.get_state_value(message.guild, self.STATE_KEY_RECENT_MESSAGES) \
			or AgeBoundList(max_age, lambda index, message : message.created_at)
		recent_messages.max_age = max_age
		recent_messages.append(message)
		Storage.set_state_value(message.guild, self.STATE_KEY_RECENT_MESSAGES, recent_messages)

		# Get all recent messages by user
		member_messages = [m for m in recent_messages if m.author.id == message.author.id]
		if len(member_messages) < self.get_guild_setting(message.guild, self.SETTING_WARN_COUNT):
			return

		# Look for repeats
		hash_to_count = {}
		max_count = 0
		for m in member_messages:
			key = hash(m.content)
			count = (hash_to_count.get(key) or 0) + 1
			hash_to_count[key] = count
			max_count = max(max_count, count)
		if max_count < self.get_guild_setting(message.guild, self.SETTING_WARN_COUNT):
			return

		# Handle the spam
		spam_lookup = Storage.get_state_value(message.guild, self.STATE_KEY_SPAM_CONTEXT) \
			or SizeBoundDict(self.max_spam_contexts, lambda key, context : context.age)
		Storage.set_state_value(message.guild, self.STATE_KEY_SPAM_CONTEXT, spam_lookup)
		for message_hash, count in hash_to_count.items():
			if count < self.get_guild_setting(message.guild, self.SETTING_WARN_COUNT):
				continue
			key = f'{message.author.id}|{message_hash}'
			context = spam_lookup.get(key)
			is_new = context is None
			if context is None:
				context = SpamContext(message.author, message_hash)
				spam_lookup[key] = context
				context.age = message.created_at
			for m in member_messages:
				if hash(m.content) == message_hash:
					context.spam_messages.add(m)
			await self.__update_from_context(context)

	async def __update_from_context(self, context: SpamContext):
		if len(context.spam_messages) >= self.get_guild_setting(context.member.guild, self.SETTING_BAN_COUNT):
			if not context.is_banned:
				count = len(context.spam_messages)
				await context.member.ban(reason=f'Autobanned by Rocketbot for posting same message {count} times', delete_message_days=1)
				context.is_kicked = True
				context.is_banned = True
				context.is_autobanned = True
				context.deleted_messages |= context.spam_messages
				self.log(context.member.guild, f'Bot autobanned {context.member.name} ({context.member.id}) for spamming')
			else:
				# Already banned. Nothing to update in the message.
				return
		await self.__update_message_from_context(context)

	async def __update_message_from_context(self, context: SpamContext) -> None:
		first_spam_message = next(iter(context.spam_messages))
		spam_count = len(context.spam_messages)
		deleted_count = len(context.deleted_messages)
		message = context.bot_message
		if message is None:
			message = BotMessage(context.member.guild, '',
				BotMessage.TYPE_MOD_WARNING, context)
			message.quote = first_spam_message.content
		if context.is_autobanned:
			text = f'User {context.member.mention} auto banned for ' + \
				f'posting the same message {deleted_count} ' + \
				'times. Messages from past 24 hours deleted.'
			await message.set_reactions([])
			await message.set_text(text)
		else:
			await message.set_text(f'User {context.member.mention} posted ' +
				f'the same message {spam_count} times.')
			await message.set_reactions(BotMessageReaction.standard_set(
				did_delete = deleted_count >= spam_count,
				message_count = spam_count,
				did_kick = context.is_kicked,
				did_ban = context.is_banned))
		if context.bot_message is None:
			await self.post_message(message)
			context.bot_message = message

	async def __delete_messages(self, context: SpamContext) -> None:
		for message in context.spam_messages - context.deleted_messages:
			await message.delete()
			context.deleted_messages.add(message)
		await self.__update_from_context(context)
		self.log(context.member.guild, f'Mod deleted messages from {context.member.name} ({context.member.id})')

	async def __kick(self, context: SpamContext) -> None:
		await context.member.kick(reason='Posting same message repeatedly')
		context.is_kicked = True
		await self.__update_from_context(context)
		self.log(context.member.guild, f'Mod kicked user {context.member.name} ({context.member.id})')

	async def __ban(self, context: SpamContext) -> None:
		await context.member.ban(reason='Posting same message repeatedly', delete_message_days=1)
		context.deleted_messages |= context.spam_messages
		context.is_kicked = True
		context.is_banned = True
		await self.__update_from_context(context)
		self.log(context.member.guild, f'Mod banned user {context.member.name} ({context.member.id})')

	async def on_mod_react(self,
			bot_message: BotMessage,
			reaction: BotMessageReaction,
			reacted_by: Member) -> None:
		context: SpamContext = bot_message.context
		if context is None:
			return

		if reaction.emoji == CONFIG['trash_emoji']:
			await self.__delete_messages(context)
		elif reaction.emoji == CONFIG['kick_emoji']:
			await self.__kick(context)
		elif reaction.emoji == CONFIG['ban_emoji']:
			await self.__ban(context)

	@commands.Cog.listener()
	async def on_message(self, message: Message):
		if message.author is None or \
				message.author.bot or \
				message.channel is None or \
				message.guild is None or \
				message.content is None or \
				message.content == '':
			return
		await self.__record_message(message)

	# -- Commands -----------------------------------------------------------

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