akn/clientmulti_bot.py

import asyncio
from datetime import datetime as dt
from pyrogram import Client
from pyrogram.types import *
from pyrogram.errors import *
from pyrogram.raw.all import layer
import pyromod
from config import API_HASH, API_ID
from akn import assistant
from akn.utils.database import db
from akn.utils.logger import LOGS
from akn.utils.expired_bot import watch_do_time
from akn import ids, gemini_bot_id
from pyrogram import __version__ as pyrogram_version
from box import Box

async def start_user() -> None:
    sessions = await db.get_all_sessions()
    if not sessions:
        return
    active_clients = []
    for i, session_data in enumerate(sessions, 1):
        api_id = session_data.get("api_id")
        api_hash = session_data.get("api_hash")
        session_str = session_data.get("session")
        user_id = session_data.get("user_id", 0)
        try:
            if not (api_id, api_hash, session_str, user_id):
                continue
            client = Client(
                name=f"UserBot_{i}_{user_id}",
                api_id=api_id,
                api_hash=api_hash,
                session_string=session_str,
                in_memory=False,
                plugins=dict(root="akn.Akeno"),
                workdir="/tmp/akn_sessions",
                app_version="UserBot/latest",
                device_model="Anonymous",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            if me.id != user_id:
                raise ValueError(f"Session user_id mismatch (expected {user_id}, got {me.id})")
            ids.append(user_id)
            LOGS.info(f"✅ Started User #{i}: Name: {me.first_name}")
            active_clients.append(client)
            asyncio.create_task(
                _check_session_health(client, user_id),
                name=f"health_monitor_{user_id}"
            )
        except (
            UserDeactivatedBan,
            AuthKeyDuplicated,
            UserDeactivated,
            AuthKeyUnregistered,
            SessionRevoked
        ) as e:
            await _handle_dead_session(user_id, e)
            continue
        except Exception as e:
            LOGS.error(f"⚠️ User #{i} failed: {type(e).__name__}: {str(e)}")
            continue

async def start_magic_bot() -> None:
    sessions = await db.get_all_magic_bot()
    if not sessions:
        LOGS.warning("No Magic bots found in database")
        return

    active_clients = []
    active_bots = []
    for i, data in enumerate(sessions):
        user_id = data.get("user_id")
        bot_token = data.get("bot_token")
        if not (bot_token, user_id):
            continue
        get_user_exp = await db.get_expired_date(user_id)
        to_obj = Box(get_user_exp or {})
        if not to_obj:
            continue
        if not to_obj.get("user_client"):
            continue
        if to_obj.user_client.get("disconnected", False):
            continue
        try:
            client = Client(
                name=f"MagicBot_{i}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                in_memory=True,
                plugins=dict(root="akn.MagicFonts"),
                workdir="/tmp/akn_magic",
                app_version="MagicBot/v1.1.0",
                device_model="SecureMagicBot",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            active_bots.append((client, user_id))
            active_clients.append(client)
            LOGS.info(f"✅ Started Magic bot {i}: {me.first_name} | {me.id}")
            asyncio.create_task(
                connection_watchdog(client),
                name=f"health_dog_{me.id}"
            )
        except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
            LOGS.warning(f"💀 Bot Magic #{i} session dead: {type(e).__name__}")
            success = await db.remove_bot_token_magic(client.bot_token)
            if not success:
                LOGS.warning(f"⚠️ Failed to revoke token: {client.bot_token}")
                raise
        except Exception as e:
            LOGS.error(f"⚠️ Failed bot Magic #{i}: {type(e).__name__}: {str(e)}")
            await _cleanup_failed_client(client)
            raise

    tasks = []
    for client, user_id in active_bots:
        tasks.append(watch_do_time(user_id, client, assistant))
    if tasks:
        await asyncio.gather(*tasks)

async def start_meta_bot() -> None:
    sessions = await db.get_all_meta_bot()
    if not sessions:
        LOGS.warning("No Meta bots found in database")
        return

    active_clients = []
    active_bots_2 = []
    for i, data in enumerate(sessions):
        user_id = data.get("user_id")
        bot_token = data.get("bot_token")
        if not (bot_token, user_id):
            continue
        get_user_exp = await db.get_expired_date(user_id)
        to_obj = Box(get_user_exp or {})
        if not to_obj:
            continue
        if not to_obj.get("user_client"):
            continue
        if to_obj.user_client.get("disconnected", False):
            continue
        try:
            client = Client(
                name=f"MetaAI_{i}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                in_memory=True,
                plugins=dict(root="akn.Meta"),
                workdir="/tmp/akn_meta",
                app_version="MetaBot/v1.1.0",
                device_model="SecureMetaBot",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            active_bots_2.append((client, user_id))
            active_clients.append(client)
            LOGS.info(f"✅ Started Meta AI bot {i}: {me.first_name} | {me.id}")
            asyncio.create_task(
                connection_watchdog(client),
                name=f"health_dog_{me.id}"
            )
        except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
            LOGS.warning(f"💀 Bot Meta #{i} session dead: {type(e).__name__}")
            success = await db.remove_bot_token_meta(client.bot_token)
            if not success:
                LOGS.warning(f"Failed to revoke token: {client.bot_token}")
                raise
        except Exception as e:
            LOGS.error(f"⚠️ Failed bot Meta #{i}: {type(e).__name__}: {str(e)}")
            await _cleanup_failed_client(client)
            raise

    taskso = []
    for client, user_id in active_bots_2:
        taskso.append(watch_do_time(user_id, client, assistant))
    if taskso:
        await asyncio.gather(*taskso)

async def start_gemini_bot() -> None:
    sessions = await db.get_all_gemini_bot()
    if not sessions:
        LOGS.warning("No Gemini bots found in database")
        return

    active_clients = []

    for i, data in enumerate(sessions, 1):
        try:
            bot_token = data.get("bot_token")
            if not bot_token:
                continue
            client = Client(
                name=f"GeminiBot_{i}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                in_memory=True,
                plugins=dict(root="akn.Gemini"),
                workdir="/tmp/akn_gemini",
                app_version="GeminiBot/v1.1.0",
                device_model="SecureGemini",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            gemini_bot_id.append(me.id)
            active_clients.append(client)
            LOGS.info(f"✅ Started Gemini bot {i}: {me.first_name} | {me.id}")
            asyncio.create_task(
                connection_watchdog(client),
                name=f"health_dog_{me.id}"
            )
        except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
            LOGS.warning(f"💀 Bot Gemini #{i} session dead: {type(e).__name__}")
            success = await db.remove_bot_token_gemini(client.bot_token)
            if not success:
                LOGS.warning(f"⚠️ Failed to revoke token: {client.bot_token}")
                raise
        except Exception as e:
            LOGS.error(f"⚠️ Failed bot Gemini #{i}: {type(e).__name__}: {str(e)}")
            await _cleanup_failed_client(client)
            raise

async def start_youtube_bot() -> None:
    sessions = await db.get_all_youtube_bot()
    if not sessions:
        LOGS.warning("No YouTube bots found in database")
        return

    active_clients = []
    for i, data in enumerate(sessions, 1):
        try:
            bot_token = data.get("bot_token")
            if not bot_token:
                continue
            client = Client(
                name=f"YoutubeBot_{i}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                in_memory=True,
                plugins=dict(root="akn.Youtube"),
                workdir="/tmp/akn_youtubebot",
                app_version="YoutubeBot/v1.1.0",
                device_model="SecureYoutubeBot",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            LOGS.info(f"✅ Started Youtube bot {i}: {me.first_name} | {me.id}")
            asyncio.create_task(
                connection_watchdog(client),
                name=f"health_dog_{me.id}"
            )
        except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
            LOGS.warning(f"💀 Bot YouTube #{i} session dead: {type(e).__name__}")
            success = await db.remove_bot_token_youtube(client.bot_token)
            if not success:
                LOGS.warning(f"⚠️ Failed to revoke token: {client.bot_token}")
                raise
        except Exception as e:
            LOGS.error(f"⚠️ Failed bot Session #{i}: {type(e).__name__}: {str(e)}")
            await _cleanup_failed_client(client)
            raise

async def start_session_bot() -> None:
    sessions = await db.get_all_session_bot()
    if not sessions:
        LOGS.warning("No Session bots found in database")
        return

    active_clients = []
    for i, data in enumerate(sessions, 1):
        try:
            bot_token = data.get("bot_token")
            if not bot_token:
                continue
            client = Client(
                name=f"SessionBot_{i}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                in_memory=True,
                plugins=dict(root="akn.SessionBot"),
                workdir="/tmp/akn_sessionbot",
                app_version="SessionBot/v1.1.0",
                device_model="SecureSessionBot",
                system_version="Linux/Kernel-6.5",
                sleep_threshold=30
            )
            await client.start()
            me = await client.get_me()
            active_clients.append(client)
            LOGS.info(f"✅ Started Session bot {i}: {me.first_name} | {me.id}")
            asyncio.create_task(
                connection_watchdog(client),
                name=f"health_dog_{me.id}"
            )
        except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
            LOGS.warning(f"💀 Bot Session #{i} session dead: {type(e).__name__}")
            success = await db.remove_bot_token_sessionbot(client.bot_token)
            if not success:
                LOGS.warning(f"⚠️ Failed to revoke token: {client.bot_token}")
                raise
        except Exception as e:
            LOGS.error(f"⚠️ Failed bot Session #{i}: {type(e).__name__}: {str(e)}")
            await _cleanup_failed_client(client)
            raise

async def start_captcha_bot() -> None:
    sessions = await db.get_all_captcha_bot()
    for i, data in enumerate(sessions):
        try:
            bot_token = data.get("bot_token")
            client = Client(
                name=f"CaptchaBot#{i + 1}",
                api_id=API_ID,
                api_hash=API_HASH,
                bot_token=bot_token,
                app_version="latest",
                device_model="Captcha Bot",
                system_version="Linux",
                plugins=dict(root="akn.ApproveBot"),
            )
            await client.start()
            me = await client.get_me()
            LOGS.info(f"✅ Started Captcha bot {i + 1}: {me.first_name} | {me.id}")
        except AuthKeyDuplicated:
            continue
        except UserDeactivated:
            continue
        except Exception as e:
            continue

async def start_all_downloader_bot() -> None:
    try:
        bot_sessions = await db.alldl_bot.find({
            "bots": {
                "$elemMatch": {
                    "status": "approved",
                    "is_active": True
                }
            }
        }).to_list(length=None)
        if not bot_sessions:
            LOGS.warning("No active downloader bots found in database")
            return
        LOGS.info(f"✅ Starting {len(bot_sessions)} downloader bots...")
        active_clients = []
        success_count = 0
        
        for i, session_data in enumerate(bot_sessions, 1):
            user_id = session_data.get("user_id")
            bots = session_data.get("bots", [])
            
            for bot_data in bots:
                if not (bot_data.get("status") == "approved" and bot_data.get("is_active")):
                    continue
                    
                bot_token = bot_data.get("bot_token")
                if not bot_token:
                    LOGS.warning(f"Skipping bot #{i} - No token found for user {user_id}")
                    continue
                    
                try:
                    client = Client(
                        name=f"AllDLBot_{user_id}_{i}",
                        api_id=API_ID,
                        api_hash=API_HASH,
                        bot_token=bot_token,
                        in_memory=True,
                        plugins=dict(root="akn.AllDownloaderBot"),
                        workdir="/tmp",
                        app_version="DLBot/v2.0.0",
                        device_model="SecureDownloaderPro",
                        system_version="Linux/Kernel-6.6",
                        sleep_threshold=30
                    )
                    
                    start_success = await _startv_client_downloader_safely(client, i, user_id)
                    if start_success:
                        active_clients.append({
                            "client": client,
                            "user_id": user_id,
                            "bot_data": bot_data
                        })
                        success_count += 1
                        
                        await db.alldl_bot.update_one(
                            {"user_id": user_id, "bots.bot_token": bot_token},
                            {"$set": {"bots.$.last_active": dt.now().isoformat()}}
                        )
                        
                except Exception as e:
                    error_type = type(e).__name__
                    LOGS.error(f"⚠️ Failed to start bot #{i} for user {user_id}: {error_type}: {str(e)}")
                    await db.alldl_bot.update_one(
                        {"user_id": user_id, "bots.bot_token": bot_token},
                        {"$set": {"bots.$.is_active": False, "bots.$.error": str(e)}}
                    )
                    
                    await _cleanupv_failed_client(client)
                    continue

        LOGS.info(f"✅ Successfully started {success_count}/{len(bot_sessions)} bots")

        global storage_running
        storage_running = active_clients
        
    except Exception as e:
        LOGS.critical(f"⚠️ Fatal error in bot startup: {type(e).__name__}: {str(e)}")
        raise

async def _startv_client_downloader_safely(client: Client, index: int, user_id: int) -> bool:
    try:
        await asyncio.wait_for(client.start(), timeout=30)
        
        me = await client.get_me()
        if not me:
            raise ConnectionError("Failed to get bot identity")
            
        LOGS.info(f"✅ Bot #{index} started for user {user_id} - @{me.username}")
        return True
        
    except asyncio.TimeoutError:
        LOGS.warning(f"⌛ Timeout starting bot #{index} for user {user_id}")
    except Exception as e:
        LOGS.error(f"⚠️ Error starting bot #{index} for user {user_id}: {str(e)}")
    return False

async def _cleanupv_failed_client(client: Client) -> None:
    try:
        if client.is_connected:
            await client.stop()
    except Exception as e:
        LOGS.warning(f"Error during client cleanup: {str(e)}")
    finally:
        await client.disconnect()

async def _start_client_downloder_safely(client: Client, bot_num: int) -> None:
    try:
        await client.start()
        me = await client.get_me()
        LOGS.info(
            f"✅ Started Downloader Bot #{bot_num}: "
            f"{me.first_name} (ID: {me.id})"
        )
        asyncio.create_task(
            connection_watchdog(client),
            name=f"health_dog_{me.id}"
        )
    except (AuthKeyDuplicated, UserDeactivated, AccessTokenExpired) as e:
        LOGS.warning(f"💀 Bot Downloader #{bot_num} session dead: {type(e).__name__}")
        success = await db.remove_bot_token_alldlbot(client.bot_token)
        if not success:
            LOGS.warning(f"⚠️ Failed to revoke token: {client.bot_token}")
            raise
    except Exception as e:
        LOGS.error(f"⚠️ Bot Downloader #{bot_num} critical failure: {str(e)}")
        raise

async def _check_session_health(client: Client, user_id: int, interval: int = 300) -> None:
    while True:
        try:
            await asyncio.wait_for(client.get_me(), timeout=10)
            
            if not client.is_connected:
                raise ConnectionError("Client disconnected")
                
            LOGS.debug(f"Session health OK: User {user_id}")
            await asyncio.sleep(interval)
            
        except (UserDeactivated, AuthKeyInvalid) as e:
            LOGS.warning(f"💀 Session dead for {user_id}: {type(e).__name__}")
            await _handle_dead_session(user_id, e)
            break
        except Exception as e:
            LOGS.error(f"Health check failed for {user_id}: {type(e).__name__}: {str(e)}")
            await asyncio.sleep(60)

async def _handle_dead_session(user_id: int, error: Exception) -> None:
    LOGS.warning(f"Cleaning up dead session: {user_id}")
    await _send_message_warning(
        user_id,
        f"🚨 Session terminated\n"
        f"User: {user_id}\n"
        f"Reason: Error: {type(error).__name__}"
    )
    await db.remove_strings_userv(user_id)
    LOGS.warning(
        f"🚨 Session terminated\n"
        f"User: {user_id}\n"
        f"Reason: {type(error).__name__}"
    )

async def check_connection(client: Client) -> bool:
    try:
        return await client.get_me() is not None
    except:
        return False

async def connection_watchdog(client: Client):
    while True:
        if not await check_connection(client):
            LOGS.warning("Reconnecting...")
            await client.disconnect()
            await client.connect()
        await asyncio.sleep(300)

async def _send_message_warning(user_id, text):
    try:
        await assistant.send_message(user_id, text)
    except:
        pass

async def _cleanup_failed_client(client: Client | None) -> None:
    if client:
        try:
            await client.stop()
        except:
            pass
        finally:
            del client