Skip to content

Commit

Permalink
Add image support for bad-words/auto-responses
Browse files Browse the repository at this point in the history
  • Loading branch information
JulianVennen committed Jul 8, 2024
1 parent f359732 commit 6855347
Show file tree
Hide file tree
Showing 23 changed files with 675 additions and 218 deletions.
76 changes: 76 additions & 0 deletions src/apis/CloudVision.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
import config from '../bot/Config.js';
import vision from '@google-cloud/vision';
import logger from '../bot/Logger.js';
import {Collection} from 'discord.js';
import GuildSettings from '../settings/GuildSettings.js';

export class CloudVision {
#imageAnnotatorClient = null;
#imageTexts = new Collection();

get isEnabled() {
return config.data.googleCloud.vision?.enabled;
}

get annotatorClient() {
if (!this.isEnabled) {
return null;
}

return this.#imageAnnotatorClient ??= new vision.ImageAnnotatorClient({
credentials: config.data.googleCloud.credentials
});
}

/**
* Get all image attachments from a message
* @param {import('discord.js').Message} message
* @returns {import('discord.js').Collection<import('discord.js').Snowflake, import('discord.js').Attachment>}
*/
getImages(message) {
return message.attachments.filter(attachment => attachment.contentType?.startsWith('image/'));
}

/**
* Get text from images in a message
* @param {import('discord.js').Message} message
* @returns {Promise<string[]>}
*/
async getImageText(message) {
if (!this.isEnabled) {
return [];
}

const guildSettings = await GuildSettings.get(message.guild.id);
if (!guildSettings.isFeatureWhitelisted) {
return [];
}

if (this.#imageTexts.has(message.id)) {
return this.#imageTexts.get(message.id);
}

const texts = [];

for (const image of this.getImages(message).values()) {
try {
const [{textAnnotations}] = await this.annotatorClient.textDetection(image.url);
for (const annotation of textAnnotations) {
texts.push(annotation.description);
}
}
catch (error) {
await logger.error(error);
}
}

if (texts.length) {
this.#imageTexts.set(message.id, texts);
setTimeout(() => this.#imageTexts.delete(message.id), 5000);
}

return texts;
}
}

export default new CloudVision();
52 changes: 41 additions & 11 deletions src/automod/AutoModManager.js
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import {formatTime} from '../util/timeutils.js';
import RepeatedMessage from './RepeatedMessage.js';
import SafeSearch from './SafeSearch.js';
import logger from '../bot/Logger.js';
import cloudVision from '../apis/CloudVision.js';

export class AutoModManager {
#safeSearchCache;
Expand Down Expand Up @@ -178,24 +179,53 @@ export class AutoModManager {

const words = (/** @type {Collection<number, BadWord>} */ await BadWord.get(channel.id, message.guild.id))
.sort((a, b) => b.priority - a.priority);

for (let word of words.values()) {
if (word.matches(message)) {
const reason = 'Using forbidden words or phrases';
const comment = `(Filter ID: ${word.id})`;
await bot.delete(message, reason + ' ' + comment);
if (word.response !== 'disabled') {
await this.#sendWarning(message, word.getResponse());
}
if (word.punishment.action !== 'none') {
const member = new Member(message.author, message.guild);
await member.executePunishment(word.punishment, reason, comment);
}
if (word.matches(message.content)) {
await this.#deleteBadWordMessage(word, message);
return true;
}
}

if (!cloudVision.isEnabled || !(await GuildSettings.get(message.guild.id)).isFeatureWhitelisted) {
return false;
}

let texts = null;
for (let word of words.values()) {
if (word.enableVision && word.trigger.supportsImages()) {
texts ??= await cloudVision.getImageText(message);
for (const text of texts) {
if (word.matches(text)) {
await this.#deleteBadWordMessage(word, message);
return true;
}
}

}
}

return false;
}

/**
* @param {BadWord} word
* @param {import('discord.js').Message} message
* @returns {Promise<void>}
*/
async #deleteBadWordMessage(word, message) {
const reason = 'Using forbidden words or phrases';
const comment = `(Filter ID: ${word.id})`;
await bot.delete(message, reason + ' ' + comment);
if (word.response !== 'disabled') {
await this.#sendWarning(message, word.getResponse());
}
if (word.punishment.action !== 'none') {
const member = new Member(message.author, message.guild);
await member.executePunishment(word.punishment, reason, comment);
}
}

/**
* @param {import('discord.js').Message} message
* @return {Promise<boolean>} has the message been deleted
Expand Down
21 changes: 4 additions & 17 deletions src/automod/SafeSearch.js
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
import config from '../bot/Config.js';
import GuildSettings from '../settings/GuildSettings.js';
import vision from '@google-cloud/vision';
import Cache from '../bot/Cache.js';
import Request from '../bot/Request.js';
import database from '../bot/Database.js';
import logger from '../bot/Logger.js';
import cloudVision from '../apis/CloudVision.js';

const CACHE_DURATION = 60 * 60 * 1000;

Expand All @@ -17,25 +16,13 @@ export default class SafeSearch {
*/
#requesting = new Map();

constructor() {
if (this.isEnabled) {
this.annotatorClient = new vision.ImageAnnotatorClient({
credentials: config.data.googleCloud.credentials
});
}
}

get isEnabled() {
return config.data.googleCloud.vision?.enabled;
}

/**
* is safe search filtering enabled in this guild
* @param {import('discord.js').Guild} guild
* @return {Promise<boolean>}
*/
async isEnabledInGuild(guild) {
if (!this.isEnabled) {
if (!cloudVision.isEnabled) {
return false;
}

Expand All @@ -50,7 +37,7 @@ export default class SafeSearch {
*/
async detect(message) {
/** @type {import('discord.js').Collection<string, import('discord.js').Attachment>} */
const images = message.attachments.filter(attachment => attachment.contentType?.startsWith('image/'));
const images = cloudVision.getImages(message);
if (!images.size) {
return null;
}
Expand Down Expand Up @@ -96,7 +83,7 @@ export default class SafeSearch {

let safeSearchAnnotation = null;
try {
[{safeSearchAnnotation}] = await this.annotatorClient.safeSearchDetection(image.url);
[{safeSearchAnnotation}] = await cloudVision.annotatorClient.safeSearchDetection(image.url);

if (safeSearchAnnotation) {
this.#cache.set(hash, safeSearchAnnotation, CACHE_DURATION);
Expand Down
10 changes: 7 additions & 3 deletions src/bot/Database.js
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ import logger from './Logger.js';
import config from './Config.js';
import CommentFieldMigration from '../database/migrations/CommentFieldMigration.js';
import {asyncFilter} from '../util/util.js';
import BadWordVisionMigration from '../database/migrations/BadWordVisionMigration.js';
import AutoResponseVisionMigration from '../database/migrations/AutoResponseVisionMigration.js';

export class Database {
/**
Expand Down Expand Up @@ -100,16 +102,18 @@ export class Database {
await this.query('CREATE TABLE IF NOT EXISTS `channels` (`id` VARCHAR(20) NOT NULL, `config` TEXT NOT NULL, PRIMARY KEY (`id`), `guildid` VARCHAR(20))');
await this.query('CREATE TABLE IF NOT EXISTS `guilds` (`id` VARCHAR(20) NOT NULL, `config` TEXT NOT NULL, PRIMARY KEY (`id`))');
await this.query('CREATE TABLE IF NOT EXISTS `users` (`id` VARCHAR(20) NOT NULL, `config` TEXT NOT NULL, PRIMARY KEY (`id`))');
await this.query('CREATE TABLE IF NOT EXISTS `responses` (`id` int PRIMARY KEY AUTO_INCREMENT, `guildid` VARCHAR(20) NOT NULL, `trigger` TEXT NOT NULL, `response` TEXT NOT NULL, `global` BOOLEAN NOT NULL, `channels` TEXT NULL DEFAULT NULL)');
await this.query('CREATE TABLE IF NOT EXISTS `badWords` (`id` int PRIMARY KEY AUTO_INCREMENT, `guildid` VARCHAR(20) NOT NULL, `trigger` TEXT NOT NULL, `punishment` TEXT NOT NULL, `response` TEXT NOT NULL, `global` BOOLEAN NOT NULL, `channels` TEXT NULL DEFAULT NULL, `priority` int NULL)');
await this.query('CREATE TABLE IF NOT EXISTS `responses` (`id` int PRIMARY KEY AUTO_INCREMENT, `guildid` VARCHAR(20) NOT NULL, `trigger` TEXT NOT NULL, `response` TEXT NOT NULL, `global` BOOLEAN NOT NULL, `channels` TEXT NULL DEFAULT NULL, `enableVision` BOOLEAN DEFAULT FALSE)');
await this.query('CREATE TABLE IF NOT EXISTS `badWords` (`id` int PRIMARY KEY AUTO_INCREMENT, `guildid` VARCHAR(20) NOT NULL, `trigger` TEXT NOT NULL, `punishment` TEXT NOT NULL, `response` TEXT NOT NULL, `global` BOOLEAN NOT NULL, `channels` TEXT NULL DEFAULT NULL, `priority` int NULL, `enableVision` BOOLEAN DEFAULT FALSE)');
await this.query('CREATE TABLE IF NOT EXISTS `moderations` (`id` int PRIMARY KEY AUTO_INCREMENT, `guildid` VARCHAR(20) NOT NULL, `userid` VARCHAR(20) NOT NULL, `action` VARCHAR(10) NOT NULL, `created` bigint NOT NULL, `value` int DEFAULT 0, `expireTime` bigint NULL DEFAULT NULL, `reason` TEXT, `comment` TEXT NULL DEFAULT NULL, `moderator` VARCHAR(20) NULL DEFAULT NULL, `active` BOOLEAN DEFAULT TRUE)');
await this.query('CREATE TABLE IF NOT EXISTS `confirmations` (`id` int PRIMARY KEY AUTO_INCREMENT, `data` TEXT NOT NULL, `expires` bigint NOT NULL)');
await this.query('CREATE TABLE IF NOT EXISTS `safeSearch` (`hash` CHAR(64) PRIMARY KEY, `data` TEXT NOT NULL)');
}

async getMigrations() {
return await asyncFilter([
new CommentFieldMigration(this)
new CommentFieldMigration(this),
new BadWordVisionMigration(this),
new AutoResponseVisionMigration(this),
], async migration => await migration.check());
}

Expand Down
41 changes: 35 additions & 6 deletions src/commands/settings/auto-response/AddAutoResponseCommand.js
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import AutoResponse from '../../../database/AutoResponse.js';
import ErrorEmbed from '../../../embeds/ErrorEmbed.js';
import colors from '../../../util/colors.js';
import {SELECT_MENU_OPTIONS_LIMIT} from '../../../util/apiLimits.js';
import config from '../../../bot/Config.js';

export default class AddAutoResponseCommand extends SubCommand {

Expand Down Expand Up @@ -40,14 +41,23 @@ export default class AddAutoResponseCommand extends SubCommand {
.setName('global')
.setDescription('Use auto-response in all channels')
.setRequired(false));

if (config.data.googleCloud.vision.enabled) {
builder.addBooleanOption(option => option
.setName('image-detection')
.setDescription('Respond to images containing text that matches the trigger')
.setRequired(false));
}

return super.buildOptions(builder);
}

async execute(interaction) {
const global = interaction.options.getBoolean('global') ?? false;
const type = interaction.options.getString('type') ?? 'include';
const global = interaction.options.getBoolean('global') ?? false,
type = interaction.options.getString('type') ?? 'include',
vision = interaction.options.getBoolean('image-detection') ?? false;

const confirmation = new Confirmation({global, type}, timeAfter('1 hour'));
const confirmation = new Confirmation({global, type, vision}, timeAfter('1 hour'));
await interaction.showModal(new ModalBuilder()
.setTitle(`Create new Auto-response of type ${type}`)
.setCustomId(`auto-response:add:${await confirmation.save()}`)
Expand Down Expand Up @@ -109,7 +119,8 @@ export default class AddAutoResponseCommand extends SubCommand {
[],
confirmation.data.type,
trigger,
response
response,
confirmation.data.vision,
);
}
else {
Expand Down Expand Up @@ -154,6 +165,7 @@ export default class AddAutoResponseCommand extends SubCommand {
confirmation.data.type,
confirmation.data.trigger,
confirmation.data.response,
confirmation.data.vision,
);
}

Expand All @@ -165,10 +177,27 @@ export default class AddAutoResponseCommand extends SubCommand {
* @param {string} type
* @param {string} trigger
* @param {string} response
* @param {?boolean} enableVision
* @return {Promise<*>}
*/
async create(interaction, global, channels, type, trigger, response) {
const result = await AutoResponse.new(interaction.guild.id, global, channels, type, trigger, response);
async create(
interaction,
global,
channels,
type,
trigger,
response,
enableVision,
) {
const result = await AutoResponse.new(
interaction.guild.id,
global,
channels,
type,
trigger,
response,
enableVision,
);
if (!result.success) {
return interaction.reply(ErrorEmbed.message(result.message));
}
Expand Down
Loading

0 comments on commit 6855347

Please sign in to comment.