api/custom_logic.py

import os
import json
import asyncio
import logging
import httpx
import html

# Настройка базового логгера для системных предупреждений
logger = logging.getLogger("cyber_agent_custom_logic")

async def process_custom_message(
    text: str, 
    chat_id: int, 
    history: str, 
    tg_message_id: int = None, 
    page_code: str = None, 
    page_cookies: str = None
) -> str:
    """
    АДМИНИСТРАТИВНЫЙ ИИ-КОНТУР АГЕНТА (ПОЛНОСТЬЮ АСИНХРОННЫЙ)
    Архитектура: Web3CyberServices Core
    Оптимизирован для отказоустойчивости, защиты разметки и расширенного инференса.
    """

    # --- Безопасный изолированный импорт смежных модулей юзербота ---
    try:
        from api.user_client import search_in_group, join_telegram_chat
    except (ImportError, ModuleNotFoundError) as e:
        logger.error(f"Критическое предупреждение среды: модуль api.user_client не найден: {e}")
        search_in_group, join_telegram_chat = None, None

    try:
        from api.assistant_factory import AssistantFactory
        factory = AssistantFactory(logger_agent=None)
    except Exception as e:
        logger.warning(f"Ошибка инициализации фабрики ассистентов: {e}")
        factory = None

    # --- Настройки среды и валидация ACL ---
    tg_token = os.environ.get("TELEGRAM_BOT_TOKEN") or os.environ.get("TELEGRAM_TOKEN")
    tg_base_url = os.environ.get("TELEGRAM_BASE_URL", "https://api.telegram.org").rstrip("/")

    try:
        admin_id_raw = os.environ.get("ADMIN_CHAT_ID", "7933414267")
        admin_id = int(admin_id_raw)
        chat_id = int(chat_id) 
    except (ValueError, TypeError):
        admin_id = 7933414267
        chat_id = 0

    is_admin = (chat_id == admin_id)
    text = str(text).strip()

    # Использование оптимизированных настроек таймаутов для предотвращения зависания потока
    limits = httpx.Limits(max_keepalive_connections=5, max_connections=10)
    async with httpx.AsyncClient(limits=limits, timeout=httpx.Timeout(15.0, connect=5.0)) as client:

        user_name = "Пользователь"
        if tg_token and chat_id != 0:
            try:
                chat_res = await client.post(f"{tg_base_url}/bot{tg_token}/getChat", json={"chat_id": chat_id})
                if chat_res.status_code == 200:
                    user_name = chat_res.json().get("result", {}).get("first_name", "Пользователь")
            except Exception as e:
                logger.warning(f"Не удалось получить имя чата: {e}")

        # Чат-логгер для пересылки сообщений администратору (Защищен HTML-экранированием)
        if not is_admin and tg_token and chat_id != 0:
            try:
                safe_user_name = html.escape(user_name)
                safe_text = html.escape(text)
                log_message = (
                    f"🔔 <b>Чат-логгер:</b>\n"
                    f"Пользователь {safe_user_name} (<code>{chat_id}</code>) написал:\n\n"
                    f"<i>{safe_text}</i>"
                )
                await client.post(
                    f"{tg_base_url}/bot{tg_token}/sendMessage",
                    json={
                        "chat_id": admin_id,
                        "text": log_message,
                        "parse_mode": "HTML"
                    }
                )
            except Exception as e:
                logger.error(f"Ошибка чат-логгера при отправке админу: {e}")

        # =====================================================================
        # БЛОК СИСТЕМНЫХ КОМАНД (ACL ЗАЩИТА)
        # =====================================================================
        if text.startswith(("/", "!")) and not is_admin:
            return f"⚠️ {user_name}, у вашего аккаунта отсутствует доступ к ядру Web3CyberAgent."

        if text.startswith("/send "):
            if not tg_token: return "❌ Ошибка среды: Токен Telegram не найден."
            cmd_body = text[6:].strip().split(" ", 1)
            if len(cmd_body) < 2: return "❌ Формат: `/send <chat_id> <текст>`"
            try:
                res = await client.post(
                    f"{tg_base_url}/bot{tg_token}/sendMessage", 
                    json={"chat_id": cmd_body[0].strip(), "text": cmd_body[1].strip(), "parse_mode": "Markdown"}
                )
                return "✅ Отправлено." if res.status_code == 200 else f"❌ Ошибка: {res.text}"
            except Exception as e: return f"❌ Ошибка транспорта: {str(e)}"

        if text.startswith("/search "):
            if search_in_group is None: 
                return "❌ Системный вызов заблокирован: модуль api.user_client недоступен на хосте."
            cmd_body = text[8:].strip().split(" ", 1)
            if not cmd_body[0]: return "❌ Укажите параметры поиска."
            target_group = cmd_body[0].strip() if cmd_body[0].startswith(("@", "t.me")) else "@web3cyberservices_info"
            query = cmd_body[1].strip() if len(cmd_body) > 1 else text[8:].strip()
            try:
                return await search_in_group(target_group, query)
            except Exception as e: return f"❌ Ошибка поиска: {str(e)}"

        if text.startswith("/join "):
            if join_telegram_chat is None: 
                return "❌ Системный вызов заблокирован: модуль api.user_client недоступен на хосте."
            invite_link = text[6:].strip()
            if not invite_link: return "❌ Укажите инвайт-ссылку."
            try:
                return await join_telegram_chat(invite_link)
            except Exception as e: return f"❌ Ошибка выполнения системного вызова: {str(e)}"

        if text == "/reload_rules":
            return factory.hot_reload_rules() if factory else "⚠️ Фабрика не активна."

        # =====================================================================
        # БЛОК ИИ-ДИАЛОГА (ОПТИМИЗИРОВАННЫЙ АСИНХРОННЫЙ СТЕК)
        # =====================================================================
        api_keys = {
            "cerebras": os.environ.get("CEREBRAS_API_KEY"),
            "sambanova": os.environ.get("SAMBANOVA_API_KEY"),
            "mistral": os.environ.get("MISTRAL_API_KEY"),
            "gemini": os.environ.get("GEMINI_API_KEY"),
            "groq": os.environ.get("GROQ_API_KEY")
        }

        # Инжектируем контекст DOM-дерева и куки в промпт для расширения Chrome
        context_payload = ""
        if page_code:
            context_payload += f"\n[КОНТЕКСТ СТРАНИЦЫ АУДИТА]:\nРазмер DOM: {len(page_code)} символов.\n"
        if page_cookies:
            context_payload += f"[СЕССИОННЫЕ COOKIES]: {page_cookies}\n"

        base_instruction = (
            f"Ты — Cyber-Agent, ИИ-помощник контура Web3CyberServices. "
            f"Помогаешь в OSINT, пентесте и автоматизации. Сейчас общаешься с {user_name}. "
            f"Отвечай развернуто, технически грамотно и структурировано.{context_payload}"
        )
        
        try:
            if factory and hasattr(factory, "create_assistant_prompt"):
                system_prompt = await factory.create_assistant_prompt(base_instruction)
            else:
                system_prompt = base_instruction
        except Exception as e:
            logger.error(f"Ошибка генерации системного промпта через фабрику: {e}")
            system_prompt = base_instruction

        # Проверка типа системного промпта на случай скрытых асинхронных вызовов
        if asyncio.iscoroutine(system_prompt):
            logger.warning("Обнаружен объект coroutine в system_prompt! Принудительное разрешение.")
            system_prompt = await system_prompt

        messages = [{"role": "system", "content": str(system_prompt)}]

        # Если history пришла как объект сопрограммы из-за пропущенного await на верхнем уровне
        if asyncio.iscoroutine(history):
            logger.warning("Параметр history передан как coroutine! Разрешаем внутри функции.")
            try:
                history = await history
            except Exception as e:
                logger.error(f"Не удалось разрешить coroutine истории: {e}")
                history = ""

        # Безопасный парсинг истории диалога
        if isinstance(history, str) and history.strip():
            try:
                json_history = json.loads(history)
                if isinstance(json_history, list):
                    messages.extend(json_history)
            except json.JSONDecodeError:
                current_role = None
                current_content = []
                for line in history.split('\n'):
                    if line.startswith('User: '):
                        if current_role:
                            messages.append({"role": current_role, "content": '\n'.join(current_content).strip()})
                        current_role = 'user'
                        current_content = [line[6:]]
                    elif line.startswith('Agent: '):
                        if current_role:
                            messages.append({"role": current_role, "content": '\n'.join(current_content).strip()})
                        current_role = 'assistant'
                        current_content = [line[7:]]
                    else:
                        if current_role:
                            current_content.append(line)
                if current_role:
                    messages.append({"role": current_role, "content": '\n'.join(current_content).strip()})

        messages.append({"role": "user", "content": text})

        # --- КРИТИЧЕСКИЙ ФИЛЬТР ВАЛИДАЦИИ СТРУКТУРЫ ДАННЫХ ---
        # Проходим по всему стеку сообщений и принудительно очищаем от любых coroutine-объектов
        sanitized_messages = []
        for msg in messages:
            if not isinstance(msg, dict):
                continue
            role = msg.get("role", "user")
            content = msg.get("content", "")
            
            # Если контент оказался объектом сопрограммы или футуры
            if asyncio.iscoroutine(content) or hasattr(content, "__await__"):
                logger.error(f"Обнаружен coroutine-объект внутри контента сообщения для роли {role}!")
                try:
                    content = await content
                except Exception:
                    content = "[Ошибка: Несериализуемый объект контекста]"
            
            sanitized_messages.append({
                "role": str(role),
                "content": str(content)
            })

        # Формирование пула доступных провайдеров
        provider_pool = []
        if api_keys.get("groq"):
            provider_pool.append({
                "name": "Groq Llama 3.3 70B", "url": "https://api.groq.com/openai/v1/chat/completions",
                "headers": {"Authorization": f"Bearer {api_keys['groq']}"},
                "payload": {"model": "llama-3.3-70b-versatile", "messages": sanitized_messages, "max_tokens": 2048, "temperature": 0.5}
            })
        if api_keys.get("cerebras"):
            provider_pool.append({
                "name": "Cerebras Systems", "url": "https://api.cerebras.ai/v1/chat/completions",
                "headers": {"Authorization": f"Bearer {api_keys['cerebras']}"},
                "payload": {"model": "llama3.3-70b", "messages": sanitized_messages, "max_tokens": 2048, "temperature": 0.5}
            })
        if api_keys.get("sambanova"):
            provider_pool.append({
                "name": "SambaNova Cloud", "url": "https://api.sambanova.ai/v1/chat/completions",
                "headers": {"Authorization": f"Bearer {api_keys['sambanova']}"},
                "payload": {"model": "Meta-Llama-3.3-70B-Instruct", "messages": sanitized_messages, "max_tokens": 2048, "temperature": 0.5}
            })
        if api_keys.get("mistral"):
            provider_pool.append({
                "name": "Mistral AI", "url": "https://api.mistral.ai/v1/chat/completions",
                "headers": {"Authorization": f"Bearer {api_keys['mistral']}"},
                "payload": {"model": "mistral-large-latest", "messages": sanitized_messages, "max_tokens": 2048, "temperature": 0.5}
            })
        if api_keys.get("gemini"):
            provider_pool.append({
                "name": "Google Gemini 2.0 Flash", "url": "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
                "headers": {"Authorization": f"Bearer {api_keys['gemini']}"},
                "payload": {"model": "gemini-2.0-flash", "messages": sanitized_messages, "max_tokens": 2048, "temperature": 0.5}
            })

        if not provider_pool:
            return "❌ Критическая ошибка конфигурации: Ни один API-ключ ИИ-провайдеров не найден в переменных окружения."

        diagnostic_logs = []

        # Каскадный опрос ИИ-моделей
        for provider in provider_pool:
            try:
                headers = {**provider["headers"], "Content-Type": "application/json"}
                response = await client.post(provider["url"], headers=headers, json=provider["payload"])

                if response.status_code == 200:
                    data = response.json()
                    answer = data.get('choices', [{}])[0].get('message', {}).get('content')
                    if answer:
                        return answer
                    else:
                        diagnostic_logs.append(f"• `{provider['name']}` ➔ Ошибка: Пустой контент в ответе")
                        continue

                diagnostic_logs.append(f"• `{provider['name']}` ➔ HTTP {response.status_code}")
            except Exception as e:
                diagnostic_logs.append(f"• `{provider['name']}` ➔ Ошибка сети: {str(e)}")
                continue

        return f"❌ Все контуры инференса недоступны.\n\n**Отчет системы:**\n" + "\n".join(diagnostic_logs)