This commit adds the loading of the lockdown cog in main.py. It also includes error handling to catch and log any exceptions that may occur during the loading process, improving the bot's robustness.
730 lines
32 KiB
Python
730 lines
32 KiB
Python
import asyncio
|
|
# Set the event loop policy to the default asyncio policy BEFORE other asyncio/discord imports
|
|
# This is to test if uvloop (if active globally) is causing issues with asyncpg.
|
|
asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
|
|
|
|
import threading
|
|
import discord
|
|
from discord.ext import commands
|
|
import os
|
|
from dotenv import load_dotenv
|
|
import sys
|
|
import asyncio
|
|
import subprocess
|
|
import importlib.util
|
|
import argparse
|
|
import logging # Add logging
|
|
import asyncpg
|
|
import redis.asyncio as aioredis
|
|
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
|
|
import settings_manager # Import the settings manager
|
|
from db import mod_log_db # Import the new mod log db functions
|
|
import command_customization # Import command customization utilities
|
|
from global_bot_accessor import set_bot_instance # Import the new accessor
|
|
import custom_bot_manager # Import the custom bot manager
|
|
|
|
# Import the unified API service runner and the sync API module
|
|
import sys
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
|
from run_unified_api import start_api_in_thread
|
|
import discord_bot_sync_api # Import the module to set the cog instance
|
|
|
|
# Import the markdown server
|
|
from run_markdown_server import start_markdown_server_in_thread
|
|
|
|
# Check if API dependencies are available
|
|
try:
|
|
import uvicorn
|
|
API_AVAILABLE = True
|
|
except ImportError:
|
|
print("uvicorn not available. API service will not be available.")
|
|
API_AVAILABLE = False
|
|
|
|
# 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, but disables mention as prefix."""
|
|
if not message.guild:
|
|
# Use default prefix in DMs
|
|
return DEFAULT_PREFIX
|
|
|
|
# Fetch prefix from settings manager (cache first, then DB)
|
|
prefix = await settings_manager.get_guild_prefix(message.guild.id, DEFAULT_PREFIX)
|
|
return prefix
|
|
|
|
# --- Bot Setup ---
|
|
# Set up intents (permissions)
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.members = True
|
|
|
|
# --- Custom Bot Class with setup_hook for async initialization ---
|
|
class MyBot(commands.Bot):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.owner_id = int(os.getenv('OWNER_USER_ID'))
|
|
self.core_cogs = CORE_COGS # Attach core cogs list to bot instance
|
|
self.settings_manager = settings_manager # Attach settings manager instance
|
|
self.pg_pool = None # Will be initialized in setup_hook
|
|
self.redis = None # Will be initialized in setup_hook
|
|
self.ai_cogs_to_skip = [] # For --disable-ai flag
|
|
|
|
async def setup_hook(self):
|
|
log.info("Running setup_hook...")
|
|
|
|
# Create Postgres pool on this loop
|
|
self.pg_pool = await asyncpg.create_pool(
|
|
dsn=settings_manager.DATABASE_URL, # Use DATABASE_URL from settings_manager
|
|
min_size=1,
|
|
max_size=10,
|
|
loop=self.loop # Explicitly use the bot's event loop
|
|
)
|
|
log.info("Postgres pool initialized and attached to bot.pg_pool.")
|
|
|
|
# Create Redis client on this loop
|
|
self.redis = await aioredis.from_url(
|
|
settings_manager.REDIS_URL, # Use REDIS_URL from settings_manager
|
|
max_connections=10,
|
|
decode_responses=True,
|
|
)
|
|
log.info("Redis client initialized and attached to bot.redis.")
|
|
|
|
# Make sure the bot instance is set in the global_bot_accessor
|
|
# This ensures settings_manager can access the pools via get_bot_instance()
|
|
set_bot_instance(self)
|
|
log.info("Bot instance set in global_bot_accessor from setup_hook.")
|
|
|
|
# Initialize database schema and run migrations using settings_manager
|
|
if self.pg_pool and self.redis:
|
|
try:
|
|
await settings_manager.initialize_database() # Uses the bot instance via get_bot_instance()
|
|
log.info("Database schema initialization called via settings_manager.")
|
|
await settings_manager.run_migrations() # Uses the bot instance via get_bot_instance()
|
|
log.info("Database migrations called via settings_manager.")
|
|
except Exception as e:
|
|
log.exception("CRITICAL: Failed during settings_manager database setup (init/migrations).")
|
|
else:
|
|
log.error("CRITICAL: pg_pool or redis_client not initialized in setup_hook. Cannot proceed with settings_manager setup.")
|
|
|
|
|
|
# Setup the moderation log table *after* pool initialization
|
|
if self.pg_pool:
|
|
try:
|
|
await mod_log_db.setup_moderation_log_table(self.pg_pool)
|
|
log.info("Moderation log table setup complete via setup_hook.")
|
|
except Exception as e:
|
|
log.exception("CRITICAL: Failed to setup moderation log table in setup_hook.")
|
|
else:
|
|
log.warning("pg_pool not available in setup_hook, skipping mod_log_db setup.")
|
|
|
|
# Load all cogs from the 'cogs' directory, skipping AI if requested
|
|
await load_all_cogs(self, skip_cogs=self.ai_cogs_to_skip)
|
|
log.info(f"Cogs loaded in setup_hook. Skipped: {self.ai_cogs_to_skip or 'None'}")
|
|
|
|
# Load the lockdown cog
|
|
try:
|
|
await self.load_extension("cogs.lockdown_cog")
|
|
log.info("Lockdown cog loaded successfully.")
|
|
except Exception as e:
|
|
log.exception(f"Failed to load lockdown cog: {e}")
|
|
|
|
# --- Share GurtCog, ModLogCog, and bot instance with the sync API ---
|
|
try:
|
|
gurt_cog = self.get_cog("Gurt")
|
|
if gurt_cog:
|
|
discord_bot_sync_api.gurt_cog_instance = gurt_cog
|
|
log.info("Successfully shared GurtCog instance with discord_bot_sync_api via setup_hook.")
|
|
else:
|
|
log.warning("GurtCog not found after loading cogs in setup_hook.")
|
|
|
|
discord_bot_sync_api.bot_instance = self
|
|
log.info("Successfully shared bot instance with discord_bot_sync_api via setup_hook.")
|
|
|
|
mod_log_cog = self.get_cog("ModLogCog")
|
|
if mod_log_cog:
|
|
discord_bot_sync_api.mod_log_cog_instance = mod_log_cog
|
|
log.info("Successfully shared ModLogCog instance with discord_bot_sync_api via setup_hook.")
|
|
else:
|
|
log.warning("ModLogCog not found after loading cogs in setup_hook.")
|
|
except Exception as e:
|
|
log.exception(f"Error sharing instances with discord_bot_sync_api in setup_hook: {e}")
|
|
|
|
# --- Manually Load FreakTetoCog (only if AI is NOT disabled) ---
|
|
if not self.ai_cogs_to_skip: # Check if list is empty (meaning AI is not disabled)
|
|
try:
|
|
freak_teto_cog_path = "freak_teto.cog"
|
|
await self.load_extension(freak_teto_cog_path)
|
|
log.info(f"Successfully loaded FreakTetoCog from {freak_teto_cog_path} in setup_hook.")
|
|
except commands.ExtensionAlreadyLoaded:
|
|
log.info(f"FreakTetoCog ({freak_teto_cog_path}) already loaded (setup_hook).")
|
|
except commands.ExtensionNotFound:
|
|
log.error(f"Error: FreakTetoCog not found at {freak_teto_cog_path} (setup_hook).")
|
|
except Exception as e:
|
|
log.exception(f"Failed to load FreakTetoCog in setup_hook: {e}")
|
|
log.info("setup_hook completed.")
|
|
|
|
# Create bot instance using the custom class
|
|
bot = MyBot(command_prefix=get_prefix, intents=intents)
|
|
|
|
# --- 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():
|
|
log.info(f'{bot.user.name} has connected to Discord!')
|
|
log.info(f'Bot ID: {bot.user.id}')
|
|
# Set the bot's status
|
|
await bot.change_presence(activity=discord.Activity(type=discord.ActivityType.listening, name="!help"))
|
|
log.info("Bot status set to 'Listening to !help'")
|
|
|
|
# --- Add current guilds to DB ---
|
|
if bot.pg_pool:
|
|
log.info("Syncing guilds with database...")
|
|
try:
|
|
async with bot.pg_pool.acquire() as conn:
|
|
# Get guilds bot is currently in
|
|
current_guild_ids = {guild.id for guild in bot.guilds}
|
|
log.debug(f"Bot is currently in {len(current_guild_ids)} guilds.")
|
|
|
|
# Get guilds currently in DB
|
|
db_records = await conn.fetch("SELECT guild_id FROM guilds")
|
|
db_guild_ids = {record['guild_id'] for record in db_records}
|
|
log.debug(f"Found {len(db_guild_ids)} guilds in database.")
|
|
|
|
# Add guilds bot joined while offline
|
|
guilds_to_add = current_guild_ids - db_guild_ids
|
|
if guilds_to_add:
|
|
log.info(f"Adding {len(guilds_to_add)} new guilds to database: {guilds_to_add}")
|
|
await conn.executemany("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT DO NOTHING;",
|
|
[(guild_id,) for guild_id in guilds_to_add])
|
|
|
|
# Remove guilds bot left while offline
|
|
guilds_to_remove = db_guild_ids - current_guild_ids
|
|
if guilds_to_remove:
|
|
log.info(f"Removing {len(guilds_to_remove)} guilds from database: {guilds_to_remove}")
|
|
await conn.execute("DELETE FROM guilds WHERE guild_id = ANY($1::bigint[])", list(guilds_to_remove))
|
|
|
|
log.info("Guild sync with database complete.")
|
|
except Exception as e:
|
|
log.exception("Error syncing guilds with database on ready.")
|
|
else:
|
|
log.warning("Bot Postgres pool not initialized, skipping guild sync.")
|
|
# -----------------------------
|
|
|
|
# Patch Discord methods to store message content
|
|
try:
|
|
patch_discord_methods()
|
|
print("Discord methods patched to store message content for error handling")
|
|
|
|
# Make the store_interaction_content function available globally
|
|
import builtins
|
|
builtins.store_interaction_content = store_interaction_content
|
|
print("Made store_interaction_content available globally")
|
|
except Exception as e:
|
|
print(f"Warning: Failed to patch Discord methods: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
try:
|
|
print("Starting command sync process...")
|
|
# List commands before sync
|
|
commands_before = [cmd.name for cmd in bot.tree.get_commands()]
|
|
print(f"Commands before sync: {commands_before}")
|
|
|
|
# Skip global command sync to avoid duplication
|
|
print("Skipping global command sync to avoid command duplication...")
|
|
|
|
# Only sync guild-specific commands with customizations
|
|
print("Syncing guild-specific command customizations...")
|
|
guild_syncs = await command_customization.register_all_guild_commands(bot)
|
|
|
|
total_guild_syncs = sum(len(cmds) for cmds in guild_syncs.values())
|
|
print(f"Synced commands for {len(guild_syncs)} guilds with a total of {total_guild_syncs} customized commands")
|
|
|
|
# List commands after sync
|
|
commands_after = [cmd.name for cmd in bot.tree.get_commands()]
|
|
print(f"Commands registered in command tree: {commands_after}")
|
|
|
|
except Exception as e:
|
|
print(f"Failed to sync commands: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# Start custom bots for users who have enabled them
|
|
try:
|
|
log.info("Starting custom bots for users...")
|
|
|
|
# Import the API database to access user settings
|
|
try:
|
|
from api_service.api_server import db
|
|
if not db:
|
|
log.warning("API database not initialized, cannot start custom bots")
|
|
return
|
|
except ImportError:
|
|
log.warning("Could not import API database, cannot start custom bots")
|
|
return
|
|
|
|
# Get all user settings from the API database
|
|
user_settings_dict = db.user_settings
|
|
enabled_bots = 0
|
|
|
|
for user_id, settings in user_settings_dict.items():
|
|
# Check if custom bot is enabled and has a token
|
|
if settings.custom_bot_enabled and settings.custom_bot_token:
|
|
enabled_bots += 1
|
|
token = settings.custom_bot_token
|
|
prefix = settings.custom_bot_prefix
|
|
status_type = settings.custom_bot_status_type
|
|
status_text = settings.custom_bot_status_text
|
|
|
|
log.info(f"Creating and starting custom bot for user {user_id}")
|
|
# Create the bot
|
|
success, message = await custom_bot_manager.create_custom_bot(
|
|
user_id=user_id,
|
|
token=token,
|
|
prefix=prefix,
|
|
status_type=status_type,
|
|
status_text=status_text
|
|
)
|
|
|
|
if success:
|
|
# Start the bot
|
|
success, message = custom_bot_manager.run_custom_bot_in_thread(
|
|
user_id=user_id,
|
|
token=token
|
|
)
|
|
|
|
if success:
|
|
log.info(f"Successfully started custom bot for user {user_id}")
|
|
else:
|
|
log.error(f"Failed to start custom bot for user {user_id}: {message}")
|
|
else:
|
|
log.error(f"Failed to create custom bot for user {user_id}: {message}")
|
|
|
|
log.info(f"Found {enabled_bots} users with custom bots enabled")
|
|
except Exception as e:
|
|
log.exception(f"Error starting custom bots: {e}")
|
|
|
|
@bot.event
|
|
async def on_shard_disconnect(shard_id):
|
|
print(f"Shard {shard_id} disconnected. Attempting to reconnect...")
|
|
try:
|
|
await bot.connect(reconnect=True)
|
|
print(f"Shard {shard_id} reconnected successfully.")
|
|
except Exception as e:
|
|
print(f"Failed to reconnect shard {shard_id}: {e}")
|
|
|
|
@bot.event
|
|
async def on_guild_join(guild: discord.Guild):
|
|
"""Adds guild to database when bot joins and syncs commands."""
|
|
log.info(f"Joined guild: {guild.name} ({guild.id})")
|
|
if bot.pg_pool:
|
|
try:
|
|
async with bot.pg_pool.acquire() as conn:
|
|
await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT DO NOTHING;", guild.id)
|
|
log.info(f"Added guild {guild.id} to database.")
|
|
|
|
# Sync commands for the new guild
|
|
try:
|
|
log.info(f"Syncing commands for new guild: {guild.name} ({guild.id})")
|
|
synced = await command_customization.register_guild_commands(bot, guild)
|
|
log.info(f"Synced {len(synced)} commands for guild {guild.id}")
|
|
except Exception as e:
|
|
log.exception(f"Failed to sync commands for new guild {guild.id}: {e}")
|
|
except Exception as e:
|
|
log.exception(f"Failed to add guild {guild.id} to database on join.")
|
|
else:
|
|
log.warning("Bot Postgres pool not initialized, cannot add guild on join.")
|
|
|
|
@bot.event
|
|
async def on_guild_remove(guild: discord.Guild):
|
|
"""Removes guild from database when bot leaves."""
|
|
log.info(f"Left guild: {guild.name} ({guild.id})")
|
|
if bot.pg_pool:
|
|
try:
|
|
async with bot.pg_pool.acquire() as conn:
|
|
# Note: Cascading deletes should handle related settings in other tables
|
|
await conn.execute("DELETE FROM guilds WHERE guild_id = $1", guild.id)
|
|
log.info(f"Removed guild {guild.id} from database.")
|
|
except Exception as e:
|
|
log.exception(f"Failed to remove guild {guild.id} from database on leave.")
|
|
else:
|
|
log.warning("Bot Postgres pool not initialized, cannot remove guild on leave.")
|
|
|
|
|
|
# Error handling - Updated to handle custom check failures
|
|
@bot.event
|
|
async def on_command_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}")
|
|
elif isinstance(error, CommandDisabledError):
|
|
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}")
|
|
# Import here to avoid circular imports
|
|
from cogs.ban_system_cog import UserBannedError
|
|
if isinstance(error, UserBannedError):
|
|
await ctx.send(error.message, ephemeral=True) # Send the custom ban message
|
|
log.warning(f"Command '{ctx.command.qualified_name}' blocked for banned user {ctx.author.id}: {error.message}")
|
|
else:
|
|
# Pass other errors to the original handler
|
|
await handle_error(ctx, error)
|
|
|
|
@bot.tree.error
|
|
async def on_app_command_error(interaction, error):
|
|
# Import here to avoid circular imports
|
|
from cogs.ban_system_cog import UserBannedError
|
|
if isinstance(error, UserBannedError):
|
|
if not interaction.response.is_done():
|
|
await interaction.response.send_message(error.message, ephemeral=True)
|
|
else:
|
|
await interaction.followup.send(error.message, ephemeral=True)
|
|
log.warning(f"Command blocked for banned user {interaction.user.id}: {error.message}")
|
|
else:
|
|
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.
|
|
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}`.")
|
|
|
|
class CommandDisabledError(commands.CheckFailure):
|
|
"""Custom exception for disabled commands."""
|
|
def __init__(self, command_name):
|
|
self.command_name = command_name
|
|
super().__init__(f"The command `{command_name}` is disabled in this server.")
|
|
|
|
@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 if the Command is enabled
|
|
# This only applies if the command has been explicitly disabled
|
|
is_cmd_enabled = await settings_manager.is_command_enabled(guild_id, command_name, default_enabled=True)
|
|
if not is_cmd_enabled:
|
|
log.warning(f"Command '{command_name}' blocked in guild {guild_id}: Command is disabled.")
|
|
raise CommandDisabledError(command_name)
|
|
|
|
# 3. 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 all 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):
|
|
"""Restarts the bot. (Owner Only)"""
|
|
await ctx.send("Restarting the bot...")
|
|
await bot.close() # Gracefully close the bot
|
|
os.execv(sys.executable, [sys.executable] + sys.argv) # Restart the bot process
|
|
|
|
bot.add_command(restart)
|
|
|
|
@commands.command(name="gitpull_restart", help="Pulls latest code from git and restarts the bot. Owner only.")
|
|
@commands.is_owner()
|
|
async def gitpull_restart(ctx):
|
|
"""Pulls latest code from git and restarts the bot. (Owner Only)"""
|
|
await ctx.send("Pulling latest code from git...")
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"git", "pull",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
output = stdout.decode().strip() + "\n" + stderr.decode().strip()
|
|
if "unstaged changes" in output or "Please commit your changes" in output:
|
|
await ctx.send("Unstaged changes detected. Committing changes before pulling...")
|
|
commit_proc = await asyncio.create_subprocess_exec(
|
|
"git", "commit", "-am", "Git pull and restart command",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
commit_stdout, commit_stderr = await commit_proc.communicate()
|
|
commit_output = commit_stdout.decode().strip() + "\n" + commit_stderr.decode().strip()
|
|
await ctx.send(f"Committed changes:\n```\n{commit_output}\n```Trying git pull again...")
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"git", "pull",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
output = stdout.decode().strip() + "\n" + stderr.decode().strip()
|
|
if proc.returncode == 0:
|
|
await ctx.send(f"Git pull successful:\n```\n{output}\n```Restarting the bot...")
|
|
await bot.close()
|
|
os.execv(sys.executable, [sys.executable] + sys.argv)
|
|
else:
|
|
await ctx.send(f"Git pull failed:\n```\n{output}\n```")
|
|
|
|
bot.add_command(gitpull_restart)
|
|
|
|
@commands.command(name="reload_cogs", help="Reloads all cogs. Owner only.")
|
|
@commands.is_owner()
|
|
async def reload_cogs(ctx):
|
|
"""Reloads all cogs. (Owner Only)"""
|
|
# Access the disable_ai flag from the bot instance or re-parse args if needed
|
|
# For simplicity, assume disable_ai is accessible; otherwise, need a way to pass it.
|
|
# Let's add it to the bot object for easier access later.
|
|
skip_list = getattr(bot, 'ai_cogs_to_skip', [])
|
|
await ctx.send(f"Reloading all cogs... (Skipping: {', '.join(skip_list) or 'None'})")
|
|
reloaded_cogs, failed_reload = await reload_all_cogs(bot, skip_cogs=skip_list)
|
|
if reloaded_cogs:
|
|
await ctx.send(f"Successfully reloaded cogs: {', '.join(reloaded_cogs)}")
|
|
if failed_reload:
|
|
await ctx.send(f"Failed to reload cogs: {', '.join(failed_reload)}")
|
|
|
|
bot.add_command(reload_cogs)
|
|
|
|
@commands.command(name="gitpull_reload", help="Pulls latest code from git and reloads all cogs. Owner only.")
|
|
@commands.is_owner()
|
|
async def gitpull_reload(ctx):
|
|
"""Pulls latest code from git and reloads all cogs. (Owner Only)"""
|
|
# Access the disable_ai flag from the bot instance or re-parse args if needed
|
|
skip_list = getattr(bot, 'ai_cogs_to_skip', [])
|
|
await ctx.send(f"Pulling latest code from git... (Will skip reloading: {', '.join(skip_list) or 'None'})")
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"git", "pull",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
output = stdout.decode().strip() + "\n" + stderr.decode().strip()
|
|
if "unstaged changes" in output or "Please commit your changes" in output:
|
|
await ctx.send("Unstaged changes detected. Committing changes before pulling...")
|
|
commit_proc = await asyncio.create_subprocess_exec(
|
|
"git", "commit", "-am", "Git pull and reload command",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
commit_stdout, commit_stderr = await commit_proc.communicate()
|
|
commit_output = commit_stdout.decode().strip() + "\n" + commit_stderr.decode().strip()
|
|
await ctx.send(f"Committed changes:\n```\n{commit_output}\n```Trying git pull again...")
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"git", "pull",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE
|
|
)
|
|
stdout, stderr = await proc.communicate()
|
|
output = stdout.decode().strip() + "\n" + stderr.decode().strip()
|
|
if proc.returncode == 0:
|
|
await ctx.send(f"Git pull successful:\n```\n{output}\n```Reloading all cogs...")
|
|
reloaded_cogs, failed_reload = await reload_all_cogs(bot, skip_cogs=skip_list)
|
|
if reloaded_cogs:
|
|
await ctx.send(f"Successfully reloaded cogs: {', '.join(reloaded_cogs)}")
|
|
if failed_reload:
|
|
await ctx.send(f"Failed to reload cogs: {', '.join(failed_reload)}")
|
|
else:
|
|
await ctx.send(f"Git pull failed:\n```\n{output}\n```")
|
|
|
|
bot.add_command(gitpull_reload)
|
|
|
|
|
|
|
|
|
|
# The unified API service is now handled by run_unified_api.py
|
|
|
|
async def main(args): # Pass parsed args
|
|
"""Main async function to load cogs and start the bot."""
|
|
TOKEN = os.getenv('DISCORD_TOKEN')
|
|
if not TOKEN:
|
|
raise ValueError("No token found. Make sure to set DISCORD_TOKEN in your .env file.")
|
|
|
|
# Start Flask server as a separate process
|
|
flask_process = subprocess.Popen([sys.executable, "flask_server.py"], cwd=os.path.dirname(__file__))
|
|
|
|
# Start the unified API service in a separate thread if available
|
|
api_thread = None
|
|
if API_AVAILABLE:
|
|
print("Starting unified API service...")
|
|
try:
|
|
# Start the API in a separate thread
|
|
api_thread = start_api_in_thread()
|
|
print("Unified API service started successfully")
|
|
except Exception as e:
|
|
print(f"Failed to start unified API service: {e}")
|
|
|
|
# Start the markdown server for TOS and Privacy Policy
|
|
markdown_thread = None
|
|
try:
|
|
print("Starting markdown server for TOS and Privacy Policy...")
|
|
markdown_thread = start_markdown_server_in_thread(host="0.0.0.0", port=5006)
|
|
print("Markdown server started successfully. TOS available at: http://localhost:5006/tos")
|
|
except Exception as e:
|
|
print(f"Failed to start markdown server: {e}")
|
|
|
|
# Configure OAuth settings from environment variables
|
|
oauth_host = os.getenv("OAUTH_HOST", "0.0.0.0")
|
|
oauth_port = int(os.getenv("OAUTH_PORT", "8080"))
|
|
oauth_redirect_uri = os.getenv("DISCORD_REDIRECT_URI", f"http://{oauth_host}:{oauth_port}/oauth/callback")
|
|
|
|
# Update the OAuth redirect URI in the environment
|
|
os.environ["DISCORD_REDIRECT_URI"] = oauth_redirect_uri
|
|
print(f"OAuth redirect URI set to: {oauth_redirect_uri}")
|
|
|
|
# --- Define AI cogs to potentially skip ---
|
|
ai_cogs_to_skip = []
|
|
if args.disable_ai:
|
|
print("AI functionality disabled via command line flag.")
|
|
ai_cogs_to_skip = [
|
|
"cogs.ai_cog",
|
|
"cogs.multi_conversation_ai_cog",
|
|
# Add any other AI-related cogs from the 'cogs' folder here
|
|
]
|
|
# Store the skip list on the bot object for reload commands
|
|
# This is now done on the bot instance directly in the MyBot class
|
|
bot.ai_cogs_to_skip = ai_cogs_to_skip
|
|
else:
|
|
bot.ai_cogs_to_skip = [] # Ensure it exists even if empty
|
|
|
|
set_bot_instance(bot) # Set the global bot instance
|
|
log.info(f"Global bot instance set in global_bot_accessor. Bot ID: {id(bot)}")
|
|
|
|
# Pool initialization and cog loading are now handled in MyBot.setup_hook()
|
|
|
|
try:
|
|
# The bot will call setup_hook internally after login but before on_ready.
|
|
await bot.start(TOKEN)
|
|
except Exception as e:
|
|
log.exception(f"An error occurred during bot.start(): {e}")
|
|
finally:
|
|
# Terminate the Flask server process when the bot stops
|
|
if flask_process and flask_process.poll() is None: # Check if process exists and is running
|
|
flask_process.terminate()
|
|
log.info("Flask server process terminated.")
|
|
else:
|
|
log.info("Flask server process was not running or already terminated.")
|
|
|
|
# Stop all custom bots using the improved cleanup function
|
|
try:
|
|
custom_bot_manager.cleanup_all_custom_bots()
|
|
except Exception as e:
|
|
log.exception(f"Error during custom bot cleanup: {e}")
|
|
|
|
# Close database/cache pools if they were initialized
|
|
if bot.pg_pool:
|
|
log.info("Closing Postgres pool in main finally block...")
|
|
await bot.pg_pool.close()
|
|
if bot.redis:
|
|
log.info("Closing Redis pool in main finally block...")
|
|
await bot.redis.close()
|
|
if not bot.pg_pool and not bot.redis:
|
|
log.info("Pools were not initialized or already closed, skipping close_pools in main.")
|
|
|
|
# Run the main async function
|
|
import signal
|
|
|
|
def handle_sighup(signum, frame):
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
try:
|
|
print("Received SIGHUP: pulling latest code from /home/git/git (branch master)...")
|
|
result = subprocess.run(
|
|
["git", "pull"],
|
|
capture_output=True, text=True
|
|
)
|
|
print(result.stdout)
|
|
print(result.stderr)
|
|
print("Restarting process after SIGHUP...")
|
|
except Exception as e:
|
|
print(f"Error during SIGHUP git pull: {e}")
|
|
os.execv(sys.executable, [sys.executable] + sys.argv)
|
|
|
|
if __name__ == '__main__':
|
|
# Write PID to .pid file for git hook use
|
|
try:
|
|
with open(".pid", "w") as f:
|
|
f.write(str(os.getpid()))
|
|
except Exception as e:
|
|
print(f"Failed to write .pid file: {e}")
|
|
|
|
# Register SIGHUP handler (Linux only)
|
|
try:
|
|
signal.signal(signal.SIGHUP, handle_sighup)
|
|
except AttributeError:
|
|
print("SIGHUP not available on this platform.")
|
|
|
|
# --- Argument Parsing ---
|
|
parser = argparse.ArgumentParser(description="Run the Discord Bot")
|
|
parser.add_argument(
|
|
"--disable-ai",
|
|
action="store_true",
|
|
help="Disable AI-related cogs and functionality."
|
|
)
|
|
args = parser.parse_args()
|
|
# ------------------------
|
|
|
|
try:
|
|
asyncio.run(main(args)) # Pass parsed args to main
|
|
except KeyboardInterrupt:
|
|
log.info("Bot stopped by user.")
|
|
except Exception as e:
|
|
log.exception(f"An error occurred running the bot: {e}")
|
|
# The finally block with pool closing is now correctly inside the main() function
|