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

"""
Base cog class and helper classes.
"""
from datetime import datetime, timedelta, timezone
from typing import Optional

from discord import Guild, Member, Message, RawReactionActionEvent, TextChannel
from discord.abc import GuildChannel
from discord.ext import commands

from config import CONFIG
from rocketbot.bot import Rocketbot
from rocketbot.botmessage import BotMessage, BotMessageReaction
from rocketbot.cogsetting import CogSetting
from rocketbot.collections import AgeBoundDict
from rocketbot.storage import Storage
from rocketbot.utils import bot_log

class WarningContext:
	def __init__(self, member: Member, warn_time: datetime):
		self.member = member
		self.last_warned = warn_time

class BaseCog(commands.Cog):
	STATE_KEY_RECENT_WARNINGS = "BaseCog.recent_warnings"

	"""
	Superclass for all Rocketbot cogs. Provides lots of conveniences for
	common tasks.
	"""
	def __init__(self, bot):
		self.bot: Rocketbot = bot
		self.are_settings_setup = False
		self.settings = []

	# Config

	@classmethod
	def get_cog_default(cls, key: str):
		"""
		Convenience method for getting a cog configuration default from
		`CONFIG['cogs'][<cog_name>][<key>]`. These values are used for
		CogSettings when no guild-specific value is configured yet.
		"""
		cogs: dict = CONFIG['cog_defaults']
		cog = cogs.get(cls.__name__)
		if cog is None:
			return None
		return cog.get(key)

	def add_setting(self, setting: CogSetting) -> None:
		"""
		Called by a subclass in __init__ to register a mod-configurable
		guild setting. A "get" and "set" command will be generated. If the
		setting is named "enabled" (exactly) then "enable" and "disable"
		commands will be created instead which set the setting to True/False.

		If the cog has a command group it will be detected automatically and
		the commands added to that. Otherwise, the commands will be added at
		the top level.

		Changes to settings can be detected by overriding `on_setting_updated`.
		"""
		self.settings.append(setting)

	@classmethod
	def get_guild_setting(cls,
			guild: Guild,
			setting: CogSetting,
			use_cog_default_if_not_set: bool = True):
		"""
		Returns the configured value for a setting for the given guild. If no
		setting is configured the default for the cog will be returned,
		unless the optional `use_cog_default_if_not_set` is `False`, then
		`None` will be returned.
		"""
		key = f'{cls.__name__}.{setting.name}'
		value = Storage.get_config_value(guild, key)
		if value is None and use_cog_default_if_not_set:
			value = cls.get_cog_default(setting.name)
		return value

	@classmethod
	def set_guild_setting(cls,
			guild: Guild,
			setting: CogSetting,
			new_value) -> None:
		"""
		Manually sets a setting for the given guild. BaseCog creates "get" and
		"set" commands for guild administrators to configure values themselves,
		but this method can be used for hidden settings from code. A ValueError
		will be raised if the new value does not pass validation specified in
		the CogSetting.
		"""
		setting.validate_value(new_value)
		key = f'{cls.__name__}.{setting.name}'
		Storage.set_config_value(guild, key, new_value)

	@commands.Cog.listener()
	async def on_ready(self):
		"""Event listener"""
		if not self.are_settings_setup:
			self.are_settings_setup = True
			CogSetting.set_up_all(self, self.bot, self.settings)

	async def on_setting_updated(self, guild: Guild, setting: CogSetting) -> None:
		"""
		Subclass override point for being notified when a CogSetting is edited.
		"""

	# Warning squelch

	def was_warned_recently(self, member: Member) -> bool:
		"""
		Tests if a given member was included in a mod warning message recently.
		Used to suppress redundant messages. Should be checked before pinging
		mods for relatively minor warnings about single users, but warnings
		about larger threats involving several members (e.g. join raids) should
		issue warnings regardless. Call record_warning or record_warnings after
		triggering a mod warning.
		"""
		recent_warns: AgeBoundDict[int, WarningContext, datetime, timedelta] = Storage.get_state_value(member.guild,
			BaseCog.STATE_KEY_RECENT_WARNINGS)
		if recent_warns is None:
			return False
		context: WarningContext = recent_warns.get(member.id)
		if context is None:
			return False
		squelch_warning_seconds: int = CONFIG['squelch_warning_seconds']
		return datetime.now() - context.last_warned < timedelta(seconds=squelch_warning_seconds)

	def record_warning(self, member: Member):
		"""
		Records that mods have been warned about a member and do not need to be
		warned about them again for a short while.
		"""
		recent_warns: AgeBoundDict[int, WarningContext, datetime, timedelta] = Storage.get_state_value(member.guild,
			BaseCog.STATE_KEY_RECENT_WARNINGS)
		if recent_warns is None:
			recent_warns = AgeBoundDict(timedelta(seconds=CONFIG['squelch_warning_seconds']),
				lambda i, context : context.last_warned)
			Storage.set_state_value(member.guild, BaseCog.STATE_KEY_RECENT_WARNINGS, recent_warns)
		context: WarningContext = recent_warns.get(member.id)
		if context is None:
			context = WarningContext(member, datetime.now())
			recent_warns[member.id] = context
		else:
			context.last_warned = datetime.now()

	def record_warnings(self, members: list):
		"""
		Records that mods have been warned about some members and do not need to
		be warned about them again for a short while.
		"""
		for member in members:
			self.record_warning(member)

	# Bot message handling

	@classmethod
	def __bot_messages(cls, guild: Guild) -> AgeBoundDict[int, BotMessage, datetime, timedelta]:
		bm: AgeBoundDict[int, BotMessage, datetime, timedelta] = Storage.get_state_value(guild, 'bot_messages')
		if bm is None:
			far_future = datetime.now(timezone.utc) + timedelta(days=1000)
			bm = AgeBoundDict(timedelta(seconds=600),
				lambda k, v : v.message_sent_at() or far_future)
			Storage.set_state_value(guild, 'bot_messages', bm)
		return bm

	async def post_message(self, message: BotMessage) -> bool:
		"""
		Posts a BotMessage to a guild. Returns whether it was successful. If
		the caller wants to listen to reactions they should be added before
		calling this method. Listen to reactions by overriding `on_mod_react`.
		"""
		message.source_cog = self
		await message.update()
		return message.is_sent()

	@commands.Cog.listener()
	async def on_raw_reaction_add(self, payload: RawReactionActionEvent):
		"""Event handler"""
		# Avoid any unnecessary requests. Gets called for every reaction
		# multiplied by every active cog.
		if payload.user_id == self.bot.user.id:
			# Ignore bot's own reactions
			return

		guild: Guild = self.bot.get_guild(payload.guild_id) or await self.bot.fetch_guild(payload.guild_id)
		if guild is None:
			# Possibly a DM
			return

		guild_messages: dict[int, BotMessage] = Storage.get_bot_messages(guild)
		bot_message = guild_messages.get(payload.message_id)
		if bot_message is None:
			# Unknown message (expired or was never tracked)
			return
		if self is not bot_message.source_cog:
			# Belongs to a different cog
			return
		reaction = bot_message.reaction_for_emoji(payload.emoji)
		if reaction is None or not reaction.is_enabled:
			# Can't use this reaction with this message
			return

		g_channel: GuildChannel = guild.get_channel(payload.channel_id) or await guild.fetch_channel(payload.channel_id)
		if g_channel is None:
			# Possibly a DM
			return
		if not isinstance(g_channel, TextChannel):
			return
		channel: TextChannel = g_channel
		member: Member = payload.member
		if member is None:
			return
		if not channel.permissions_for(member).ban_members:
			# Not a mod (could make permissions configurable per BotMessageReaction some day)
			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

		await self.on_mod_react(bot_message, reaction, member)

	async def on_mod_react(self,
			bot_message: BotMessage,
			reaction: BotMessageReaction,
			reacted_by: Member) -> None:
		"""
		Subclass override point for receiving mod reactions to bot messages sent
		via `post_message()`.
		"""

	# Helpers

	@classmethod
	def log(cls, guild: Optional[Guild], message) -> None:
		"""
		Writes a message to the console. Intended for significant events only.
		"""
		bot_log(guild, cls, message)