python-app-rocketbot/cogs/patterncog.py

from discord import Guild, Message
from discord.ext import commands
from datetime import timedelta

from cogs.basecog import BaseCog, BotMessage
from storage import Storage

class Criterion:
	def __init__(self, type, **kwargs):
		self.type = type
		if type == 'contains':
			text = kwargs['text']
			self.text = text
			self.test = lambda m : text.lower() in m.content.lower()
		elif type == 'joinage':
			min = kwargs['min']
			self.min = min
			self.test = lambda m : m.created_at - m.author.joined_at < min
		else:
			raise RuntimeError(f'Unknown criterion type "{type}"')

	def matches(self, message: Message) -> bool:
		return self.test(message)

	@classmethod
	def decode(cls, val: dict):
		type = val['type']
		if type == 'contains':
			return Criterion(type, text=val['text'])
		elif type == 'joinage':
			return Criterion(type, min=timedelta(seconds=val['min']))

class Pattern:
	def __init__(self, criteria: list, action: str, must_match_all: bool = True):
		self.criteria = criteria
		self.action = action
		self.must_match_all = must_match_all

	def matches(self, message: Message) -> bool:
		for criterion in self.criteria:
			crit_matches = criterion.matches(message)
			if crit_matches and not self.must_match_all:
				return True
			if not crit_matches and self.must_match_all:
				return False
		return self.must_match_all

	@classmethod
	def decode(cls, val: dict):
		match_all = val.get('must_match_all')
		action = val.get('action')
		encoded_criteria = val.get('criteria')
		criteria = []
		for ec in encoded_criteria:
			criteria.append(Criterion.decode(ec))
		return Pattern(criteria, action, match_all if isinstance(match_all, bool) else True)

class PatternCog(BaseCog):
	def __init__(self, bot):
		super().__init__(bot)

	def __patterns(self, guild: Guild) -> list:
		patterns = Storage.get_state_value(guild, 'pattern_patterns')
		if patterns is None:
			patterns_encoded = Storage.get_config_value(guild, 'pattern_patterns')
			if patterns_encoded:
				patterns = []
				for pe in patterns_encoded:
					patterns.append(Pattern.decode(pe))
				Storage.set_state_value(guild, 'pattern_patterns', patterns)
		return patterns

	@commands.Cog.listener()
	async def on_message(self, message: Message) -> None:
		if message.guild is None or message.content is None or message.channel is None:
			return
		if message.author.id == self.bot.user.id:
			# Ignore self
			return
		if message.author.permissions_in(message.channel).ban_members:
			# Ignore mods
			return
		patterns = self.__patterns(message.guild)
		for pattern in patterns:
			if pattern.matches(message):
				msg = None
				if pattern.action == 'delete':
					await message.delete()
					msg = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Deleted.'
				elif pattern.action == 'kick':
					await message.delete()
					await message.author.kick(reason='Rocketbot: Message matched banned pattern')
					msg = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Message deleted and user kicked.'
				elif pattern.action == 'ban':
					await message.delete()
					await message.author.ban(reason='Rocketbot: Message matched banned pattern')
					msg = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Message deleted and user banned.'
				if msg:
					m = BotMessage(message.guild,
						text = msg,
						type = BotMessage.TYPE_MOD_WARNING)
					m.quote = message.content
					await self.post_message(m)
				break