feat: limit whitelisted drops to 3/day and export drop methods

src/database.js

@@ -34,6 +34,20 @@ db.exec(`
         key           TEXT PRIMARY KEY,
         value         TEXT NOT NULL
     );
+
+    CREATE TABLE IF NOT EXISTS whitelist (
+        user_id       TEXT PRIMARY KEY,
+        added_by      TEXT NOT NULL,
+        added_at      DATETIME DEFAULT CURRENT_TIMESTAMP
+    );
+
+    CREATE TABLE IF NOT EXISTS drop_log (
+        id            INTEGER PRIMARY KEY AUTOINCREMENT,
+        user_id       TEXT NOT NULL,
+        title         TEXT NOT NULL,
+        channel_id    TEXT NOT NULL,
+        dropped_at    DATETIME DEFAULT CURRENT_TIMESTAMP
+    );
 `);
 
 // ── Prepared Statements ───────────────────────────────────────
@@ -60,6 +74,17 @@ const stmts = {
     setState: db.prepare('INSERT OR REPLACE INTO bot_state (key, value) VALUES (?, ?)'),
     getState: db.prepare('SELECT value FROM bot_state WHERE key = ?'),
     delState: db.prepare('DELETE FROM bot_state WHERE key = ?'),
+
+    // Whitelist
+    addWhitelist: db.prepare('INSERT OR REPLACE INTO whitelist (user_id, added_by) VALUES (?, ?)'),
+    removeWhitelist: db.prepare('DELETE FROM whitelist WHERE user_id = ?'),
+    getWhitelist: db.prepare('SELECT * FROM whitelist WHERE user_id = ?'),
+    getAllWhitelist: db.prepare('SELECT * FROM whitelist'),
+
+    // Drop log
+    logDrop: db.prepare('INSERT INTO drop_log (user_id, title, channel_id) VALUES (?, ?, ?)'),
+    getDropCount24h: db.prepare(`SELECT COUNT(*) as count FROM drop_log WHERE user_id = ? AND dropped_at > datetime('now', '-24 hours')`),
+    getLastDrop: db.prepare(`SELECT dropped_at FROM drop_log WHERE user_id = ? ORDER BY dropped_at DESC LIMIT 1`),
 };
 
 module.exports = { db, stmts };

src/events/messageCreate.js

@@ -1,7 +1,9 @@
 const { ChannelType } = require('discord.js');
-const { errorEmbed, infoEmbed } = require('../utils/embeds');
+const { errorEmbed, infoEmbed, successEmbed, createEmbed } = require('../utils/embeds');
 const { logCommand } = require('../systems/logger');
-const { startDropSession, hasSession, getPrompt, handleDropMessage } = require('../systems/drops');
+const { startDropSession, hasSession, getPrompt, handleDropMessage, canDrop } = require('../systems/drops');
+const { stmts } = require('../database');
+const { Colors } = require('../config');
 
 // Import command handlers
 const setup = require('../commands/setup');
@@ -17,7 +19,7 @@ const applyUpdates = require('../commands/applyUpdates');
 
 const OWNER_ID = process.env.OWNER_ID;
 
-// Command registry
+// Owner-only command registry
 const commands = {
     'setup server': setup,
     'create roles': createRoles,
@@ -34,15 +36,14 @@ const commands = {
 module.exports = {
     name: 'messageCreate',
     async execute(client, message) {
-        // ── SECURITY: Only accept DMs from the Owner ──
         if (message.author.bot) return;
         if (message.channel.type !== ChannelType.DM) return;
-        if (message.author.id !== OWNER_ID) return;
 
+        const userId = message.author.id;
         const content = message.content.trim().toLowerCase();
 
-        // ── Handle active drop session first ──
-        if (hasSession(message.author.id)) {
+        // ── Handle active drop session (owner OR whitelisted) ──
+        if (hasSession(userId)) {
             try {
                 const handled = await handleDropMessage(message);
                 if (handled) return;
@@ -53,17 +54,80 @@ module.exports = {
             }
         }
 
-        // ── Start a new drop ──
+        // ── Whitelisted user: only allow "drop" ──
+        if (userId !== OWNER_ID) {
+            if (content === 'drop') {
+                const check = canDrop(userId);
+                if (!check.allowed) {
+                    if (check.reason === 'not_whitelisted') {
+                        return; // Silently ignore non-whitelisted users
+                    }
+                    // Rate limited
+                    return message.reply({
+                        embeds: [createEmbed({
+                            title: '⏳ Drop Cooldown',
+                            description: `You've used all **3 drops** in the last 24 hours.\n\n> Resets in **${check.resetIn}**`,
+                            color: Colors.WARNING,
+                        })],
+                    });
+                }
+                const session = startDropSession(userId);
+                await message.reply({
+                    ...getPrompt(session),
+                    content: `📦 **Drop ${check.remaining}/3 remaining today**`,
+                });
+                return;
+            }
+            return; // Whitelisted users can only use "drop"
+        }
+
+        // ── Owner-only commands below ──────────────────────────
+
+        // Drop (no rate limit for owner)
         if (content === 'drop') {
             await logCommand(client, 'drop');
-            const session = startDropSession(message.author.id);
+            const session = startDropSession(userId);
             await message.reply(getPrompt(session));
             return;
         }
 
+        // Whitelist management
+        if (content.startsWith('whitelist ')) {
+            const targetId = content.split(' ')[1]?.trim();
+            if (!targetId || !/^\d{17,20}$/.test(targetId)) {
+                return message.reply({ embeds: [errorEmbed('Invalid ID', 'Usage: `whitelist <user_id>`')] });
+            }
+            stmts.addWhitelist.run(targetId, userId);
+            const user = await client.users.fetch(targetId).catch(() => null);
+            const name = user ? `**${user.tag}**` : `ID \`${targetId}\``;
+            return message.reply({ embeds: [successEmbed('✅ Whitelisted', `${name} can now use \`drop\` (3/day)`)] });
+        }
+
+        if (content.startsWith('unwhitelist ')) {
+            const targetId = content.split(' ')[1]?.trim();
+            if (!targetId) return message.reply({ embeds: [errorEmbed('Invalid ID', 'Usage: `unwhitelist <user_id>`')] });
+            stmts.removeWhitelist.run(targetId);
+            return message.reply({ embeds: [successEmbed('✅ Removed', `User \`${targetId}\` removed from whitelist`)] });
+        }
+
+        if (content === 'whitelist') {
+            const all = stmts.getAllWhitelist.all();
+            if (all.length === 0) {
+                return message.reply({ embeds: [infoEmbed('Whitelist', 'No users whitelisted.\n\nUse `whitelist <user_id>` to add.')] });
+            }
+            const lines = [];
+            for (const w of all) {
+                const user = await client.users.fetch(w.user_id).catch(() => null);
+                const name = user ? user.tag : 'Unknown';
+                const drops = stmts.getDropCount24h.get(w.user_id);
+                lines.push(`• **${name}** (\`${w.user_id}\`) — ${drops.count}/3 drops today`);
+            }
+            return message.reply({ embeds: [infoEmbed('📋 Whitelist', lines.join('\n'))] });
+        }
+
         // Show help
         if (content === 'help') {
-            const allCommands = [...Object.keys(commands), 'drop'];
+            const allCommands = [...Object.keys(commands), 'drop', 'whitelist', 'whitelist <id>', 'unwhitelist <id>'];
             const embed = infoEmbed('WSB Commands', [
                 'All commands are **DM-only** and **Owner-only**.\n',
                 ...allCommands.map(cmd => `\`${cmd}\``),
@@ -77,10 +141,11 @@ module.exports = {
             if (content.startsWith('setup') || content.startsWith('create') ||
                 content.startsWith('backup') || content.startsWith('shutdown') ||
                 content.startsWith('ticket') || content.startsWith('permission') ||
-                content.startsWith('post') || content.startsWith('drop')) {
+                content.startsWith('post') || content.startsWith('drop') ||
+                content.startsWith('apply')) {
                 return message.reply({ embeds: [errorEmbed('Unknown Command', `Did you mean one of these?\n${[...Object.keys(commands), 'drop'].map(c => `\`${c}\``).join('\n')}`)] });
             }
-            return; // Silently ignore non-command messages
+            return;
         }
 
         try {

src/systems/drops.js

@@ -6,12 +6,40 @@ const {
 } = require('discord.js');
 const { createEmbed } = require('../utils/embeds');
 const { Colors } = require('../config');
+const { stmts } = require('../database');
+
+const MAX_DROPS_PER_DAY = 3;
+const OWNER_ID = process.env.OWNER_ID;
 
 // Active drop sessions: Map<userId, session>
 const activeSessions = new Map();
 
 const STEPS = ['title', 'file', 'warnings', 'status', 'about', 'image', 'preview'];
 
+/**
+ * Check if a user can drop (owner = unlimited, whitelisted = 3/day).
+ * Returns { allowed: boolean, remaining?: number, resetIn?: string }
+ */
+function canDrop(userId) {
+    if (userId === OWNER_ID) return { allowed: true, remaining: Infinity };
+
+    const wl = stmts.getWhitelist.get(userId);
+    if (!wl) return { allowed: false, reason: 'not_whitelisted' };
+
+    const { count } = stmts.getDropCount24h.get(userId);
+    if (count >= MAX_DROPS_PER_DAY) {
+        const last = stmts.getLastDrop.get(userId);
+        const resetTime = new Date(last.dropped_at + 'Z');
+        resetTime.setHours(resetTime.getHours() + 24);
+        const diff = resetTime - new Date();
+        const hours = Math.floor(diff / 3600000);
+        const mins = Math.floor((diff % 3600000) / 60000);
+        return { allowed: false, reason: 'rate_limited', resetIn: `${hours}h ${mins}m` };
+    }
+
+    return { allowed: true, remaining: MAX_DROPS_PER_DAY - count };
+}
+
 /**
  * Start a new drop session for the owner.
  */
@@ -268,6 +296,9 @@ async function handleDropMessage(message) {
                     components: [downloadRow],
                 });
 
+                // Log the drop for rate limiting
+                stmts.logDrop.run(userId, session.title, channelId);
+
                 activeSessions.delete(userId);
                 await message.reply({
                     embeds: [createEmbed({
@@ -355,4 +386,5 @@ module.exports = {
     handleDropMessage,
     handleDropButton,
     buildDropEmbed,
+    canDrop,
 };