faucet/main.py

import discord
from discord.ext import commands
import asyncio
import os
from dotenv import load_dotenv
from log import logger
import signal
import argparse

# Load .env
load_dotenv()
TOKEN = os.getenv("TOKEN")
PREFIX = os.getenv("PREFIX", "! .").split(" ") 

# Set up argument parsing
parser = argparse.ArgumentParser(description="Selfbot for Discord")
parser.add_argument("--useragent", type=str, help="Specify the user agent")
parser.add_argument("--loglevel", type=str, choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Set the logging level")
args = parser.parse_args()

# Configure logging level based on argument
if args.loglevel:
    logger.setLevel(args.loglevel)

async def load_cogs(bot, directory="cogs"):
    tasks = []
    for filename in os.listdir(directory):
        if filename.endswith(".py") and not filename.startswith("_"):
            ext = f"{directory}.{filename[:-3]}"
            task = asyncio.create_task(bot.load_extension(ext))
            task.ext = ext  # attach for error reporting
            tasks.append(task)
            logger.info(f"Prepared to load {ext}")

    for task in tasks:
        try:
            await task
            logger.info(f"Successfully loaded {task.ext}")
        except Exception as e:
            logger.error(f"Failed to load {task.ext}", exc_info=e)

class SelfBot(commands.Bot):
    def __init__(self):        
        super().__init__(command_prefix=PREFIX, self_bot=True)
        self.logger = logger
        self.self_bot = True
        self.timeout = 10.0  # Default timeout for commands in seconds
        self.config = {}

    async def setup_hook(self):
        await load_cogs(self)
        logger.info("Cogs loaded, bot is ready")

    async def on_ready(self):
        logger.info(f"Logged in as {self.user} (ID: {self.user.id})")
        logger.info("------")
        await self.change_presence(afk=True)

    async def on_message(self, message):
        if message.author != self.user:
            return

        await self.process_commands(message)

    def get_config(self, key, default=None):
        if not hasattr(self, 'config'):
            self.logger.warning(f"Config not available, using default for {key}")
            return default
        
        config_entry = self.config.get(key, None)
        return default if config_entry is None else config_entry

async def shutdown_bot():
    if bot:
        logger.info("Unloading all cogs")
        try:
            await asyncio.gather(*(bot.unload_extension(ext) for ext in list(bot.extensions)))
        except Exception as e:
            logger.error("Error unloading cogs during shutdown", exc_info=e)

        logger.info("Closing bot connection...")
        await bot.close()

def handle_exit(sig, frame):
    logger.info(f"Received signal {sig}, shutting down gracefully...")
    loop = asyncio.get_event_loop()
    loop.create_task(shutdown_bot())

async def main():
    global bot
    if not TOKEN:
        raise RuntimeError("DISCORD_TOKEN not found in .env file")

    bot = SelfBot()
    if getattr(args, "useragent", None):
        bot.http.super_properties["browser_user_agent"] = args.useragent   
    async with bot:
        max_retries = 5
        retry_count = 0
        backoff_time = 5

        while retry_count < max_retries:
            try:
                await bot.start(TOKEN)
                break
            except discord.errors.ConnectionClosed as e:
                retry_count += 1
                logger.error(f"Connection closed. Retrying ({retry_count}/{max_retries}) in {backoff_time}s: {e}")
                if retry_count < max_retries:
                    await asyncio.sleep(backoff_time)
                    backoff_time *= 2
                else:
                    logger.critical("Maximum retries reached. Shutting down.")
            except discord.errors.HTTPException as e:
                if hasattr(e, "retry_after"):
                    retry_after = getattr(e, "retry_after", 5)
                    logger.warning(f"Rate limited. Retrying in {retry_after}s")
                    await asyncio.sleep(retry_after)
                else:
                    logger.error(f"HTTP Error: {e}")
                    break
            except Exception as e:
                logger.error("Unexpected exception during bot start", exc_info=e)
                break

if __name__ == "__main__":
    signal.signal(signal.SIGINT, handle_exit)
    signal.signal(signal.SIGTERM, handle_exit)
    asyncio.run(main())