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.
	"""

	shared: Optional['GeneralCog'] = None

	def __init__(self, bot: Rocketbot):
		super().__init__(
			bot,
			config_prefix=None,
			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)
		GeneralCog.shared = self

	@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': 'If a warning channel is configured, it will be posted '
								'there. If a warning role/user is configured, they will be '
								'tagged in the message.',
		},
	)
	@guild_only()
	@default_permissions(ban_members=True)
	async def test_warn(self, interaction: Interaction):
		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,
			)
		self.log(interaction.guild, f'{interaction.user.name} used /test_warn')

	@command(
		description='Responds to the user with a greeting.',
		extras={
			'long_description': 'Useful for checking bot responsiveness. Message '
								'is only visible to the user.',
		},
	)
	async def hello(self, interaction: Interaction):
		await interaction.response.send_message(
			f'Hey, {interaction.user.name}!',
		 	ephemeral=True,
		)
		self.log(interaction.guild, f'{interaction.user.name} used /hello')

	@command(
		description='Shuts down the bot.',
		extras={
			'long_description': 'For emergency use if the bot gains sentience. Only usable '
								'by a server administrator.',
		},
	)
	@guild_only()
	@default_permissions(administrator=True)
	async def shutdown(self, interaction: Interaction):
		"""Command handler"""
		await interaction.response.send_message('👋', ephemeral=True)
		self.log(interaction.guild, f'{interaction.user.name} used /shutdown')
		await self.bot.close()

	@command(
		description='Mass deletes recent messages by a user.',
		extras={
			'long_description': 'The age is a duration, such as "30s", "5m", "1h30m", "7d". '
								'Only the most recent 100 messages in each channel are searched.\n\n'
								"The author can be a numeric ID if they aren't showing up in autocomplete.",
			'usage': '<user:id|mention> <age:timespan>',
		},
	)
	@guild_only()
	@default_permissions(manage_messages=True)
	async def delete_messages(self, interaction: Interaction, author: User, age: Transform[timedelta, TimeDeltaTransformer]) -> None:
		"""
		Mass deletes messages.

		Parameters
		----------
		interaction: :class:`Interaction`
		author: :class:`User`
			author of messages to delete
		age: :class:`timedelta`
			maximum age of messages to delete (e.g. 30s, 5m, 1h30s, 7d)
		"""
		member_id = author.id
		cutoff: datetime = datetime.now(timezone.utc) - age

		# Finding and deleting messages takes time but interaction needs a timely acknowledgement.
		resp = await interaction.response.defer(ephemeral=True, thinking=True)

		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

		if len(deleted_messages) > 0:
			await resp.resource.edit(
				content=f'{CONFIG["success_emoji"]} Deleted {len(deleted_messages)} '
						f'messages by {author.mention} from the past {describe_timedelta(age)}.',
			)
		else:
			await resp.resource.edit(
				content=f'{CONFIG["success_emoji"]} No messages found for {author.mention} '
						f'from the past {describe_timedelta(age)}.',
			)
		self.log(interaction.guild, f'{interaction.user.name} used /delete_messages {author.id} {age}')