feat(bots/discord): add source

bots/discord/src/events/api/disconnect.ts

@@ -0,0 +1,30 @@
+import { on } from '$utils/api/events'
+import { DisconnectReason, HumanizedDisconnectReason } from '@revanced/bot-shared'
+import { api, logger } from 'src/context'
+
+on('disconnect', (reason, msg) => {
+    if (reason === DisconnectReason.PlannedDisconnect && api.isStopping) return
+
+    const ws = api.client.ws
+    if (!ws.disconnected) ws.disconnect()
+
+    logger.fatal(
+        `Disconnected from the bot API ${
+            reason in HumanizedDisconnectReason
+                ? `because ${HumanizedDisconnectReason[reason as keyof typeof HumanizedDisconnectReason]}`
+                : 'for an unknown reason'
+        }`,
+    )
+
+    // TODO: move to config
+    if (api.disconnectCount >= 3) {
+        console.error(new Error('Disconnected from bot API too many times'))
+        // We don't want the process hanging
+        process.exit(1)
+    }
+
+    logger.info(
+        `Disconnected from bot API ${++api.disconnectCount} times (this time because: ${reason}, ${msg}), reconnecting again...`,
+    )
+    setTimeout(() => ws.connect(), 10000)
+})

bots/discord/src/events/api/ready.ts

@@ -0,0 +1,6 @@
+import { on } from '$utils/api/events'
+import { logger } from 'src/context'
+
+on('ready', () => {
+    logger.info('Connected to the bot API')
+})

bots/discord/src/events/discord/guildCreate.ts

@@ -0,0 +1,7 @@
+import { on } from '$utils/discord/events'
+import { leaveDisallowedGuild } from '$utils/discord/security'
+
+on('guildCreate', async ({ config }, guild) => {
+    if (config.allowedGuilds.includes(guild.id)) return
+    await leaveDisallowedGuild(guild)
+})

bots/discord/src/events/discord/interactionCreate/chat-commmand.ts

@@ -0,0 +1,78 @@
+import { createErrorEmbed } from '$utils/discord/embeds'
+import { on } from '$utils/discord/events'
+
+export default on('interactionCreate', async (context, interaction) => {
+    if (!interaction.isChatInputCommand()) return
+
+    const { logger, discord, config } = context
+    const command = discord.commands[interaction.commandName]
+
+    logger.debug(`Command ${interaction.commandName} being invoked by ${interaction.user.tag}`)
+
+    if (!command) {
+        logger.error(`Command ${interaction.commandName} not implemented but registered!!!`)
+        return void interaction.reply({
+            embeds: [
+                createErrorEmbed(
+                    'Command not implemented',
+                    'This command has not been implemented yet. Please report this to the developers.',
+                ),
+            ],
+            ephemeral: true,
+        })
+    }
+
+    const userIsOwner = config.owners.includes(interaction.user.id)
+
+    if ((command.ownerOnly ?? true) && !userIsOwner)
+        return void (await interaction.reply({
+            embeds: [createErrorEmbed('Massive skill issue', 'This command can only be used by the bot owners.')],
+            ephemeral: true,
+        }))
+
+    if (interaction.inGuild()) {
+        // Bot owners get bypass
+        if (command.memberRequirements && !userIsOwner) {
+            const { permissions = -1n, roles = [], mode } = command.memberRequirements
+
+            const member = await interaction.guild!.members.fetch(interaction.user.id)
+
+            const [missingPermissions, missingRoles] = [
+                // This command is an owner-only command (the user is not an owner, checked above)
+                permissions < 0n ||
+                    // or the user doesn't have the required permissions
+                    (permissions >= 0n && !interaction.memberPermissions.has(permissions)),
+
+                // If not:
+                !roles.some(x => member.roles.cache.has(x)),
+            ]
+
+            if ((mode === 'any' && missingPermissions && missingRoles) || missingPermissions || missingRoles)
+                return void interaction.reply({
+                    embeds: [
+                        createErrorEmbed(
+                            'Missing roles or permissions',
+                            "You don't have the required roles or permissions to use this command.",
+                        ),
+                    ],
+                    ephemeral: true,
+                })
+        }
+    }
+
+    try {
+        logger.debug(`Command ${interaction.commandName} being executed`)
+        await command.execute(context, interaction)
+    } catch (err) {
+        logger.error(`Error while executing command ${interaction.commandName}:`, err)
+        await interaction.reply({
+            embeds: [
+                createErrorEmbed(
+                    'An error occurred while executing this command',
+                    'Please report this to the developers.',
+                ),
+            ],
+            ephemeral: true,
+        })
+    }
+})

bots/discord/src/events/discord/interactionCreate/correct-response.ts

@@ -0,0 +1,111 @@
+import { handleUserResponseCorrection } from '$/utils/discord/messageScan'
+import { createErrorEmbed, createStackTraceEmbed, createSuccessEmbed } from '$utils/discord/embeds'
+import { on } from '$utils/discord/events'
+
+import type { ButtonInteraction, StringSelectMenuInteraction, TextBasedChannel } from 'discord.js'
+
+// No permission check required as it is already done when the user reacts to a bot response
+export default on('interactionCreate', async (context, interaction) => {
+    const {
+        logger,
+        database: db,
+        config: { messageScan: msConfig },
+    } = context
+
+    if (!msConfig?.humanCorrections) return
+    if (!interaction.isStringSelectMenu() && !interaction.isButton()) return
+    if (!interaction.customId.startsWith('cr_')) return
+
+    const [, key, action] = interaction.customId.split('_') as ['cr', string, 'select' | 'cancel' | 'delete']
+    if (!key || !action) return
+
+    const response = db.labeledResponses.get(key)
+    // If the message isn't saved in my DB (unrelated message)
+    if (!response)
+        return void (await interaction.reply({
+            content: "I don't recall having sent this response, so I cannot correct it.",
+            ephemeral: true,
+        }))
+
+    try {
+        // We're gonna pretend reactionChannel is a text-based channel, but it can be many more
+        // But `messages` should always exist as a property
+        const reactionGuild = await interaction.client.guilds.fetch(response.guild)
+        const reactionChannel = (await reactionGuild.channels.fetch(response.channel)) as TextBasedChannel | null
+        const reactionMessage = await reactionChannel?.messages.fetch(key)
+
+        if (!reactionMessage) {
+            await interaction.deferUpdate()
+            await interaction.message.edit({
+                content: null,
+                embeds: [
+                    createErrorEmbed(
+                        'Response not found',
+                        'Thank you for your feedback! Unfortunately, the response message could not be found (most likely deleted).',
+                    ),
+                ],
+                components: [],
+            })
+
+            return
+        }
+
+        const editMessage = (content: string, description?: string) =>
+            editInteractionMessage(interaction, reactionMessage.url, content, description)
+        const handleCorrection = (label: string) =>
+            handleUserResponseCorrection(context, response, reactionMessage, label, interaction.user)
+
+        if (response.correctedBy)
+            return await editMessage(
+                'Response already corrected',
+                'Thank you for your feedback! Unfortunately, this response has already been corrected by someone else.',
+            )
+
+        // We immediately know that the action is `select`
+        if (interaction.isStringSelectMenu()) {
+            const selectedLabel = interaction.values[0]!
+
+            await handleCorrection(selectedLabel)
+            await editMessage(
+                'Message being trained',
+                `Thank you for your feedback! I've edited the response according to the selected label (\`${selectedLabel}\`). The message is now being trained. 🎉`,
+            )
+        } else {
+            switch (action) {
+                case 'cancel':
+                    await editMessage('Canceled', 'You canceled this interaction. 😞')
+                    break
+                case 'delete':
+                    await handleCorrection(msConfig.humanCorrections.falsePositiveLabel)
+                    await editMessage(
+                        'Marked as false positive',
+                        'The response has been deleted and marked as a false positive. Thank you for your feedback. 🎉',
+                    )
+
+                    break
+            }
+        }
+    } catch (e) {
+        logger.error('Failed to handle correct response interaction:', e)
+        await interaction.reply({
+            embeds: [createStackTraceEmbed(e)],
+            ephemeral: true,
+        })
+    }
+})
+
+const editInteractionMessage = async (
+    interaction: StringSelectMenuInteraction | ButtonInteraction,
+    replyUrl: string,
+    title: string,
+    description?: string,
+) => {
+    if (!interaction.deferred) await interaction.deferUpdate()
+    await interaction.message.edit({
+        content: null,
+        embeds: [
+            createSuccessEmbed(title, `${description ?? ''}\n\n**⬅️ Back to bot response**: ${replyUrl}`.trimStart()),
+        ],
+        components: [],
+    })
+}

bots/discord/src/events/discord/messageCreate/scan.ts

@@ -0,0 +1,70 @@
+import { MessageScanLabeledResponseReactions } from '$/constants'
+import { getResponseFromContent, shouldScanMessage } from '$/utils/discord/messageScan'
+import { createMessageScanResponseEmbed } from '$utils/discord/embeds'
+import { on } from '$utils/discord/events'
+
+on('messageCreate', async (ctx, msg) => {
+    const {
+        api,
+        config: { messageScan: config },
+        database: db,
+        logger,
+    } = ctx
+
+    if (!config || !config.responses) return
+    if (!shouldScanMessage(msg, config)) return
+
+    if (msg.content.length) {
+        logger.debug(`Classifying message ${msg.id}`)
+
+        const { response, label } = await getResponseFromContent(msg.content, ctx)
+
+        if (response) {
+            logger.debug('Response found')
+
+            const reply = await msg.reply({
+                embeds: [createMessageScanResponseEmbed(response)],
+            })
+
+            if (label)
+                db.labeledResponses.save({
+                    reply: reply.id,
+                    channel: reply.channel.id,
+                    guild: reply.guild.id,
+                    referenceMessage: msg.id,
+                    label,
+                    text: msg.content,
+                })
+
+            if (label) {
+                for (const reaction of Object.values(MessageScanLabeledResponseReactions)) {
+                    await reply.react(reaction)
+                }
+            }
+        }
+    }
+
+    if (msg.attachments.size > 0) {
+        logger.debug(`Classifying message attachments for ${msg.id}`)
+
+        for (const attachment of msg.attachments.values()) {
+            if (attachment.contentType && !config.allowedAttachmentMimeTypes.includes(attachment.contentType)) continue
+
+            try {
+                const { text: content } = await api.client.parseImage(attachment.url)
+                const { response } = await getResponseFromContent(content, ctx)
+
+                if (response) {
+                    logger.debug(`Response found for attachment: ${attachment.url}`)
+                    await msg.reply({
+                        embeds: [createMessageScanResponseEmbed(response)],
+                    })
+
+                    break
+                }
+            } catch {
+                logger.error(`Failed to parse image: ${attachment.url}`)
+            }
+        }
+    }
+})

bots/discord/src/events/discord/messageReactionAdd/correct-response.ts

@@ -0,0 +1,130 @@
+import { MessageScanLabeledResponseReactions as Reactions } from '$/constants'
+import { createErrorEmbed, createStackTraceEmbed, createSuccessEmbed } from '$/utils/discord/embeds'
+import { on } from '$/utils/discord/events'
+
+import {
+    ActionRowBuilder,
+    ButtonBuilder,
+    ButtonStyle,
+    StringSelectMenuBuilder,
+    StringSelectMenuOptionBuilder,
+} from 'discord.js'
+
+import { handleUserResponseCorrection } from '$/utils/discord/messageScan'
+import type { ConfigMessageScanResponseLabelConfig } from 'config.example'
+
+const PossibleReactions = Object.values(Reactions) as string[]
+
+on('messageReactionAdd', async (context, rct, user) => {
+    if (user.bot) return
+
+    const { database: db, logger, config } = context
+    const { messageScan: msConfig } = config
+
+    // If there's no config, we can't do anything
+    if (!msConfig?.humanCorrections) return
+
+    const reaction = await rct.fetch()
+    const reactionMessage = await reaction.message.fetch()
+
+    if (reactionMessage.author.id !== reaction.client.user!.id) return
+    if (!PossibleReactions.includes(reaction.emoji.name!)) return
+
+    if (reactionMessage.inGuild() && msConfig.humanCorrections.memberRequirements) {
+        const {
+            memberRequirements: { roles, permissions },
+        } = msConfig.humanCorrections
+
+        if (!roles && !permissions)
+            return void logger.warn(
+                'No member requirements specified for human corrections, ignoring this request for security reasons',
+            )
+
+        const member = await reactionMessage.guild.members.fetch(user.id)
+
+        if (
+            permissions &&
+            !member.permissions.has(permissions) &&
+            roles &&
+            !roles.some(role => member.roles.cache.has(role))
+        )
+            return
+        // User is not owner, and not included in allowUsers
+    } else if (!config.owners.includes(user.id) && !msConfig.humanCorrections.allowUsers?.includes(user.id)) return
+
+    // Sanity check
+    const response = db.labeledResponses.get(rct.message.id)
+    if (!response || response.correctedBy) return
+
+    const handleCorrection = (label: string) =>
+        handleUserResponseCorrection(context, response, reactionMessage, label, user)
+
+    try {
+        if (reaction.emoji.name === Reactions.train) {
+            // Bot is right, nice!
+
+            await handleCorrection(response.label)
+            await user.send({ embeds: [createSuccessEmbed('Trained message', 'Thank you for your feedback.')] })
+        } else if (reaction.emoji.name === Reactions.edit) {
+            // Bot is wrong :(
+
+            const labels = msConfig.responses!.flatMap(r =>
+                r.triggers.filter((t): t is ConfigMessageScanResponseLabelConfig => 'label' in t).map(t => t.label),
+            )
+
+            const componentPrefix = `cr_${reactionMessage.id}`
+            const select = new StringSelectMenuBuilder().setCustomId(`${componentPrefix}_select`)
+
+            for (const label of labels) {
+                const opt = new StringSelectMenuOptionBuilder().setLabel(label).setValue(label)
+
+                if (label === response.label) {
+                    opt.setDefault(true)
+                    opt.setLabel(`${label} (current)`)
+                    opt.setDescription('This is the current label of the message')
+                }
+
+                select.addOptions(opt)
+            }
+
+            const rows = [
+                new ActionRowBuilder<StringSelectMenuBuilder>().addComponents(select),
+                new ActionRowBuilder<ButtonBuilder>().addComponents(
+                    new ButtonBuilder()
+                        .setEmoji('⬅️')
+                        .setLabel('Cancel')
+                        .setStyle(ButtonStyle.Secondary)
+                        .setCustomId(`${componentPrefix}_cancel`),
+                    new ButtonBuilder()
+                        .setEmoji(Reactions.delete)
+                        .setLabel('Delete (mark as false positive)')
+                        .setStyle(ButtonStyle.Danger)
+                        .setCustomId(`${componentPrefix}_delete`),
+                ),
+            ]
+
+            await user.send({
+                content: 'Please pick the right label for the message (you can only do this once!)',
+                components: rows,
+            })
+        } else if (reaction.emoji.name === Reactions.delete) {
+            await handleCorrection(msConfig.humanCorrections.falsePositiveLabel)
+            await user.send({ content: 'The response has been deleted and marked as a false positive.' })
+        }
+    } catch (e) {
+        logger.error('Failed to correct response:', e)
+        user.send({
+            embeds: [createStackTraceEmbed(e)],
+        }).catch(() => {
+            reactionMessage.reply({
+                content: `<@${user.id}>`,
+                embeds: [
+                    createErrorEmbed(
+                        'Enable your DMs!',
+                        'I cannot send you messages. Please enable your DMs to use this feature.',
+                    ),
+                ],
+            })
+        })
+    }
+})

bots/discord/src/events/discord/ready.ts

@@ -0,0 +1,8 @@
+import { on } from 'src/utils/discord/events'
+
+export default on('ready', ({ logger }, client) => {
+    logger.info(`Connected to Discord API, logged in as ${client.user.displayName} (@${client.user.tag})!`)
+    logger.info(
+        `Bot is in ${client.guilds.cache.size} guilds, if this is not expected, please run the /leave-unknowns command`,
+    )
+})