From 2ab239a9a32ae285148b176e8ca80520165a5090 Mon Sep 17 00:00:00 2001 From: diamante0018 Date: Sun, 16 Mar 2025 17:07:11 +0100 Subject: [PATCH] feat: store punished users --- bot/events_handlers/member_events.py | 10 +++++++ bot/events_handlers/message_events.py | 10 ++++++- bot/tasks.py | 16 ++++++++++ database/__init__.py | 43 +++++++++++++++++++++++++++ database/schema.sql | 8 ++++- 5 files changed, 85 insertions(+), 2 deletions(-) diff --git a/bot/events_handlers/member_events.py b/bot/events_handlers/member_events.py index 491757d..0dd8bec 100644 --- a/bot/events_handlers/member_events.py +++ b/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 diff --git a/bot/events_handlers/message_events.py b/bot/events_handlers/message_events.py index d1a5bca..1f92144 100644 --- a/bot/events_handlers/message_events.py +++ b/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}", diff --git a/bot/tasks.py b/bot/tasks.py index dd12c55..59b204d 100644 --- a/bot/tasks.py +++ b/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() diff --git a/database/__init__.py b/database/__init__.py index 5843278..760439f 100644 --- a/database/__init__.py +++ b/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 diff --git a/database/schema.sql b/database/schema.sql index ea1c1d1..c69f54e 100644 --- a/database/schema.sql +++ b/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 +);