feat: store punished users

bot/events_handlers/member_events.py

@@ -1,4 +1,9 @@
 from bot.utils import generate_random_nickname, is_valid_username, is_numeric_name
+from database import user_has_role
+
+import discord
+
+SPAM_ROLE_ID = 1350511935677927514
 
 
 async def handle_member_join(member):
@@ -15,6 +20,11 @@ async def handle_member_join(member):
         new_nick = generate_random_nickname()
         await member.edit(nick=new_nick)
 
+    if user_has_role(member.id):
+        spam_role = discord.utils.get(member.guild.roles, id=SPAM_ROLE_ID)
+        if spam_role:
+            await member.add_roles(spam_role)
+
 
 async def handle_member_update(before, after):
     name_to_check = after.name

bot/events_handlers/message_events.py

@@ -1,9 +1,12 @@
 from datetime import timedelta
 import re
 import discord
+
 from bot.utils import timeout_member, aware_utcnow
 from bot.config import message_patterns
 
+from database import add_user_to_role
+
 BOT_LOG = 1112049391482703873
 
 CRAZY_USER_ID = 1319364607487512658
@@ -210,7 +213,12 @@ async def handle_message(message, bot):
             spam_role = message.guild.get_role(SPAM_ROLE_ID)
             member = message.guild.get_member(message.author.id)
 
-            await member.add_roles(spam_role)
+            # Check if the member already has the spam role
+            if spam_role not in member.roles:
+                await member.add_roles(spam_role)
+
+                # Add the user to the database
+                add_user_to_role(member.id, SPAM_ROLE_ID)
 
             await message.reply(
                 f"Dink Donk! Time to ping everyone! {spam_role.mention}",

bot/tasks.py

@@ -6,9 +6,13 @@ from discord.ext import tasks, commands
 
 from bot.utils import aware_utcnow, fetch_api_data
 
+from database import migrate_users_with_role
+
 TARGET_DATE = datetime(2036, 8, 12, tzinfo=timezone.utc)
 OFFTOPIC_CHANNEL = 1112048063448617142
 
+SPAM_ROLE_ID = 1350511935677927514
+
 COD_GAMES = {
     10180: {"name": "Modern Warfare 2 (2009)", "channel": 1145458108190163014},
     42680: {"name": "Modern Warfare 3 (2011)", "channel": 1145459504436220014},
@@ -18,6 +22,16 @@ COD_GAMES = {
 }
 
 
+async def migrate_all_users(bot):
+    # Fetch users with the SPAM_ROLE_ID and migrate them to the database
+    guild = discord.utils.get(bot.guilds)
+    if guild:
+        spam_role = discord.utils.get(guild.roles, id=SPAM_ROLE_ID)
+        if spam_role:
+            for member in spam_role.members:
+                migrate_users_with_role(member.id, SPAM_ROLE_ID)
+
+
 class SteamSaleChecker(commands.Cog):
     def __init__(self, bot):
         self.bot = bot
@@ -132,6 +146,8 @@ async def setup(bot):
         except Exception as e:
             print(f"An error occurred in heat_death task: {e}")
 
+    await migrate_all_users(bot)
+
     update_status.start()
     heat_death.start()
 

database/__init__.py

@@ -1,6 +1,8 @@
 import sqlite3
 import os
 
+from bot.utils import aware_utcnow
+
 DB_DIR = os.getenv("BOT_DATA_DIR", "/bot-data")
 DB_PATH = os.path.join(DB_DIR, "database.db")
 
@@ -55,3 +57,44 @@ def remove_pattern(pattern_id: int):
     cursor.execute("DELETE FROM message_patterns WHERE id = ?", (pattern_id,))
     conn.commit()
     conn.close()
+
+
+def migrate_users_with_role(user_id: int, role_id: int):
+    """Migrates existing users with the role to the new table."""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        "INSERT OR IGNORE INTO user_roles (user_id, role_id, date_assigned) VALUES (?, ?, ?)",
+        (user_id, role_id, aware_utcnow().isoformat()),
+    )
+
+    conn.commit()
+    conn.close()
+
+
+def add_user_to_role(user_id: int, role_id: int):
+    """Adds a new user when they receive the role."""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    cursor.execute(
+        "INSERT OR REPLACE INTO user_roles (user_id, role_id, date_assigned) VALUES (?, ?, ?)",
+        (user_id, role_id, aware_utcnow().isoformat()),
+    )
+
+    conn.commit()
+    conn.close()
+
+
+def user_has_role(user_id: int):
+    """Checks if a user is in the database."""
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    cursor.execute("SELECT * FROM user_roles WHERE user_id = ?", (user_id,))
+    result = cursor.fetchone()
+
+    conn.close()
+
+    return result is not None

database/schema.sql

@@ -2,4 +2,10 @@ CREATE TABLE IF NOT EXISTS message_patterns (
     id INTEGER PRIMARY KEY AUTOINCREMENT,
     regex TEXT NOT NULL,
     response TEXT NOT NULL
-)
+);
+
+CREATE TABLE IF NOT EXISTS user_roles (
+    user_id INTEGER PRIMARY KEY,
+    role_id INTEGER NOT NULL,
+    date_assigned TEXT
+);