python-app-rocketbot/cogs/basecog.py

from discord import Guild, Message, PartialEmoji, RawReactionActionEvent, TextChannel
from discord.ext import commands
from datetime import timedelta

from rscollections import AgeBoundDict
from storage import ConfigKey, Storage
import json

class BaseCog(commands.Cog):
	def __init__(self, bot):
		self.bot = bot
		self.listened_mod_react_message_ids = AgeBoundDict(timedelta(minutes=5), lambda message_id, tpl : tpl[0])

	@classmethod
	def get_cog_default(cls, key: str):
		"""
		Convenience method for getting a cog configuration default from
		`CONFIG['cogs'][<cogname>][<key>]`.
		"""
		cogs: dict = CONFIG['cog_defaults']
		cog = cogs.get(cls.__name__)
		if cog is None:
			return None
		return cog.get(key)

	def listen_for_reactions_to(self, message: Message, context = None) -> None:
		"""
		Registers a warning message as something a mod may react to to enact
		some action. `context` will be passed back in `on_mod_react` and can be
		any value that helps give the cog context about the action being
		performed.
		"""
		self.listened_mod_react_message_ids[message.id] = (message.created_at, context)

	@commands.Cog.listener()
	async def on_raw_reaction_add(self, payload: RawReactionActionEvent):
		'Event handler'
		if payload.user_id == self.bot.user.id:
			# Ignore bot's own reactions
			return
		member: Member = payload.member
		if member is None:
			return
		guild: Guild = self.bot.get_guild(payload.guild_id)
		if guild is None:
			# Possibly a DM
			return
		channel: GuildChannel = guild.get_channel(payload.channel_id)
		if channel is None:
			# Possibly a DM
			return
		message: Message = await channel.fetch_message(payload.message_id)
		if message is None:
			# Message deleted?
			return
		if message.author.id != self.bot.user.id:
			# Bot didn't author this
			return
		if not member.permissions_in(channel).ban_members:
			# Not a mod
			return
		tpl = self.listened_mod_react_message_ids.get(message.id)
		if tpl is None:
			# Not a message we're listening for
			return
		context = tpl[1]
		await self.on_mod_react(message, payload.emoji, context)

	async def on_mod_react(self, message: Message, emoji: PartialEmoji, context) -> None:
		"""
		Override point for getting a mod's emote on a bot message. Used to take
		action on a warning, such as banning an offending user. This event is
		only triggered for registered bot messages and reactions by members
		with the proper permissions. The given `context` value is whatever was
		passed in `listen_to_reactions_to()`.
		"""
		pass

	async def validate_param(self, context: commands.Context, param_name: str, value,
		allowed_types: tuple = None,
		min_value = None,
		max_value = None) -> bool:
		"""
		Convenience method for validating a command parameter is of the expected
		type and in the expected range. Bad values will cause a reply to be sent
		to the original message and a False will be returned. If all checks
		succeed, True will be returned.
		"""
		if allowed_types is not None and not isinstance(value, allowed_types):
			if len(allowed_types) == 1:
				await context.message.reply(f'⚠️ `{param_name}` must be of type ' +
					f'{allowed_types[0]}.', mention_author=False)
			else:
				await context.message.reply(f'⚠️ `{param_name}` must be of types ' +
					f'{allowed_types}.', mention_author=False)
			return False
		if min_value is not None and value < min_value:
			await context.message.reply(f'⚠️ `{param_name}` must be >= {min_value}.',
				mention_author=False)
			return False
		if max_value is not None and value > max_value:
			await context.message.reply(f'⚠️ `{param_name}` must be <= {max_value}.',
				mention_author=False)
		return True

	@classmethod
	async def warn(cls, guild: Guild, message: str) -> Message:
		"""
		Sends a warning message to the configured warning channel for the
		given guild. If no warning channel is configured no action is taken.
		Returns the Message if successful or None if not.
		"""
		channel_id = Storage.get_config_value(guild, ConfigKey.WARNING_CHANNEL_ID)
		if channel_id is None:
			cls.guild_trace(guild, 'No warning channel set! No warning issued.')
			return None
		channel: TextChannel = guild.get_channel(channel_id)
		if channel is None:
			cls.guild_trace(guild, 'Configured warning channel does not exist!')
			return None
		mention: str = Storage.get_config_value(guild, ConfigKey.WARNING_MENTION)
		text: str = message
		if mention is not None:
			text = f'{mention} {text}'
		msg: Message = await channel.send(text)
		return msg

	@classmethod
	async def update_warn(cls, warn_message: Message, new_text: str) -> None:
		"""
		Updates the text of a previously posted `warn`. Includes configured
		mentions if necessary.
		"""
		text: str = new_text
		mention: str = Storage.get_config_value(
			warn_message.guild,
			ConfigKey.WARNING_MENTION)
		if mention is not None:
			text = f'{mention} {text}'
		await warn_message.edit(content=text)

	@classmethod
	def guild_trace(cls, guild: Guild, message: str) -> None:
		print(f'[guild {guild.id}|{guild.name}] {message}')