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.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
		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):
				text = None
				if pattern.action == 'delete':
					await message.delete()
					text = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Deleted.'
					self.log(message.guild, 'Message matched pattern. Deleted.')
				elif pattern.action == 'kick':
					await message.delete()
					await message.author.kick(reason='Rocketbot: Message matched banned pattern')
					text = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Message deleted and user kicked.'
					self.log(message.guild, 'Message matched pattern. Kicked user.')
				elif pattern.action == 'ban':
					await message.delete()
					await message.author.ban(reason='Rocketbot: Message matched banned pattern')
					text = f'Message from {message.author.mention} matched ' + \
						'banned pattern. Message deleted and user banned.'
					self.log(message.guild, 'Message matched pattern. Banned user.')
				if text:
					m = BotMessage(message.guild,
						text = msg,
						type = BotMessage.TYPE_MOD_WARNING)
					m.quote = message.content
					await self.post_message(m)
				break

	"""
	Expression language samples:

		content contains "poop"
		content contains "poop" and content contains "tinkle"
		joinage < 600s
		(content contains "this" and content contains "that") or content contains "whatever"

		<field> <op> <value>

		Fields:	
			content
			author.id
			author.name
			author.joinage

		Ops:
			==
			!=
			<
			>
			<=
			>=
			contains, !contains  -- plain strings
			matches, !matches  -- regexes

		Value types:
			timedelta  (600, 600s, 10m, 5m30s)
			number
			string
			regex
			mention
	"""

	@commands.group(
		brief='Manages message pattern matching',
	)
	@commands.has_permissions(ban_members=True)
	@commands.guild_only()
	async def patterns(self, context: commands.Context):
		'Message pattern matching'
		if context.invoked_subcommand is None:
			await context.send_help()

	@patterns.command()
	async def addpattern(self, context: commands.Context, name: str, expression: str, *args):
		print(f'Pattern name: {name}')
		tokens = []
		tokens.append(expression)
		tokens += args
		print('Expression: ' + (' '.join(tokens)))