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

"""
Cog for handling most ungrouped commands and basic behaviors.
"""
from datetime import datetime, timedelta, timezone
from typing import Optional

from discord import Interaction, Message, User
from discord.app_commands import command, default_permissions, guild_only, Transform
from discord.errors import DiscordException
from discord.ext.commands import Cog

from config import CONFIG
from rocketbot.bot import Rocketbot
from rocketbot.cogs.basecog import BaseCog, BotMessage
from rocketbot.utils import describe_timedelta, TimeDeltaTransformer
from rocketbot.storage import ConfigKey, Storage

class GeneralCog(BaseCog, name='General'):
	"""
	Cog for handling high-level bot functionality and commands. Should be the
	first cog added to the bot.
	"""
	def __init__(self, bot: Rocketbot):
		super().__init__(
			bot,
			config_prefix=None,
			name='',
			short_description='',
		)
		self.is_connected = False
		self.is_first_connect = True
		self.last_disconnect_time: Optional[datetime] = None
		self.noteworthy_disconnect_duration = timedelta(seconds=5)

	@Cog.listener()
	async def on_connect(self):
		"""Event handler"""
		if self.is_first_connect:
			self.log(None, 'Connected')
			self.is_first_connect = False
		else:
			disconnect_duration = datetime.now(
				timezone.utc) - self.last_disconnect_time if self.last_disconnect_time else None
			if disconnect_duration is not None and disconnect_duration > self.noteworthy_disconnect_duration:
				self.log(None, f'Reconnected after {disconnect_duration.total_seconds()} seconds')
		self.is_connected = True

	@Cog.listener()
	async def on_disconnect(self):
		"""Event handler"""
		self.last_disconnect_time = datetime.now(timezone.utc)
		# self.log(None, 'Disconnected')

	@Cog.listener()
	async def on_resumed(self):
		"""Event handler"""
		disconnect_duration = datetime.now(timezone.utc) - self.last_disconnect_time if self.last_disconnect_time else None
		if disconnect_duration is not None and disconnect_duration > self.noteworthy_disconnect_duration:
			self.log(None, f'Session resumed after {disconnect_duration.total_seconds()} seconds')

	@command(
		description='Posts a test warning',
		extras={
			'long_description': 'Tests whether a warning channel is configured for this ' + \
				'guild by posting a test warning. If a mod mention is ' + \
				'configured, that user/role will be tagged in the test warning.',
		},
	)
	@guild_only()
	@default_permissions(ban_members=True)
	async def test_warn(self, interaction: Interaction):
		"""Command handler"""
		if Storage.get_config_value(interaction.guild, ConfigKey.WARNING_CHANNEL_ID) is None:
			await interaction.response.send_message(
				f'{CONFIG["warning_emoji"]} No warning channel set!',
				ephemeral=True,
			)
		else:
			bm = BotMessage(
				interaction.guild,
				f'Test warning message (requested by {interaction.user.name})',
				type=BotMessage.TYPE_MOD_WARNING)
			await self.post_message(bm)
			await interaction.response.send_message(
				'Warning issued',
				ephemeral=True,
			)

	@command(
		description='Simple test reply',
		extras={
			'long_description': 'Replies to the command message. Useful to ensure the ' + \
				'bot is working properly.',
		},
	)
	async def hello(self, interaction: Interaction):
		"""Command handler"""
		await interaction.response.send_message(
			f'Hey, {interaction.user.name}!',
		 	ephemeral=True,
		)

	@command(
		description='Shuts down the bot',
		extras={
			'long_description': 'Causes the bot script to terminate. Only usable by a ' + \
				'user with server admin permissions.',
		},
	)
	@guild_only()
	@default_permissions(administrator=True)
	async def shutdown(self, interaction: Interaction):
		"""Command handler"""
		await interaction.response.send_message('👋', ephemeral=True)
		await self.bot.close()

	@command(
		description='Mass deletes messages',
		extras={
			'long_description': 'Deletes recent messages by the given user. The user ' +
				'can be either an @ mention or a numeric user ID. The age is ' +
				'a duration, such as "30s", "5m", "1h30m". Only the most ' +
				'recent 100 messages in each channel are searched.',
			'usage': '<user:id|mention> <age:timespan>',
		},
	)
	@guild_only()
	@default_permissions(manage_messages=True)
	async def delete_messages(self, interaction: Interaction, user: User, age: Transform[timedelta, TimeDeltaTransformer]) -> None:
		"""Command handler"""
		member_id = user.id
		cutoff: datetime = datetime.now(timezone.utc) - age
		def predicate(message: Message) -> bool:
			return str(message.author.id) == member_id and message.created_at >= cutoff
		deleted_messages = []
		for channel in interaction.guild.text_channels:
			try:
				deleted_messages += await channel.purge(limit=100, check=predicate)
			except DiscordException:
				# XXX: Sloppily glossing over access errors instead of checking access
				pass
		await interaction.response.send_message(
			f'{CONFIG["success_emoji"]} Deleted {len(deleted_messages)} ' + \
			f'messages by {user.mention}> from the past {describe_timedelta(age)}.',
			ephemeral=True,
		)