883 lines
34 KiB
Python
883 lines
34 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
|
|
intents.presences = True # Required for .status / .activity
|
|
|
|
|
|
# --- 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 separately if needed
|
|
# Note: load_all_cogs already loads all cogs in the directory. The
|
|
# manual load below caused the cog to be loaded twice resulting in an
|
|
# ExtensionAlreadyLoaded error.
|
|
# await self.load_extension("cogs.lockdown_cog")
|
|
|
|
# --- 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.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 https://git.slipstreamm.dev/slipstream/discordbot.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
|