python-app-rocketbot/rocketbot/cogs/usernamecog.py

"""
Cog for detecting username patterns.
"""
from typing import Optional

from discord import Guild, Member
from discord.ext import commands

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

class UsernamePatternContext:
	"""
	BotMessage context for a flagged username
	"""
	def __init__(self, member: Member) -> None:
		self.member: Member = member
		self.kicked_by: Optional[Member] = None
		self.banned_by: Optional[Member] = None
		self.ignored_by: Optional[Member] = None

	def reactions(self) -> list[BotMessageReaction]:
		"""
		Generates updated BotMessageReactions based on context state.
		"""
		r: list[BotMessageReaction] = []
		if self.ignored_by:
			r.append(BotMessageReaction(CONFIG['ignore_emoji'], False, f'Ignored by {self.ignored_by.name}'))
		elif self.banned_by:
			r.append(BotMessageReaction(CONFIG['ban_emoji'], False, f'Banned by {self.banned_by.name}'))
		elif self.kicked_by:
			r.append(BotMessageReaction(CONFIG['kick_emoji'], False, f'Kicked by {self.kicked_by.name}'))
			r.append(BotMessageReaction(CONFIG['ban_emoji'], True, 'Ban user'))
		else:
			r.append(BotMessageReaction(CONFIG['kick_emoji'], True, 'Kick user'))
			r.append(BotMessageReaction(CONFIG['ban_emoji'], True, 'Ban user'))
			r.append(BotMessageReaction(CONFIG['ignore_emoji'], True, 'Ignore warning'))
		return r

class UsernamePatternCog(BaseCog, name='Username Pattern'):
	"""
	Detects usernames that match certain flagged patterns. Posts a mod warning
	message on a match.
	"""

	SETTING_ENABLED = CogSetting('enabled', bool,
			brief='username pattern detection',
			description='Whether new users are checked for common patterns.')

	SETTING_PATTERNS = CogSetting('patterns', None)

	def __init__(self, bot):
		super().__init__(bot, 'username')
		self.add_setting(UsernamePatternCog.SETTING_ENABLED)

	def __get_patterns(self, guild: Guild) -> list[str]:
		"""
		Returns an array of username patterns.
		"""
		patterns: list[str] = self.get_guild_setting(guild, self.SETTING_PATTERNS)
		if patterns is None:
			patterns = []
			Storage.set_config_value(guild, 'UsernamePatternCog.patterns', patterns)
		return patterns

	@classmethod
	def __save_patterns(cls,
			guild: Guild,
			patterns: list[str]) -> None:
		"""
		Saves username pattern array.
		"""
		cls.set_guild_setting(guild, cls.SETTING_PATTERNS, patterns)

	@commands.group(
		brief='Manages username pattern detection'
	)
	@commands.has_permissions(ban_members=True)
	@commands.guild_only()
	async def username(self, context: commands.Context):
		"""Username pattern command group"""
		if context.invoked_subcommand is None:
			await context.send_help()

	@username.command(
		brief='Adds a username pattern',
		description='Adds a username pattern.',
		usage='<pattern>'
	)
	async def add(self, context: commands.Context, pattern: str) -> None:
		"""Command handler"""
		norm_pattern = pattern.lower()
		patterns: list[str] = self.__get_patterns(context.guild)
		if norm_pattern in patterns:
			await context.reply(f'Pattern `{norm_pattern}` already added.', mention_author=False)
			return
		patterns.append(norm_pattern)
		self.__save_patterns(context.guild, patterns)
		await context.reply(f'Pattern `{norm_pattern}` added.', mention_author=False)

	@username.command(
		brief='Removes a username pattern',
		description='Removes an existing username pattern',
		usage='<pattern>'
	)
	async def remove(self, context: commands.Context, pattern: str) -> None:
		"""Command handler"""
		norm_pattern = pattern.lower()
		guild: Guild = context.guild
		patterns: list[str] = self.__get_patterns(guild)
		len_before = len(patterns)
		patterns = list(filter(lambda p: p != norm_pattern, patterns))
		if len(patterns) == len_before:
			await context.reply(f'Pattern `{norm_pattern}` not found.', mention_author=False)
			return
		self.__save_patterns(guild, patterns)
		await context.reply(f'Pattern `{norm_pattern}` removed.', mention_author=False)

	@username.command(
		brief='Lists username patterns'
	)
	async def list(self, context: commands.Context) -> None:
		"""Command handler"""
		guild: Guild = context.guild
		patterns: list[str] = self.__get_patterns(guild)
		if len(patterns) == 0:
			await context.reply('No patterns defined', mention_author=False)
		else:
			msg = 'Patterns:\n\n> `' + '`\n> `'.join(patterns) + '`'
			await context.reply(msg, mention_author=False)

	@commands.Cog.listener()
	async def on_member_join(self, member: Member) -> None:
		"""Event handler"""
		for pattern in self.__get_patterns(member.guild):
			if self.matches(pattern, member.name):
				await self.handle_match(member, pattern)
			elif self.matches(pattern, member.display_name):
				await self.handle_match(member, pattern)

	def matches(self, pattern: str, subject: str) -> bool:
		"""Checks if a username matches a given pattern"""
		if pattern is None:
			return False
		if subject is None:
			return False
		return pattern.lower() in subject.lower()

	async def handle_match(self, member: Member, pattern: str) -> None:
		"""
		Handles a username match.
		"""
		# TODO: Prevent double handling?
		self.log(member.guild, f'User {member.id} {member.display_name} matches pattern "{pattern}"')
		context = UsernamePatternContext(member)
		bm = BotMessage(
			member.guild,
			f'User {member.mention} ({str(member.id)}, {member.display_name}) has ' +
			f'username matching pattern `{pattern}`.',
			BotMessage.TYPE_INFO if self.was_warned_recently(member) else BotMessage.TYPE_MOD_WARNING,
			context)
		self.record_warning(member)
		await bm.set_reactions(context.reactions())
		await self.post_message(bm)

	async def on_mod_react(self,
			bot_message: BotMessage,
			reaction: BotMessageReaction,
			reacted_by: Member) -> None:
		context: UsernamePatternContext = bot_message.context
		if reaction.emoji == CONFIG['kick_emoji']:
			await context.member.kick(
				reason=f'Rocketbot: Flagged username pattern. Kicked by {reacted_by.name}.')
			context.kicked_by = reacted_by
			self.log(context.member.guild, f'User {context.member.name} kicked by {reacted_by.name}')
			await bot_message.set_reactions(context.reactions())
		elif reaction.emoji == CONFIG['ban_emoji']:
			await context.member.ban(
				reason=f'Rocketbot: Flagged username pattern. Banned by {reacted_by.name}.',
				delete_message_days=0)
			context.banned_by = reacted_by
			self.log(context.member.guild, f'User {context.member.name} banned by {reacted_by.name}')
			await bot_message.set_reactions(context.reactions())
		elif reaction.emoji == CONFIG['ignore_emoji']:
			context.ignored_by = reacted_by
			self.log(context.member.guild, f'Warning ignored by {reacted_by.name}')
			await bot_message.set_reactions(context.reactions())