SileNce/discord_bot
1
0
Fork 0
Code Issues Projects Releases Packages Wiki Activity Actions

Improve whitelist feature
All checks were successful
CI / CI (push) Successful in 1m24s
Details

Browse source 
* 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.
This commit is contained in:
SileNce5k 2025-04-30 15:55:53 +02:00
parent c9d7a54e25
commit a066fd0662
Signed by: SileNce
GPG key ID: B0A142BB4291B204
6 changed files with 81 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

26
commands/admin/whitelist.js
View file

@ -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}) {
if(args < 2){
async execute({message, args, prefix, client}) {
if(args.length < 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)
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.")
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 that server.")
} else message.channel.send("Could not whitelist the server with that command. Check the logs.")
}
};

28
commands/misc/chat.js
View file

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

8
server.js
View file

@ -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();
}
@ -18,8 +19,8 @@ async function checkAndConvertJSONToSQL(){
}
const createDatabaseTables = require('./server/createDatabaseTables');
const createLastfmTable = require('./server/createLastfmTable');
const createWhitelistTable = require('./server/createWhitelistTable.js')
createWhitelistTable();
const createAndLoadWhitelistTable = require('./server/createAndLoadWhitelistTable.js')
createLastfmTable();
checkAndConvertJSONToSQL();
const { Collection, Client, GatewayIntentBits, Partials } = require('discord.js');
@ -43,8 +44,11 @@ const {
client.settings = new Collection();
client.commands = new Collection();
client.serverPrefixes = new Collection();
client.whitelist = new Collection();
createAndLoadWhitelistTable(client.whitelist);
client.settings.set("presenceType", presenceType);
client.settings.set("presenceText", presenceText);

48
server/createAndLoadWhitelistTable.js Normal file
View file

@ -0,0 +1,48 @@
const sqlite3 = require('sqlite3').verbose();
module.exports = async function (clientWhitelist) {
const db = new sqlite3.Database('data/database.db');
await new Promise ((resolve, reject)=>{
db.run(
`CREATE TABLE IF NOT EXISTS whitelist (
Id INTEGER PRIMARY KEY AUTOINCREMENT,
serverId TEXT,
command TEXT,
dateAdded INTEGER)`,
(err) => {
if (err) {
console.error(`Error while creating table 'whitelist': ${err}`);
reject(err);
} else {
console.log("Table 'whitelist' created successfully.");
resolve();
}
db.close();
}
);
})
loadWhitelist(clientWhitelist);
}
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])
}
});
}

22
server/createWhitelistTable.js
View file

@ -1,22 +0,0 @@
const sqlite3 = require('sqlite3').verbose();
module.exports = async function () {
const db = new sqlite3.Database('data/database.db');
return new Promise ((resolve, reject)=>{
db.run(
`CREATE TABLE IF NOT EXISTS whitelist (
serverId TEXT PRIMARY KEY,
command TEXT,
dateAdded INTEGER)`,
(err) => {
if (err) {
console.error(`Error while creating table 'whitelist': ${err}`);
reject(err);
} else {
console.log("Table 'whitelist' created successfully.");
resolve();
}
db.close();
}
);
})
}

20
server/message.js
View file

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