From ba88243deae95d7db6c59b744a9b0a21a2e8cdbb Mon Sep 17 00:00:00 2001 From: SileNce5k Date: Wed, 30 Apr 2025 15:55:53 +0200 Subject: [PATCH] Improve whitelist feature * Use a primary autoincrementing primary key because primary keys need to be unique * Cache whitelist on bot startup / only read from database once * "Externalize" whitelist checking to the messageCreate function. --- commands/admin/whitelist.js | 22 ++++++++++++++++++---- commands/misc/chat.js | 28 ++-------------------------- server.js | 29 +++++++++++++++++++++++++++++ server/createWhitelistTable.js | 3 ++- server/message.js | 20 +++++--------------- 5 files changed, 56 insertions(+), 46 deletions(-) diff --git a/commands/admin/whitelist.js b/commands/admin/whitelist.js index d915114..9dcc072 100644 --- a/commands/admin/whitelist.js +++ b/commands/admin/whitelist.js @@ -2,15 +2,22 @@ const sqlite3 = require('sqlite3').verbose(); module.exports = { name: 'whitelist', - description: 'Whitelist different commands that need whitelisting.', + description: 'Whitelist a command in a specific server.', admin: true, - async execute({message, args, prefix}) { + async execute({message, args, prefix, client}) { if(args < 2){ message.channel.send(`You need to supply the command and channel that you want to whitelist\n\`${prefix}whitelist \``); return; } let command = args[0]; let guild = args[1]; + if(guild === "this") guild = message.guild.id; + // TODO: Add ability to remove server from whitelist. + const whitelistedCommands = client.whitelist.get("guild"); + if(whitelistedCommands && whitelistedCommands.includes(command)){ + message.channel.send("Command is already whitelisted in that server.") + return; + } // TODO: First check if the bot has access to that guild before whitelisting. const databasePath = 'data/database.db' const db = new sqlite3.Database(databasePath) @@ -35,8 +42,15 @@ module.exports = { } }); }) - if(!err) + if(!err){ + if(whitelistedCommands){ + whitelistedCommands.push(command); + client.whitelist.set(guild, whitelistedCommands); + }else { + client.whitelist.set(guild, command); + } message.channel.send("Command has been whitelisted in this server.") - else message.channel.send("Could not whitelist the server with that command. Check the logs.") + + } else message.channel.send("Could not whitelist the server with that command. Check the logs.") } }; \ No newline at end of file diff --git a/commands/misc/chat.js b/commands/misc/chat.js index 3ef4c94..e01a4dc 100644 --- a/commands/misc/chat.js +++ b/commands/misc/chat.js @@ -1,36 +1,11 @@ require("dotenv").config(); -const sqlite3 = require('sqlite3').verbose(); module.exports = { name: 'chat', description: 'A chat command that uses an LLM to answer your prompts (server must be whitelisted)', + needsWhitelist: true, async execute({ message, args }) { - // TODO: Externalize the whitelist checking into a the message function with a variable export. And cache it. - // Have a global collection that gets created on launch from the database, the collection uses the guild ID as a key and has an array of commands that are whitelisted. - // Just write to both the database and this collection when a new command gets whitelisted. - const db = new sqlite3.Database('data/database.db'); - let isWhitelisted = false; - await new Promise((resolve, reject) => { - db.get(`SELECT * FROM whitelist WHERE serverId = ? AND command = ?`, [message.guild.id, this.name], - function (error, row){ - if(error){ - console.error(error); - resolve(""); - }else{ - if(row === undefined){ - resolve(); - }else { - isWhitelisted = true; - resolve(); - } - } - }) - }) - if(!isWhitelisted){ - message.channel.send("This server is not whitelisted. The bot admin needs to whitelist the server for this command to work."); - return; - } if(args.length === 0){ message.channel.send("You have to set your prompt in the arguments"); return; @@ -39,6 +14,7 @@ module.exports = { let answer = ""; const initialMessage = await message.channel.send("Generating response... This may take a moment.") message.channel.sendTyping(); + // TODO: More configuration. Have a basic setup but allow setting system prompt, max tokens and model. await fetch(`https://openrouter.ai/api/v1/chat/completions`, { method: `POST`, headers: { diff --git a/server.js b/server.js index 51d2f55..4e2c6bd 100644 --- a/server.js +++ b/server.js @@ -1,6 +1,7 @@ const fs = require('fs'); const createInitialConfig = require("./util/createInitialConfig") const convertJSONToSQL = require('./util/timer/convertJSONToSQL'); +const sqlite3 = require('sqlite3').verbose(); if(!fs.existsSync("./data/config.json")) { createInitialConfig(); } @@ -43,8 +44,36 @@ const { client.settings = new Collection(); client.commands = new Collection(); client.serverPrefixes = new Collection(); +client.whitelist = new Collection(); + +async function loadWhitelist(clientWhitelist) { + const db = new sqlite3.Database('data/database.db'); + let rows = await new Promise((resolve) => { + db.all(`SELECT * FROM whitelist`, function (error, rows){ + if(error){ + console.error("Failed to read whitelist table") + console.error(error); + resolve([]); + }else{ + resolve(rows); + } + }); + }); + rows.forEach(row => { + if(clientWhitelist.has(row.serverId)){ + let oldEntry = clientWhitelist.get(row.serverId); + oldEntry.push(row.command) + clientWhitelist.set(row.serverId, oldEntry) + }else { + clientWhitelist.set(row.serverId, [row.command]) + } + }); +} + +loadWhitelist(client.whitelist); + client.settings.set("presenceType", presenceType); client.settings.set("presenceText", presenceText); diff --git a/server/createWhitelistTable.js b/server/createWhitelistTable.js index b65319d..f191503 100644 --- a/server/createWhitelistTable.js +++ b/server/createWhitelistTable.js @@ -4,7 +4,8 @@ module.exports = async function () { return new Promise ((resolve, reject)=>{ db.run( `CREATE TABLE IF NOT EXISTS whitelist ( - serverId TEXT PRIMARY KEY, + Id INTEGER PRIMARY KEY AUTOINCREMENT, + serverId TEXT, command TEXT, dateAdded INTEGER)`, (err) => { diff --git a/server/message.js b/server/message.js index 9178afc..5e1b498 100644 --- a/server/message.js +++ b/server/message.js @@ -1,6 +1,3 @@ -const fs = require('fs'); -const customReplaceWithVariables = require('../util/custom_commands/customReplaceWithVariables'); - module.exports = function(client, owners, message, globalPrefix){ let prefix = globalPrefix; let serverPrefix = client.serverPrefixes.get(message.guild.id); @@ -22,19 +19,12 @@ module.exports = function(client, owners, message, globalPrefix){ const commandName = args.shift().toLowerCase(); const command = client.commands.get(commandName); - if (!command){ - const customPath = './data/customCommands.json'; - if(fs.existsSync(customPath)){ - let json = fs.readFileSync(customPath, 'utf8'); - let customCommands = JSON.parse(json) - customCommands.forEach(function (customCommand) { - if (customCommand.customName === commandName) { - let customMessage = customReplaceWithVariables(customCommand.customMessage, message, prefix, globalPrefix) - message.channel.send(customMessage) - } - }); + if(command.needsWhitelist){ + let isWhitelisted = client.whitelist.get(message.guild.id)?.includes(command.name); + if(!isWhitelisted){ + message.channel.send(`\`${command.name}\` is not whitelisted in this server. The bot admin needs to whitelist the command in this server for it to work`) + return; } - return; } if (command.admin && owners.indexOf(message.author.id.toString()) == -1) return; try {