diff --git a/commands/roletrack.py b/commands/roletrack.py new file mode 100644 index 0000000..0a8ca11 --- /dev/null +++ b/commands/roletrack.py @@ -0,0 +1,125 @@ +import discord +from discord import app_commands +from discord.ext import commands +import aiosqlite +import os + +class RoleTrack(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.db_path = os.path.join(os.path.dirname(__file__), 'roletrack.db') + bot.loop.create_task(self.init_db()) + + async def init_db(self): + """Initialize the database tables if they don't exist""" + async with aiosqlite.connect(self.db_path) as db: + # Table for tracking opt-in status + await db.execute(''' + CREATE TABLE IF NOT EXISTS role_tracking_users ( + user_id INTEGER, + guild_id INTEGER, + opted_in INTEGER DEFAULT 1, + PRIMARY KEY (user_id, guild_id) + ) + ''') + + # Table for storing role data + await db.execute(''' + CREATE TABLE IF NOT EXISTS tracked_roles ( + user_id INTEGER, + guild_id INTEGER, + role_id INTEGER, + role_name TEXT, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (user_id, guild_id, role_id) + ) + ''') + + await db.commit() + + async def is_opted_in(self, user_id: int, guild_id: int) -> bool: + """Check if a user is opted into role tracking""" + async with aiosqlite.connect(self.db_path) as db: + async with db.execute(''' + SELECT opted_in FROM role_tracking_users + WHERE user_id = ? AND guild_id = ? + ''', (user_id, guild_id)) as cursor: + result = await cursor.fetchone() + return result[0] == 1 if result else False + + async def save_user_roles(self, member: discord.Member): + """Save all roles for a user (excluding @everyone)""" + if not await self.is_opted_in(member.id, member.guild.id): + return + + async with aiosqlite.connect(self.db_path) as db: + # Clear existing roles for this user in this guild + await db.execute(''' + DELETE FROM tracked_roles + WHERE user_id = ? AND guild_id = ? + ''', (member.id, member.guild.id)) + + # Save current roles (excluding @everyone) + for role in member.roles: + if role.id != member.guild.id: # Skip @everyone role + await db.execute(''' + INSERT INTO tracked_roles (user_id, guild_id, role_id, role_name) + VALUES (?, ?, ?, ?) + ''', (member.id, member.guild.id, role.id, role.name)) + + await db.commit() + + @app_commands.command(name="roletrack", description="Manage role tracking opt-in status") + @app_commands.describe( + setting="Choose to opt in or opt out of role tracking" + ) + @app_commands.choices(setting=[ + app_commands.Choice(name="opt_in", value="opt_in"), + app_commands.Choice(name="opt_out", value="opt_out") + ]) + async def roletrack(self, interaction: discord.Interaction, setting: app_commands.Choice[str]): + """Command to opt in or opt out of role tracking""" + async with aiosqlite.connect(self.db_path) as db: + if setting.value == "opt_in": + # Opt the user in + await db.execute(''' + INSERT OR REPLACE INTO role_tracking_users (user_id, guild_id, opted_in) + VALUES (?, ?, 1) + ''', (interaction.user.id, interaction.guild.id)) + + await db.commit() + + # Save current roles (outside the db context) + await self.save_user_roles(interaction.user) + + embed = discord.Embed( + title="✅ Role Tracking Enabled", + description="You've been opted into role tracking. If you leave and rejoin this server, your roles will be automatically restored.", + color=discord.Color.green() + ) + await interaction.response.send_message(embed=embed, ephemeral=True) + + elif setting.value == "opt_out": + # Opt the user out + await db.execute(''' + INSERT OR REPLACE INTO role_tracking_users (user_id, guild_id, opted_in) + VALUES (?, ?, 0) + ''', (interaction.user.id, interaction.guild.id)) + + # Delete their tracked roles + await db.execute(''' + DELETE FROM tracked_roles + WHERE user_id = ? AND guild_id = ? + ''', (interaction.user.id, interaction.guild.id)) + + await db.commit() + + embed = discord.Embed( + title="❌ Role Tracking Disabled", + description="You've been opted out of role tracking. Your saved roles have been deleted.", + color=discord.Color.red() + ) + await interaction.response.send_message(embed=embed, ephemeral=True) + +async def setup(bot): + await bot.add_cog(RoleTrack(bot)) \ No newline at end of file diff --git a/events/roletrack.py b/events/roletrack.py new file mode 100644 index 0000000..c848fc2 --- /dev/null +++ b/events/roletrack.py @@ -0,0 +1,140 @@ +import discord +from discord.ext import commands +import sqlite3 +import os +import asyncio + +class RoleTrackEvents(commands.Cog): + def __init__(self, bot): + self.bot = bot + self.db_path = os.path.join(os.path.dirname(__file__), '..', 'roletrack.db') + + def is_opted_in(self, user_id: int, guild_id: int) -> bool: + """Check if a user is opted into role tracking""" + conn = sqlite3.connect(self.db_path) + c = conn.cursor() + + c.execute(''' + SELECT opted_in FROM role_tracking_users + WHERE user_id = ? AND guild_id = ? + ''', (user_id, guild_id)) + + result = c.fetchone() + conn.close() + + return result[0] == 1 if result else False + + def get_tracked_roles(self, user_id: int, guild_id: int): + """Retrieve tracked roles for a user""" + conn = sqlite3.connect(self.db_path) + c = conn.cursor() + + c.execute(''' + SELECT role_id, role_name FROM tracked_roles + WHERE user_id = ? AND guild_id = ? + ''', (user_id, guild_id)) + + roles = c.fetchall() + conn.close() + return roles + + def save_user_roles(self, member: discord.Member): + """Save all roles for a user (excluding @everyone)""" + if not self.is_opted_in(member.id, member.guild.id): + return + + conn = sqlite3.connect(self.db_path) + c = conn.cursor() + + # Clear existing roles for this user in this guild + c.execute(''' + DELETE FROM tracked_roles + WHERE user_id = ? AND guild_id = ? + ''', (member.id, member.guild.id)) + + # Save current roles (excluding @everyone) + for role in member.roles: + if role.id != member.guild.id: # Skip @everyone role + c.execute(''' + INSERT INTO tracked_roles (user_id, guild_id, role_id, role_name) + VALUES (?, ?, ?, ?) + ''', (member.id, member.guild.id, role.id, role.name)) + + conn.commit() + conn.close() + + @commands.Cog.listener() + async def on_member_join(self, member: discord.Member): + """Restore roles when a tracked user rejoins""" + if member.bot: + return + + # Check if user is opted in + if not self.is_opted_in(member.id, member.guild.id): + return + + # Get their tracked roles + tracked_roles = self.get_tracked_roles(member.id, member.guild.id) + + if not tracked_roles: + return + + # Wait a moment for Discord to fully process the join + await asyncio.sleep(1) + + # Restore roles + roles_to_add = [] + missing_roles = [] + + for role_id, role_name in tracked_roles: + role = member.guild.get_role(role_id) + if role: + # Check if bot can assign this role (bot's highest role must be higher) + if role < member.guild.me.top_role and not role.is_default(): + roles_to_add.append(role) + else: + missing_roles.append(role_name) + + # Add roles + if roles_to_add: + try: + await member.add_roles(*roles_to_add, reason="Role tracking: User rejoined server") + print(f"Restored {len(roles_to_add)} roles for {member.name} ({member.id})") + except discord.Forbidden: + print(f"Failed to restore roles for {member.name}: Missing permissions") + except discord.HTTPException as e: + print(f"Failed to restore roles for {member.name}: {e}") + + # Log if there were missing roles + if missing_roles: + print(f"Could not restore the following roles for {member.name} (roles no longer exist): {', '.join(missing_roles)}") + + @commands.Cog.listener() + async def on_member_update(self, before: discord.Member, after: discord.Member): + """Update tracked roles when a user's roles change""" + if after.bot: + return + + # Check if user is opted in + if not self.is_opted_in(after.id, after.guild.id): + return + + # Check if roles actually changed + if before.roles == after.roles: + return + + # Save updated roles + self.save_user_roles(after) + print(f"Updated tracked roles for {after.name} ({after.id})") + + @commands.Cog.listener() + async def on_member_remove(self, member: discord.Member): + """Optional: Log when a tracked user leaves""" + if member.bot: + return + + if self.is_opted_in(member.id, member.guild.id): + print(f"Tracked user {member.name} ({member.id}) left the server. Roles are saved for restoration.") + +async def setup(bot): + await bot.add_cog(RoleTrackEvents(bot)) \ No newline at end of file