python-app-rocketbot/cogs/crosspostcog.py

from discord import Guild, Message
from cogs.basecog import BaseCog

class CrossPostCog(BaseCog):
	class SpamContext:
		def __init__(self, member):
			self.member = member
			self.is_warned = False

	STATE_KEY_RECENT_MESSAGES = "crosspost_recent_messages"
	STATE_KEY_SPAM_CONTEXT = "crosspost_spam_context"

	def __init__(self, bot):
		super().__init__(bot)
		self.max_recent_messages = 20
		self.messages_per_user = 3
		self.min_message_length = 10

	def __record_message(self, message: Message) -> None:
		recent_messages = Storage.get_state_value(message.guild, self.STATE_KEY_RECENT_MESSAGES) or []
		recent_messages.append(message)
		while len(recent_messages) > self.max_recent_messages:
			recent_messages.pop(0)
		Storage.set_state_value(message.guild, self.STATE_KEY_RECENT_MESSAGES, recent_messages)

	def __check_for_spam(self, guild: Guild) -> None:
		recent_messages = Storage.get_state_value(guild, self.STATE_KEY_RECENT_MESSAGES)
		if recent_messages is None:
			return
		user_id_to_count = {}
		spamming_members = set()
		for message in recent_messages:
			content = message.content
			if len(content) < self.min_message_length:
				continue
			key = str(message.author.id) + '|' + hash(content)
			count = (user_id_to_count.get(key) or 0) + 1
			user_id_to_count[key] = count
			if count >= self.messages_per_user:
				spamming_members.add(message.author)
		for member in spamming_members:
			context = self.__spam_context_for_user(member, True)

	def __spam_context_for_user(self,
			member: Member,
			create_if_missing: bool = False) -> SpamContext:
		spam_lookup = Storage.get_state_value(member.guild, self.STATE_KEY_SPAM_CONTEXT) or {}
		context = spam_lookup.get(member.id)
		if context:
			return context
		if not create_if_missing:
			return None
		context = SpamContext(member)
		spam_lookup[member.id] = context
		Storage.set_state_value(member.guild, self.STATE_KEY_SPAM_CONTEXT, spam_lookup)
		return context

	@commands.Cog.listener()
	def on_message(self, message: Message):
		self.__record_message(message)
		self.__check_for_spam(message.guild)