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

@@ -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 <command> <server_id>\``);
             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.")
     }
 };

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: {

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);
 

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) => {

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){
+	if(command.needsWhitelist){
-		const customPath = './data/customCommands.json';
+		let isWhitelisted = client.whitelist.get(message.guild.id)?.includes(command.name);
-		if(fs.existsSync(customPath)){
+		if(!isWhitelisted){
-			let json = fs.readFileSync(customPath, 'utf8');
+			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`)
-			let customCommands = JSON.parse(json)
+			return;
-			customCommands.forEach(function (customCommand) {
-				if (customCommand.customName === commandName) {
-					let customMessage = customReplaceWithVariables(customCommand.customMessage, message, prefix, globalPrefix)
-					message.channel.send(customMessage)
-				}
-			});
 		}
-		return;
 	}
 	if (command.admin && owners.indexOf(message.author.id.toString()) == -1) return;
 	try {