From 5c3f0b98105471791179cc9c60c7b5477dcb6d94 Mon Sep 17 00:00:00 2001 From: Slipstream Date: Sat, 3 May 2025 13:47:49 -0600 Subject: [PATCH] hh --- .env.example | 13 ++ cogs/settings_cog.py | 208 ++++++++++++++++++ cogs/welcome_cog.py | 212 ++++++++++++++++++ main.py | 128 ++++++++++- requirements.txt | 3 +- settings_manager.py | 496 +++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 1051 insertions(+), 9 deletions(-) create mode 100644 cogs/settings_cog.py create mode 100644 cogs/welcome_cog.py create mode 100644 settings_manager.py diff --git a/.env.example b/.env.example index 5828da5..51d99a9 100644 --- a/.env.example +++ b/.env.example @@ -42,6 +42,19 @@ UNIFIED_API_PORT=5005 PISTON_API_URL=https://emkc.org/api/v2/piston/execute # Example public Piston instance URL # PISTON_API_KEY=YOUR_PISTON_API_KEY_IF_NEEDED # Optional, depending on the Piston instance used +# PostgreSQL Configuration (If using PostgreSQL for features like economy, settings, etc.) +# POSTGRES_HOST=localhost +# POSTGRES_PORT=5432 +# POSTGRES_USER=your_postgres_user +# POSTGRES_PASSWORD=your_postgres_password +# POSTGRES_DB=your_primary_database_name # e.g., bot-economy +# POSTGRES_SETTINGS_DB=your_settings_database_name # e.g., discord_bot_settings + +# Redis Configuration (If using Redis for caching, etc.) +# REDIS_HOST=localhost +# REDIS_PORT=6379 +# REDIS_PASSWORD=your_redis_password # Optional + # Terminal Command Execution Configuration (For GurtCog run_terminal_command tool) GURT_SAFETY_CHECK_MODEL=openai/gpt-4.1-nano # Model for AI safety check (e.g., openai/gpt-4.1-nano) DOCKER_EXEC_IMAGE=alpine:latest # Docker image for command execution (e.g., alpine:latest) diff --git a/cogs/settings_cog.py b/cogs/settings_cog.py new file mode 100644 index 0000000..e3f6ed0 --- /dev/null +++ b/cogs/settings_cog.py @@ -0,0 +1,208 @@ +import discord +from discord.ext import commands +import logging +from discordbot import settings_manager # Assuming settings_manager is accessible + +log = logging.getLogger(__name__) + +# CORE_COGS definition moved to main.py + +class SettingsCog(commands.Cog, name="Settings"): + """Commands for server administrators to configure the bot.""" + + def __init__(self, bot: commands.Bot): + self.bot = bot + + # --- Prefix Management --- + @commands.command(name='setprefix', help="Sets the command prefix for this server. Usage: `setprefix `") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def set_prefix(self, ctx: commands.Context, new_prefix: str): + """Sets the command prefix for the current guild.""" + if not new_prefix: + await ctx.send("Prefix cannot be empty.") + return + if len(new_prefix) > 10: # Arbitrary limit + await ctx.send("Prefix cannot be longer than 10 characters.") + return + if new_prefix.isspace(): + await ctx.send("Prefix cannot be just whitespace.") + return + + guild_id = ctx.guild.id + success = await settings_manager.set_guild_prefix(guild_id, new_prefix) + + if success: + await ctx.send(f"Command prefix for this server has been set to: `{new_prefix}`") + log.info(f"Prefix updated for guild {guild_id} to '{new_prefix}' by {ctx.author.name}") + else: + await ctx.send("Failed to set the prefix. Please check the logs.") + log.error(f"Failed to save prefix for guild {guild_id}") + + @commands.command(name='showprefix', help="Shows the current command prefix for this server.") + @commands.guild_only() + async def show_prefix(self, ctx: commands.Context): + """Shows the current command prefix.""" + # We need the bot's default prefix as a fallback + # This might need access to the bot instance's initial config or a constant + default_prefix = self.bot.command_prefix # This might not work if command_prefix is the callable + # Use the constant defined in main.py if possible, or keep a local fallback + default_prefix_fallback = "!" # TODO: Get default prefix reliably if needed elsewhere + + guild_id = ctx.guild.id + current_prefix = await settings_manager.get_guild_prefix(guild_id, default_prefix_fallback) + await ctx.send(f"The current command prefix for this server is: `{current_prefix}`") + + + # --- Cog Management --- + @commands.command(name='enablecog', help="Enables a specific module (cog) for this server. Usage: `enablecog `") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def enable_cog(self, ctx: commands.Context, cog_name: str): + """Enables a cog for the current guild.""" + # Validate if cog exists + if cog_name not in self.bot.cogs: + await ctx.send(f"Error: Cog `{cog_name}` not found.") + return + + if cog_name in CORE_COGS: + await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled/enabled.") + return + + guild_id = ctx.guild.id + success = await settings_manager.set_cog_enabled(guild_id, cog_name, enabled=True) + + if success: + await ctx.send(f"Module `{cog_name}` has been enabled for this server.") + log.info(f"Cog '{cog_name}' enabled for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send(f"Failed to enable module `{cog_name}`. Check logs.") + log.error(f"Failed to enable cog '{cog_name}' for guild {guild_id}") + + @commands.command(name='disablecog', help="Disables a specific module (cog) for this server. Usage: `disablecog `") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def disable_cog(self, ctx: commands.Context, cog_name: str): + """Disables a cog for the current guild.""" + if cog_name not in self.bot.cogs: + await ctx.send(f"Error: Cog `{cog_name}` not found.") + return + + if cog_name in CORE_COGS: + await ctx.send(f"Error: Core cog `{cog_name}` cannot be disabled.") + return + + guild_id = ctx.guild.id + success = await settings_manager.set_cog_enabled(guild_id, cog_name, enabled=False) + + if success: + await ctx.send(f"Module `{cog_name}` has been disabled for this server.") + log.info(f"Cog '{cog_name}' disabled for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send(f"Failed to disable module `{cog_name}`. Check logs.") + log.error(f"Failed to disable cog '{cog_name}' for guild {guild_id}") + + @commands.command(name='listcogs', help="Lists all available modules (cogs) and their status for this server.") + @commands.guild_only() + async def list_cogs(self, ctx: commands.Context): + """Lists available cogs and their enabled/disabled status.""" + guild_id = ctx.guild.id + # Note: Default enabled status might need adjustment based on desired behavior + # If a cog has no entry in the DB, should it be considered enabled or disabled by default? + # Let's assume default_enabled=True for now. + default_behavior = True + + embed = discord.Embed(title="Available Modules (Cogs)", color=discord.Color.blue()) + lines = [] + # Use the CORE_COGS defined at the top of this file + core_cogs_list = CORE_COGS + + for cog_name in sorted(self.bot.cogs.keys()): + is_enabled = await settings_manager.is_cog_enabled(guild_id, cog_name, default_enabled=default_behavior) + status = "✅ Enabled" if is_enabled else "❌ Disabled" + if cog_name in core_cogs_list: + status += " (Core)" + lines.append(f"`{cog_name}`: {status}") + + embed.description = "\n".join(lines) if lines else "No cogs found." + await ctx.send(embed=embed) + + + # --- Command Permission Management (Basic Role-Based) --- + @commands.command(name='allowcmd', help="Allows a role to use a specific command. Usage: `allowcmd <@Role>`") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def allow_command(self, ctx: commands.Context, command_name: str, role: discord.Role): + """Allows a specific role to use a command.""" + command = self.bot.get_command(command_name) + if not command: + await ctx.send(f"Error: Command `{command_name}` not found.") + return + + guild_id = ctx.guild.id + role_id = role.id + success = await settings_manager.add_command_permission(guild_id, command_name, role_id) + + if success: + await ctx.send(f"Role `{role.name}` is now allowed to use command `{command_name}`.") + log.info(f"Permission added for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}") + else: + await ctx.send(f"Failed to add permission for command `{command_name}`. Check logs.") + log.error(f"Failed to add permission for command '{command_name}', role {role_id} in guild {guild_id}") + + + @commands.command(name='disallowcmd', help="Disallows a role from using a specific command. Usage: `disallowcmd <@Role>`") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def disallow_command(self, ctx: commands.Context, command_name: str, role: discord.Role): + """Disallows a specific role from using a command.""" + command = self.bot.get_command(command_name) + if not command: + await ctx.send(f"Error: Command `{command_name}` not found.") + return + + guild_id = ctx.guild.id + role_id = role.id + success = await settings_manager.remove_command_permission(guild_id, command_name, role_id) + + if success: + await ctx.send(f"Role `{role.name}` is no longer allowed to use command `{command_name}`.") + log.info(f"Permission removed for command '{command_name}', role '{role.name}' ({role_id}) in guild {guild_id} by {ctx.author.name}") + else: + await ctx.send(f"Failed to remove permission for command `{command_name}`. Check logs.") + log.error(f"Failed to remove permission for command '{command_name}', role {role_id} in guild {guild_id}") + + # TODO: Add command to list permissions? + + # --- Error Handling for this Cog --- + @set_prefix.error + @enable_cog.error + @disable_cog.error + @allow_command.error + @disallow_command.error + async def on_command_error(self, ctx: commands.Context, error): + if isinstance(error, commands.MissingPermissions): + await ctx.send("You need Administrator permissions to use this command.") + elif isinstance(error, commands.BadArgument): + await ctx.send(f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`") + elif isinstance(error, commands.MissingRequiredArgument): + await ctx.send(f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`") + elif isinstance(error, commands.NoPrivateMessage): + await ctx.send("This command cannot be used in private messages.") + else: + log.error(f"Unhandled error in SettingsCog command '{ctx.command.name}': {error}") + await ctx.send("An unexpected error occurred. Please check the logs.") + + +async def setup(bot: commands.Bot): + # Ensure pools are initialized before adding the cog + if settings_manager.pg_pool is None or settings_manager.redis_pool is None: + log.warning("Settings Manager pools not initialized before loading SettingsCog. Attempting initialization.") + try: + await settings_manager.initialize_pools() + except Exception as e: + log.exception("Failed to initialize Settings Manager pools during SettingsCog setup. Cog will not load.") + return # Prevent loading if pools fail + + await bot.add_cog(SettingsCog(bot)) + log.info("SettingsCog loaded.") diff --git a/cogs/welcome_cog.py b/cogs/welcome_cog.py new file mode 100644 index 0000000..151c2ff --- /dev/null +++ b/cogs/welcome_cog.py @@ -0,0 +1,212 @@ +import discord +from discord.ext import commands +import logging +from discordbot import settings_manager # Assuming settings_manager is accessible + +log = logging.getLogger(__name__) + +class WelcomeCog(commands.Cog): + """Handles welcome and goodbye messages for guilds.""" + + def __init__(self, bot: commands.Bot): + self.bot = bot + + @commands.Cog.listener() + async def on_member_join(self, member: discord.Member): + """Sends a welcome message when a new member joins.""" + guild = member.guild + if not guild: + return + + log.debug(f"Member {member.name} joined guild {guild.name} ({guild.id})") + + # --- Fetch settings --- + welcome_channel_id_str = await settings_manager.get_setting(guild.id, 'welcome_channel_id') + welcome_message_template = await settings_manager.get_setting(guild.id, 'welcome_message', default="Welcome {user} to {server}!") + + # Handle the "__NONE__" marker for potentially unset values + if not welcome_channel_id_str or welcome_channel_id_str == "__NONE__": + log.debug(f"Welcome channel not configured for guild {guild.id}") + return + + try: + welcome_channel_id = int(welcome_channel_id_str) + channel = guild.get_channel(welcome_channel_id) + if not channel or not isinstance(channel, discord.TextChannel): + log.warning(f"Welcome channel ID {welcome_channel_id} not found or not text channel in guild {guild.id}") + # Maybe remove the setting here if the channel is invalid? + return + + # --- Format and send message --- + # Basic formatting, can be expanded + formatted_message = welcome_message_template.format( + user=member.mention, + username=member.name, + server=guild.name + ) + + await channel.send(formatted_message) + log.info(f"Sent welcome message for {member.name} in guild {guild.id}") + + except ValueError: + log.error(f"Invalid welcome_channel_id '{welcome_channel_id_str}' configured for guild {guild.id}") + except discord.Forbidden: + log.error(f"Missing permissions to send welcome message in channel {welcome_channel_id} for guild {guild.id}") + except Exception as e: + log.exception(f"Error sending welcome message for guild {guild.id}: {e}") + + @commands.Cog.listener() + async def on_member_remove(self, member: discord.Member): + """Sends a goodbye message when a member leaves.""" + guild = member.guild + if not guild: + return + + log.debug(f"Member {member.name} left guild {guild.name} ({guild.id})") + + # --- Fetch settings --- + goodbye_channel_id_str = await settings_manager.get_setting(guild.id, 'goodbye_channel_id') + goodbye_message_template = await settings_manager.get_setting(guild.id, 'goodbye_message', default="{username} has left the server.") + + # Handle the "__NONE__" marker + if not goodbye_channel_id_str or goodbye_channel_id_str == "__NONE__": + log.debug(f"Goodbye channel not configured for guild {guild.id}") + return + + try: + goodbye_channel_id = int(goodbye_channel_id_str) + channel = guild.get_channel(goodbye_channel_id) + if not channel or not isinstance(channel, discord.TextChannel): + log.warning(f"Goodbye channel ID {goodbye_channel_id} not found or not text channel in guild {guild.id}") + return + + # --- Format and send message --- + formatted_message = goodbye_message_template.format( + user=member.mention, # Might not be mentionable after leaving + username=member.name, + server=guild.name + ) + + await channel.send(formatted_message) + log.info(f"Sent goodbye message for {member.name} in guild {guild.id}") + + except ValueError: + log.error(f"Invalid goodbye_channel_id '{goodbye_channel_id_str}' configured for guild {guild.id}") + except discord.Forbidden: + log.error(f"Missing permissions to send goodbye message in channel {goodbye_channel_id} for guild {guild.id}") + except Exception as e: + log.exception(f"Error sending goodbye message for guild {guild.id}: {e}") + + + @commands.command(name='setwelcome', help="Sets the welcome message and channel. Usage: `setwelcome #channel [message template]`") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def set_welcome(self, ctx: commands.Context, channel: discord.TextChannel, *, message_template: str = "Welcome {user} to {server}!"): + """Sets the channel and template for welcome messages.""" + guild_id = ctx.guild.id + key_channel = 'welcome_channel_id' + key_message = 'welcome_message' + + # Use settings_manager.set_setting + success_channel = await settings_manager.set_setting(guild_id, key_channel, str(channel.id)) + success_message = await settings_manager.set_setting(guild_id, key_message, message_template) + + if success_channel and success_message: # Both need to succeed + await ctx.send(f"Welcome messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```") + log.info(f"Welcome settings updated for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send("Failed to save welcome settings. Check logs.") + log.error(f"Failed to save welcome settings for guild {guild_id}") + + @commands.command(name='disablewelcome', help="Disables welcome messages for this server.") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def disable_welcome(self, ctx: commands.Context): + """Disables welcome messages by removing the channel setting.""" + guild_id = ctx.guild.id + key_channel = 'welcome_channel_id' + key_message = 'welcome_message' # Also clear the message template + + # Use set_setting with None to delete the settings + success_channel = await settings_manager.set_setting(guild_id, key_channel, None) + success_message = await settings_manager.set_setting(guild_id, key_message, None) + + if success_channel and success_message: # Both need to succeed + await ctx.send("Welcome messages have been disabled.") + log.info(f"Welcome messages disabled for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send("Failed to disable welcome messages. Check logs.") + log.error(f"Failed to disable welcome settings for guild {guild_id}") + + + @commands.command(name='setgoodbye', help="Sets the goodbye message and channel. Usage: `setgoodbye #channel [message template]`") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def set_goodbye(self, ctx: commands.Context, channel: discord.TextChannel, *, message_template: str = "{username} has left the server."): + """Sets the channel and template for goodbye messages.""" + guild_id = ctx.guild.id + key_channel = 'goodbye_channel_id' + key_message = 'goodbye_message' + + # Use settings_manager.set_setting + success_channel = await settings_manager.set_setting(guild_id, key_channel, str(channel.id)) + success_message = await settings_manager.set_setting(guild_id, key_message, message_template) + + if success_channel and success_message: # Both need to succeed + await ctx.send(f"Goodbye messages will now be sent to {channel.mention} with the template:\n```\n{message_template}\n```") + log.info(f"Goodbye settings updated for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send("Failed to save goodbye settings. Check logs.") + log.error(f"Failed to save goodbye settings for guild {guild_id}") + + @commands.command(name='disablegoodbye', help="Disables goodbye messages for this server.") + @commands.has_permissions(administrator=True) + @commands.guild_only() + async def disable_goodbye(self, ctx: commands.Context): + """Disables goodbye messages by removing the channel setting.""" + guild_id = ctx.guild.id + key_channel = 'goodbye_channel_id' + key_message = 'goodbye_message' + + # Use set_setting with None to delete the settings + success_channel = await settings_manager.set_setting(guild_id, key_channel, None) + success_message = await settings_manager.set_setting(guild_id, key_message, None) + + if success_channel and success_message: # Both need to succeed + await ctx.send("Goodbye messages have been disabled.") + log.info(f"Goodbye messages disabled for guild {guild_id} by {ctx.author.name}") + else: + await ctx.send("Failed to disable goodbye messages. Check logs.") + log.error(f"Failed to disable goodbye settings for guild {guild_id}") + + # Error Handling for this Cog + @set_welcome.error + @disable_welcome.error + @set_goodbye.error + @disable_goodbye.error + async def on_command_error(self, ctx: commands.Context, error): + if isinstance(error, commands.MissingPermissions): + await ctx.send("You need Administrator permissions to use this command.") + elif isinstance(error, commands.BadArgument): + await ctx.send(f"Invalid argument provided. Check the command help: `{ctx.prefix}help {ctx.command.name}`") + elif isinstance(error, commands.MissingRequiredArgument): + await ctx.send(f"Missing required argument. Check the command help: `{ctx.prefix}help {ctx.command.name}`") + elif isinstance(error, commands.NoPrivateMessage): + await ctx.send("This command cannot be used in private messages.") + else: + log.error(f"Unhandled error in WelcomeCog command '{ctx.command.name}': {error}") + await ctx.send("An unexpected error occurred. Please check the logs.") + + +async def setup(bot: commands.Bot): + # Ensure pools are initialized before adding the cog + if settings_manager.pg_pool is None or settings_manager.redis_pool is None: + log.warning("Settings Manager pools not initialized before loading WelcomeCog. Attempting initialization.") + try: + await settings_manager.initialize_pools() + except Exception as e: + log.exception("Failed to initialize Settings Manager pools during WelcomeCog setup. Cog will not load.") + return # Prevent loading if pools fail + + await bot.add_cog(WelcomeCog(bot)) + log.info("WelcomeCog loaded.") diff --git a/main.py b/main.py index 83f6345..f3722b8 100644 --- a/main.py +++ b/main.py @@ -7,10 +7,12 @@ import sys import asyncio import subprocess import importlib.util -import argparse # Import argparse +import argparse +import logging # Add logging from commands import load_all_cogs, reload_all_cogs from error_handler import handle_error, patch_discord_methods, store_interaction_content from utils import reload_script +from discordbot import settings_manager # Import the settings manager # Import the unified API service runner and the sync API module import sys @@ -29,15 +31,38 @@ except ImportError: # Load environment variables from .env file load_dotenv() +# --- Constants --- +DEFAULT_PREFIX = "!" +CORE_COGS = {'SettingsCog', 'HelpCog'} # Cogs that cannot be disabled + +# --- Dynamic Prefix Function --- +async def get_prefix(bot_instance, message): + """Determines the command prefix based on guild settings or default.""" + if not message.guild: + # Use default prefix in DMs + return commands.when_mentioned_or(DEFAULT_PREFIX)(bot_instance, message) + + # Fetch prefix from settings manager (cache first, then DB) + prefix = await settings_manager.get_guild_prefix(message.guild.id, DEFAULT_PREFIX) + return commands.when_mentioned_or(prefix)(bot_instance, message) + +# --- Bot Setup --- # Set up intents (permissions) intents = discord.Intents.default() intents.message_content = True intents.members = True -# Create bot instance with command prefix '!' and enable the application commands -bot = commands.Bot(command_prefix='!', intents=intents) +# Create bot instance with the dynamic prefix function +bot = commands.Bot(command_prefix=get_prefix, intents=intents) bot.owner_id = int(os.getenv('OWNER_USER_ID')) +bot.core_cogs = CORE_COGS # Attach core cogs list to bot instance +# --- Logging Setup --- +# Configure logging (adjust level and format as needed) +logging.basicConfig(level=logging.INFO, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s') +log = logging.getLogger(__name__) # Logger for main.py + +# --- Events --- @bot.event async def on_ready(): print(f'{bot.user.name} has connected to Discord!') @@ -87,15 +112,96 @@ async def on_shard_disconnect(shard_id): except Exception as e: print(f"Failed to reconnect shard {shard_id}: {e}") -# Error handling +# Error handling - Updated to handle custom check failures @bot.event async def on_command_error(ctx, error): - await handle_error(ctx, error) + if isinstance(error, CogDisabledError): + await ctx.send(str(error), ephemeral=True) # Send the error message from the exception + log.warning(f"Command '{ctx.command.qualified_name}' blocked for user {ctx.author.id} in guild {ctx.guild.id}: {error}") + elif isinstance(error, CommandPermissionError): + await ctx.send(str(error), ephemeral=True) # Send the error message from the exception + log.warning(f"Command '{ctx.command.qualified_name}' blocked for user {ctx.author.id} in guild {ctx.guild.id}: {error}") + else: + # Pass other errors to the original handler + await handle_error(ctx, error) @bot.tree.error async def on_app_command_error(interaction, error): await handle_error(interaction, error) +# --- Global Command Checks --- + +# Need to import SettingsCog to access CORE_COGS, or define CORE_COGS here. +# Let's import it, assuming it's safe to do so at the top level. +# If it causes circular imports, CORE_COGS needs to be defined elsewhere or passed differently. +try: + from discordbot.cogs import settings_cog # Import the cog itself +except ImportError: + log.error("Could not import settings_cog.py for CORE_COGS definition. Cog checks might fail.") + settings_cog = None # Define as None to avoid NameError later + +class CogDisabledError(commands.CheckFailure): + """Custom exception for disabled cogs.""" + def __init__(self, cog_name): + self.cog_name = cog_name + super().__init__(f"The module `{cog_name}` is disabled in this server.") + +class CommandPermissionError(commands.CheckFailure): + """Custom exception for insufficient command permissions based on roles.""" + def __init__(self, command_name): + self.command_name = command_name + super().__init__(f"You do not have the required role to use the command `{command_name}`.") + +@bot.before_invoke +async def global_command_checks(ctx: commands.Context): + """Global check run before any command invocation.""" + # Ignore checks for DMs (or apply different logic if needed) + if not ctx.guild: + return + + # Ignore checks for the bot owner + if await bot.is_owner(ctx.author): + return + + command = ctx.command + if not command: # Should not happen with prefix commands, but good practice + return + + cog = command.cog + cog_name = cog.qualified_name if cog else None + command_name = command.qualified_name + guild_id = ctx.guild.id + + # Ensure author is a Member to get roles + if not isinstance(ctx.author, discord.Member): + log.warning(f"Could not perform permission check for user {ctx.author.id} (not a Member object). Allowing command '{command_name}'.") + return # Cannot check roles if not a Member object + + member_roles_ids = [role.id for role in ctx.author.roles] + + # 1. Check if the Cog is enabled + # Use CORE_COGS attached to the bot instance + if cog_name and cog_name not in bot.core_cogs: # Don't disable core cogs + # Assuming default is True if not explicitly set in DB + is_enabled = await settings_manager.is_cog_enabled(guild_id, cog_name, default_enabled=True) + if not is_enabled: + log.warning(f"Command '{command_name}' blocked in guild {guild_id}: Cog '{cog_name}' is disabled.") + raise CogDisabledError(cog_name) + + # 2. Check command permissions based on roles + # This check only applies if specific permissions HAVE been set for this command. + # If no permissions are set in the DB, check_command_permission returns True. + has_perm = await settings_manager.check_command_permission(guild_id, command_name, member_roles_ids) + if not has_perm: + log.warning(f"Command '{command_name}' blocked for user {ctx.author.id} in guild {guild_id}: Insufficient role permissions.") + raise CommandPermissionError(command_name) + + # If both checks pass, the command proceeds. + log.debug(f"Command '{command_name}' passed global checks for user {ctx.author.id} in guild {guild_id}.") + + +# --- Bot Commands --- + @commands.command(name="restart", help="Restarts the bot. Owner only.") @commands.is_owner() async def restart(ctx): @@ -252,10 +358,13 @@ async def main(args): # Pass parsed args else: bot.ai_cogs_to_skip = [] # Ensure it exists even if empty + # Initialize pools before starting the bot logic + await settings_manager.initialize_pools() try: async with bot: # Load all cogs from the 'cogs' directory, skipping AI if requested + # This should now include WelcomeCog and SettingsCog if they are in the cogs dir await load_all_cogs(bot, skip_cogs=ai_cogs_to_skip) # --- Share GurtCog instance with the sync API --- @@ -297,7 +406,9 @@ async def main(args): # Pass parsed args finally: # Terminate the Flask server process when the bot stops flask_process.terminate() - print("Flask server process terminated.") + log.info("Flask server process terminated.") + # Close database/cache pools + await settings_manager.close_pools() # Run the main async function if __name__ == '__main__': @@ -314,6 +425,7 @@ if __name__ == '__main__': try: asyncio.run(main(args)) # Pass parsed args to main except KeyboardInterrupt: - print("Bot stopped by user.") + log.info("Bot stopped by user.") except Exception as e: - print(f"An error occurred running the bot: {e}") + log.exception(f"An error occurred running the bot: {e}") + # The finally block with pool closing is now correctly inside the main() function diff --git a/requirements.txt b/requirements.txt index d0e13a6..dcc9a98 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,8 @@ aiohttp discord.py python-dotenv -psycopg2-binary +asyncpg +redis>=4.2 # For redis.asyncio Flask Flask-Cors gunicorn diff --git a/settings_manager.py b/settings_manager.py new file mode 100644 index 0000000..ad4ed0d --- /dev/null +++ b/settings_manager.py @@ -0,0 +1,496 @@ +import asyncpg +import redis.asyncio as redis +import os +import logging +from dotenv import load_dotenv + +# Load environment variables +load_dotenv(dotenv_path=os.path.join(os.path.dirname(__file__), '.env')) + +# --- Configuration --- +POSTGRES_USER = os.getenv("POSTGRES_USER") +POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD") +POSTGRES_HOST = os.getenv("POSTGRES_HOST") +POSTGRES_DB = os.getenv("POSTGRES_SETTINGS_DB") # Use the new settings DB +REDIS_HOST = os.getenv("REDIS_HOST") +REDIS_PORT = os.getenv("REDIS_PORT", 6379) +REDIS_PASSWORD = os.getenv("REDIS_PASSWORD") # Optional + +DATABASE_URL = f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{POSTGRES_HOST}/{POSTGRES_DB}" +REDIS_URL = f"redis://{':' + REDIS_PASSWORD + '@' if REDIS_PASSWORD else ''}{REDIS_HOST}:{REDIS_PORT}/0" # Use DB 0 for settings cache + +# --- Global Connection Pools --- +pg_pool = None +redis_pool = None + +# --- Logging --- +log = logging.getLogger(__name__) + +# --- Connection Management --- +async def initialize_pools(): + """Initializes the PostgreSQL and Redis connection pools.""" + global pg_pool, redis_pool + log.info("Initializing database and cache connection pools...") + try: + pg_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=10) + log.info(f"PostgreSQL pool connected to {POSTGRES_HOST}/{POSTGRES_DB}") + + redis_pool = redis.from_url(REDIS_URL, decode_responses=True) + await redis_pool.ping() # Test connection + log.info(f"Redis pool connected to {REDIS_HOST}:{REDIS_PORT}") + + await initialize_database() # Ensure tables exist + except Exception as e: + log.exception(f"Failed to initialize connection pools: {e}") + # Depending on bot structure, might want to raise or exit here + raise + +async def close_pools(): + """Closes the PostgreSQL and Redis connection pools gracefully.""" + global pg_pool, redis_pool + log.info("Closing database and cache connection pools...") + if redis_pool: + try: + await redis_pool.close() + log.info("Redis pool closed.") + except Exception as e: + log.exception(f"Error closing Redis pool: {e}") + redis_pool = None # Ensure it's marked as closed + + if pg_pool: + try: + await pg_pool.close() + log.info("PostgreSQL pool closed.") + except Exception as e: + log.exception(f"Error closing PostgreSQL pool: {e}") + pg_pool = None # Ensure it's marked as closed + + +# --- Database Schema Initialization --- +async def initialize_database(): + """Creates necessary tables in the PostgreSQL database if they don't exist.""" + if not pg_pool: + log.error("PostgreSQL pool not initialized. Cannot initialize database.") + return + + log.info("Initializing database schema...") + async with pg_pool.acquire() as conn: + async with conn.transaction(): + # Guilds table (to track known guilds, maybe store basic info later) + await conn.execute(""" + CREATE TABLE IF NOT EXISTS guilds ( + guild_id BIGINT PRIMARY KEY + ); + """) + + # Guild Settings table (key-value store for various settings) + await conn.execute(""" + CREATE TABLE IF NOT EXISTS guild_settings ( + guild_id BIGINT NOT NULL, + setting_key TEXT NOT NULL, + setting_value TEXT, + PRIMARY KEY (guild_id, setting_key), + FOREIGN KEY (guild_id) REFERENCES guilds(guild_id) ON DELETE CASCADE + ); + """) + # Example setting_keys: 'prefix', 'welcome_channel_id', 'welcome_message', 'goodbye_channel_id', 'goodbye_message' + + # Enabled Cogs table - Stores the explicit enabled/disabled state + await conn.execute(""" + CREATE TABLE IF NOT EXISTS enabled_cogs ( + guild_id BIGINT NOT NULL, + cog_name TEXT NOT NULL, + enabled BOOLEAN NOT NULL, + PRIMARY KEY (guild_id, cog_name), + FOREIGN KEY (guild_id) REFERENCES guilds(guild_id) ON DELETE CASCADE + ); + """) + + # Command Permissions table (simple role-based for now) + await conn.execute(""" + CREATE TABLE IF NOT EXISTS command_permissions ( + guild_id BIGINT NOT NULL, + command_name TEXT NOT NULL, + allowed_role_id BIGINT NOT NULL, + PRIMARY KEY (guild_id, command_name, allowed_role_id), + FOREIGN KEY (guild_id) REFERENCES guilds(guild_id) ON DELETE CASCADE + ); + """) + # Consider adding indexes later for performance on large tables + # await conn.execute("CREATE INDEX IF NOT EXISTS idx_guild_settings_guild ON guild_settings (guild_id);") + # await conn.execute("CREATE INDEX IF NOT EXISTS idx_enabled_cogs_guild ON enabled_cogs (guild_id);") + # await conn.execute("CREATE INDEX IF NOT EXISTS idx_command_permissions_guild ON command_permissions (guild_id);") + + log.info("Database schema initialization complete.") + + +# --- Helper Functions --- +def _get_redis_key(guild_id: int, key_type: str, identifier: str = None) -> str: + """Generates a standardized Redis key.""" + if identifier: + return f"guild:{guild_id}:{key_type}:{identifier}" + return f"guild:{guild_id}:{key_type}" + +# --- Settings Access Functions (Placeholders with Cache Logic) --- + +async def get_guild_prefix(guild_id: int, default_prefix: str) -> str: + """Gets the command prefix for a guild, checking cache first.""" + if not pg_pool or not redis_pool: + log.warning("Pools not initialized, returning default prefix.") + return default_prefix + + cache_key = _get_redis_key(guild_id, "prefix") + try: + cached_prefix = await redis_pool.get(cache_key) + if cached_prefix is not None: + log.debug(f"Cache hit for prefix (Guild: {guild_id})") + return cached_prefix + except Exception as e: + log.exception(f"Redis error getting prefix for guild {guild_id}: {e}") + + log.debug(f"Cache miss for prefix (Guild: {guild_id})") + async with pg_pool.acquire() as conn: + prefix = await conn.fetchval( + "SELECT setting_value FROM guild_settings WHERE guild_id = $1 AND setting_key = 'prefix'", + guild_id + ) + + final_prefix = prefix if prefix is not None else default_prefix + + # Cache the result (even if it's the default, to avoid future DB lookups) + try: + await redis_pool.set(cache_key, final_prefix, ex=3600) # Cache for 1 hour + except Exception as e: + log.exception(f"Redis error setting prefix for guild {guild_id}: {e}") + + return final_prefix + +async def set_guild_prefix(guild_id: int, prefix: str): + """Sets the command prefix for a guild and updates the cache.""" + if not pg_pool or not redis_pool: + log.error("Pools not initialized, cannot set prefix.") + return False # Indicate failure + + cache_key = _get_redis_key(guild_id, "prefix") + try: + async with pg_pool.acquire() as conn: + # Ensure guild exists + await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT (guild_id) DO NOTHING;", guild_id) + # Upsert the setting + await conn.execute( + """ + INSERT INTO guild_settings (guild_id, setting_key, setting_value) + VALUES ($1, 'prefix', $2) + ON CONFLICT (guild_id, setting_key) DO UPDATE SET setting_value = $2; + """, + guild_id, prefix + ) + + # Update cache + await redis_pool.set(cache_key, prefix, ex=3600) # Cache for 1 hour + log.info(f"Set prefix for guild {guild_id} to '{prefix}'") + return True # Indicate success + except Exception as e: + log.exception(f"Database or Redis error setting prefix for guild {guild_id}: {e}") + # Attempt to invalidate cache on error to prevent stale data + try: + await redis_pool.delete(cache_key) + except Exception as redis_err: + log.exception(f"Failed to invalidate Redis cache for prefix (Guild: {guild_id}): {redis_err}") + return False # Indicate failure + +# --- Generic Settings Functions --- + +async def get_setting(guild_id: int, key: str, default=None): + """Gets a specific setting for a guild, checking cache first.""" + if not pg_pool or not redis_pool: + log.warning(f"Pools not initialized, returning default for setting '{key}'.") + return default + + cache_key = _get_redis_key(guild_id, "setting", key) + try: + cached_value = await redis_pool.get(cache_key) + if cached_value is not None: + # Note: Redis stores everything as strings. Consider type conversion if needed. + log.debug(f"Cache hit for setting '{key}' (Guild: {guild_id})") + return cached_value + except Exception as e: + log.exception(f"Redis error getting setting '{key}' for guild {guild_id}: {e}") + + log.debug(f"Cache miss for setting '{key}' (Guild: {guild_id})") + async with pg_pool.acquire() as conn: + value = await conn.fetchval( + "SELECT setting_value FROM guild_settings WHERE guild_id = $1 AND setting_key = $2", + guild_id, key + ) + + final_value = value if value is not None else default + + # Cache the result (even if None or default, cache the absence or default value) + # Store None as a special marker, e.g., "None" string, or handle appropriately + value_to_cache = final_value if final_value is not None else "__NONE__" # Marker for None + try: + await redis_pool.set(cache_key, value_to_cache, ex=3600) # Cache for 1 hour + except Exception as e: + log.exception(f"Redis error setting cache for setting '{key}' for guild {guild_id}: {e}") + + return final_value + + +async def set_setting(guild_id: int, key: str, value: str | None): + """Sets a specific setting for a guild and updates/invalidates the cache. + Setting value to None effectively deletes the setting.""" + if not pg_pool or not redis_pool: + log.error(f"Pools not initialized, cannot set setting '{key}'.") + return False + + cache_key = _get_redis_key(guild_id, "setting", key) + try: + async with pg_pool.acquire() as conn: + # Ensure guild exists + await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT (guild_id) DO NOTHING;", guild_id) + + if value is not None: + # Upsert the setting + await conn.execute( + """ + INSERT INTO guild_settings (guild_id, setting_key, setting_value) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, setting_key) DO UPDATE SET setting_value = $3; + """, + guild_id, key, str(value) # Ensure value is string + ) + # Update cache + await redis_pool.set(cache_key, str(value), ex=3600) + log.info(f"Set setting '{key}' for guild {guild_id}") + else: + # Delete the setting if value is None + await conn.execute( + "DELETE FROM guild_settings WHERE guild_id = $1 AND setting_key = $2", + guild_id, key + ) + # Invalidate cache + await redis_pool.delete(cache_key) + log.info(f"Deleted setting '{key}' for guild {guild_id}") + + return True + except Exception as e: + log.exception(f"Database or Redis error setting setting '{key}' for guild {guild_id}: {e}") + # Attempt to invalidate cache on error + try: + await redis_pool.delete(cache_key) + except Exception as redis_err: + log.exception(f"Failed to invalidate Redis cache for setting '{key}' (Guild: {guild_id}): {redis_err}") + return False + +# --- Cog Enablement Functions --- + +async def is_cog_enabled(guild_id: int, cog_name: str, default_enabled: bool = True) -> bool: + """Checks if a cog is enabled for a guild, checking cache first. + Uses default_enabled if no specific setting is found.""" + if not pg_pool or not redis_pool: + log.warning(f"Pools not initialized, returning default for cog '{cog_name}'.") + return default_enabled + + cache_key = _get_redis_key(guild_id, "cog_enabled", cog_name) + try: + cached_value = await redis_pool.get(cache_key) + if cached_value is not None: + log.debug(f"Cache hit for cog enabled status '{cog_name}' (Guild: {guild_id})") + return cached_value == "True" # Redis stores strings + except Exception as e: + log.exception(f"Redis error getting cog enabled status for '{cog_name}' (Guild: {guild_id}): {e}") + + log.debug(f"Cache miss for cog enabled status '{cog_name}' (Guild: {guild_id})") + db_enabled_status = None + try: + async with pg_pool.acquire() as conn: + db_enabled_status = await conn.fetchval( + "SELECT enabled FROM enabled_cogs WHERE guild_id = $1 AND cog_name = $2", + guild_id, cog_name + ) + except Exception as e: + log.exception(f"Database error getting cog enabled status for '{cog_name}' (Guild: {guild_id}): {e}") + # Fallback to default on DB error after cache miss + return default_enabled + + final_status = db_enabled_status if db_enabled_status is not None else default_enabled + + # Cache the result (True or False) + try: + await redis_pool.set(cache_key, str(final_status), ex=3600) # Cache for 1 hour + except Exception as e: + log.exception(f"Redis error setting cache for cog enabled status '{cog_name}' (Guild: {guild_id}): {e}") + + return final_status + + +async def set_cog_enabled(guild_id: int, cog_name: str, enabled: bool): + """Sets the enabled status for a cog in a guild and updates the cache.""" + if not pg_pool or not redis_pool: + log.error(f"Pools not initialized, cannot set cog enabled status for '{cog_name}'.") + return False + + cache_key = _get_redis_key(guild_id, "cog_enabled", cog_name) + try: + async with pg_pool.acquire() as conn: + # Ensure guild exists + await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT (guild_id) DO NOTHING;", guild_id) + # Upsert the enabled status + await conn.execute( + """ + INSERT INTO enabled_cogs (guild_id, cog_name, enabled) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, cog_name) DO UPDATE SET enabled = $3; + """, + guild_id, cog_name, enabled + ) + + # Update cache + await redis_pool.set(cache_key, str(enabled), ex=3600) + log.info(f"Set cog '{cog_name}' enabled status to {enabled} for guild {guild_id}") + return True + except Exception as e: + log.exception(f"Database or Redis error setting cog enabled status for '{cog_name}' in guild {guild_id}: {e}") + # Attempt to invalidate cache on error + try: + await redis_pool.delete(cache_key) + except Exception as redis_err: + log.exception(f"Failed to invalidate Redis cache for cog enabled status '{cog_name}' (Guild: {guild_id}): {redis_err}") + return False + +# --- Command Permission Functions --- + +async def add_command_permission(guild_id: int, command_name: str, role_id: int) -> bool: + """Adds permission for a role to use a command and invalidates cache.""" + if not pg_pool or not redis_pool: + log.error(f"Pools not initialized, cannot add permission for command '{command_name}'.") + return False + + cache_key = _get_redis_key(guild_id, "cmd_perms", command_name) + try: + async with pg_pool.acquire() as conn: + # Ensure guild exists + await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT (guild_id) DO NOTHING;", guild_id) + # Add the permission rule + await conn.execute( + """ + INSERT INTO command_permissions (guild_id, command_name, allowed_role_id) + VALUES ($1, $2, $3) + ON CONFLICT (guild_id, command_name, allowed_role_id) DO NOTHING; + """, + guild_id, command_name, role_id + ) + + # Invalidate cache after DB operation succeeds + await redis_pool.delete(cache_key) + log.info(f"Added permission for role {role_id} to use command '{command_name}' in guild {guild_id}") + return True + except Exception as e: + log.exception(f"Database or Redis error adding permission for command '{command_name}' in guild {guild_id}: {e}") + # Attempt to invalidate cache even on error + try: + await redis_pool.delete(cache_key) + except Exception as redis_err: + log.exception(f"Failed to invalidate Redis cache for command permissions '{command_name}' (Guild: {guild_id}): {redis_err}") + return False + + +async def remove_command_permission(guild_id: int, command_name: str, role_id: int) -> bool: + """Removes permission for a role to use a command and invalidates cache.""" + if not pg_pool or not redis_pool: + log.error(f"Pools not initialized, cannot remove permission for command '{command_name}'.") + return False + + cache_key = _get_redis_key(guild_id, "cmd_perms", command_name) + try: + async with pg_pool.acquire() as conn: + # Ensure guild exists (though unlikely to be needed for delete) + # await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT (guild_id) DO NOTHING;", guild_id) + # Remove the permission rule + await conn.execute( + """ + DELETE FROM command_permissions + WHERE guild_id = $1 AND command_name = $2 AND allowed_role_id = $3; + """, + guild_id, command_name, role_id + ) + + # Invalidate cache after DB operation succeeds + await redis_pool.delete(cache_key) + log.info(f"Removed permission for role {role_id} to use command '{command_name}' in guild {guild_id}") + return True + except Exception as e: + log.exception(f"Database or Redis error removing permission for command '{command_name}' in guild {guild_id}: {e}") + # Attempt to invalidate cache even on error + try: + await redis_pool.delete(cache_key) + except Exception as redis_err: + log.exception(f"Failed to invalidate Redis cache for command permissions '{command_name}' (Guild: {guild_id}): {redis_err}") + return False + + +async def check_command_permission(guild_id: int, command_name: str, member_roles_ids: list[int]) -> bool: + """Checks if any of the member's roles have permission for the command. + Returns True if allowed, False otherwise. + If no permissions are set for the command in the DB, it defaults to allowed by this check. + """ + if not pg_pool or not redis_pool: + log.warning(f"Pools not initialized, defaulting to allowed for command '{command_name}'.") + return True # Default to allowed if system isn't ready + + cache_key = _get_redis_key(guild_id, "cmd_perms", command_name) + allowed_role_ids_str = set() + + try: + # Check cache first - stores a set of allowed role IDs as strings + if await redis_pool.exists(cache_key): + cached_roles = await redis_pool.smembers(cache_key) + # Handle the empty set marker + if cached_roles == {"__EMPTY_SET__"}: + log.debug(f"Cache hit (empty set) for cmd perms '{command_name}' (Guild: {guild_id}). Command allowed by default.") + return True # No specific restrictions found + allowed_role_ids_str = cached_roles + log.debug(f"Cache hit for cmd perms '{command_name}' (Guild: {guild_id})") + else: + # Cache miss - fetch from DB + log.debug(f"Cache miss for cmd perms '{command_name}' (Guild: {guild_id})") + async with pg_pool.acquire() as conn: + records = await conn.fetch( + "SELECT allowed_role_id FROM command_permissions WHERE guild_id = $1 AND command_name = $2", + guild_id, command_name + ) + # Convert fetched role IDs (BIGINT) to strings for Redis set + allowed_role_ids_str = {str(record['allowed_role_id']) for record in records} + + # Cache the result (even if empty) + try: + async with redis_pool.pipeline(transaction=True) as pipe: + pipe.delete(cache_key) # Ensure clean state + if allowed_role_ids_str: + pipe.sadd(cache_key, *allowed_role_ids_str) + else: + pipe.sadd(cache_key, "__EMPTY_SET__") # Marker for empty set + pipe.expire(cache_key, 3600) # Cache for 1 hour + await pipe.execute() + except Exception as e: + log.exception(f"Redis error setting cache for cmd perms '{command_name}' (Guild: {guild_id}): {e}") + + except Exception as e: + log.exception(f"Error checking command permission for '{command_name}' (Guild: {guild_id}): {e}") + return True # Default to allowed on error + + # --- Permission Check Logic --- + if not allowed_role_ids_str or allowed_role_ids_str == {"__EMPTY_SET__"}: + # If no permissions are defined in our system for this command, allow it. + # Other checks (like @commands.is_owner()) might still apply. + return True + else: + # Check if any of the member's roles intersect with the allowed roles + member_roles_ids_str = {str(role_id) for role_id in member_roles_ids} + if member_roles_ids_str.intersection(allowed_role_ids_str): + log.debug(f"Permission granted for '{command_name}' (Guild: {guild_id}) via role intersection.") + return True # Member has at least one allowed role + else: + log.debug(f"Permission denied for '{command_name}' (Guild: {guild_id}). Member roles {member_roles_ids_str} not in allowed roles {allowed_role_ids_str}.") + return False # Member has none of the specifically allowed roles