This commit is contained in:
Slipstream 2025-05-03 13:47:49 -06:00
parent 6d6a6a6f29
commit 5c3f0b9810
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
6 changed files with 1051 additions and 9 deletions

View File

@ -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)

208
cogs/settings_cog.py Normal file
View File

@ -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 <new_prefix>`")
@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 <CogName>`")
@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 <CogName>`")
@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 <command_name> <@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 <command_name> <@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.")

212
cogs/welcome_cog.py Normal file
View File

@ -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.")

128
main.py
View File

@ -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

View File

@ -1,7 +1,8 @@
aiohttp
discord.py
python-dotenv
psycopg2-binary
asyncpg
redis>=4.2 # For redis.asyncio
Flask
Flask-Cors
gunicorn

496
settings_manager.py Normal file
View File

@ -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