mib-new/config.py

import asyncio
import contextlib
import logging
import os
import re

import aiohttp
from aiogram import Bot
from aiogram.client.session.aiohttp import AiohttpSession
from dotenv import load_dotenv

try:
    from aiohttp_socks import ProxyConnector
except ImportError:  # pragma: no cover
    ProxyConnector = None

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("config")

BOT_TOKEN = os.getenv("BOT_TOKEN", "").strip()
YSK_API_KEY = os.getenv("YSK_API_KEY", "").strip()
API_WEATHER = os.getenv("API_WEATHER", "").strip()

ACCESS_KEY = os.getenv("ACCESS_KEY", "").strip()
SECRET_KEY = os.getenv("SECRET_KEY", "").strip()
BUCKET_NAME = os.getenv("BUCKET_NAME", "").strip()

DB_PATH = os.getenv("DB_PATH", "bot_database.db").strip() or "bot_database.db"

# Можно использовать старую переменную PROXY_URL или новую PROXY_URLS.
# Примеры:
# PROXY_URL=socks5://127.0.0.1:12553
# PROXY_URLS=socks5://127.0.0.1:12553,socks5://127.0.0.1:12554
_PROXY_RAW = (
    os.getenv("PROXY_URLS", "").strip()
    or os.getenv("PROXY_URL", "").strip()
)


class NetworkUnavailableError(RuntimeError):
    """Нет рабочего подключения ни через прокси, ни напрямую."""


def _parse_proxy_urls(raw_value: str) -> list[str]:
    """Разбирает список прокси из .env.

    Поддерживаются разделители:
    - запятая
    - точка с запятой
    - перенос строки

    Важно: AiohttpSession принимает только один proxy URL за раз,
    поэтому список нужно разбирать и проверять по очереди.
    """
    if not raw_value:
        return []

    items = []
    for item in re.split(r"[,;\r\n]+", raw_value):
        proxy = item.strip().strip("'\"")
        if proxy:
            items.append(proxy)

    # Убираем дубли, но сохраняем порядок.
    result = []
    seen = set()
    for proxy in items:
        if proxy not in seen:
            result.append(proxy)
            seen.add(proxy)

    return result


PROXY_URLS = _parse_proxy_urls(_PROXY_RAW)

# Оставляем PROXY_URL для обратной совместимости со старым кодом.
# Но новый код должен брать именно get_active_proxy_url().
PROXY_URL = PROXY_URLS[0] if PROXY_URLS else None

_ACTIVE_PROXY_URL: str | None = None
_CONNECTION_READY = False

if not BOT_TOKEN:
    raise RuntimeError("BOT_TOKEN is not set. Put it into .env")


def _mask_proxy(proxy_url: str | None) -> str:
    """Прячет логин/пароль в логах."""
    if not proxy_url:
        return "напрямую"

    return re.sub(
        r"(://)([^:/@\s]+):([^/@\s]+)@",
        r"\1***:***@",
        proxy_url,
    )


def _is_socks_proxy(proxy_url: str | None) -> bool:
    if not proxy_url:
        return False
    return proxy_url.lower().startswith(("socks4://", "socks5://", "socks5h://"))


def get_active_proxy_url() -> str | None:
    """Возвращает прокси, который реально прошёл проверку.

    Если подключение выбрано напрямую, вернёт None.
    Это нужно для других модулей, которые используют aiohttp.
    """
    return _ACTIVE_PROXY_URL


def ensure_connection_ready() -> None:
    """Защита от сетевых вызовов до setup_bot_connection()."""
    if not _CONNECTION_READY:
        raise NetworkUnavailableError(
            "Подключение к Telegram не настроено. Сначала вызовите setup_bot_connection()."
        )


def create_aiohttp_session(timeout: aiohttp.ClientTimeout | None = None) -> aiohttp.ClientSession:
    """Создаёт aiohttp-сессию с учётом активного прокси.

    Почему так:
    - aiogram.AiohttpSession умеет SOCKS-прокси сам;
    - обычный aiohttp не должен получать socks5:// через параметр proxy=;
    - для SOCKS в aiohttp нужен aiohttp_socks.ProxyConnector.
    """
    ensure_connection_ready()

    proxy_url = get_active_proxy_url()
    if not proxy_url:
        return aiohttp.ClientSession(timeout=timeout)

    if _is_socks_proxy(proxy_url):
        if ProxyConnector is None:
            raise RuntimeError(
                "Для SOCKS-прокси нужен пакет aiohttp-socks. Установите: pip install aiohttp-socks"
            )
        connector = ProxyConnector.from_url(proxy_url)
        return aiohttp.ClientSession(timeout=timeout, connector=connector)

    # Для HTTP/HTTPS-прокси aiohttp использует параметр proxy= в конкретном запросе.
    return aiohttp.ClientSession(timeout=timeout)


def get_aiohttp_request_kwargs() -> dict:
    """Аргументы для session.get/post у обычного aiohttp.

    SOCKS-прокси уже встроен в connector внутри create_aiohttp_session(),
    поэтому proxy= для socks5:// здесь не передаём.
    """
    proxy_url = get_active_proxy_url()
    if proxy_url and not _is_socks_proxy(proxy_url):
        return {"proxy": proxy_url}
    return {}


async def _close_session_safely(session: AiohttpSession | None) -> None:
    if session is None:
        return

    with contextlib.suppress(Exception):
        await session.close()


async def _build_checked_session(proxy_url: str | None) -> AiohttpSession:
    """Создаёт сессию и проверяет связь с Telegram методом get_me()."""
    session = AiohttpSession(proxy=proxy_url) if proxy_url else AiohttpSession()
    test_bot = Bot(token=BOT_TOKEN, session=session)

    try:
        await test_bot.get_me()
    except Exception:
        await _close_session_safely(session)
        raise

    return session


async def setup_bot_connection() -> bool:
    """Подбирает рабочее подключение для bot.

    Алгоритм:
    1. Проверить все прокси из .env по очереди.
    2. Если ни один не подошёл — попробовать прямое подключение.
    3. Если прямое подключение тоже не работает — вернуть False.

    Функция не бросает исключение наружу из-за нерабочей сети/прокси.
    main.py сможет аккуратно завершить работу без traceback.
    """
    global _ACTIVE_PROXY_URL, _CONNECTION_READY

    _CONNECTION_READY = False

    for proxy_url in PROXY_URLS:
        shown_proxy = _mask_proxy(proxy_url)
        logger.info("Пробую подключиться через прокси %s", shown_proxy)

        try:
            checked_session = await _build_checked_session(proxy_url)
        except Exception as exc:
            logger.warning(
                "Связь через прокси %s не удалась: %s. Пробую следующий вариант подключения.",
                shown_proxy,
                exc,
            )
            continue

        old_session = bot.session
        bot.session = checked_session
        await _close_session_safely(old_session)

        _ACTIVE_PROXY_URL = proxy_url
        _CONNECTION_READY = True
        logger.info("Подключение через прокси %s успешно.", shown_proxy)
        return True

    if PROXY_URLS:
        logger.warning("Ни один прокси не помог. Пробую подключиться напрямую.")
    else:
        logger.info("Прокси не указаны. Пробую подключиться напрямую.")

    try:
        checked_session = await _build_checked_session(None)
    except Exception as exc:
        _ACTIVE_PROXY_URL = None
        _CONNECTION_READY = False
        logger.error("Прямое подключение не удалось: %s", exc)
        logger.error("Нет возможности соединиться с Telegram. Завершаю работу.")
        return False

    old_session = bot.session
    bot.session = checked_session
    await _close_session_safely(old_session)

    _ACTIVE_PROXY_URL = None
    _CONNECTION_READY = True
    logger.info("Прямое подключение успешно.")
    return True


# Важно: не создаём AiohttpSession(proxy=PROXY_URL) на импорте.
# Если PROXY_URL кривой, программа должна не падать при импорте,
# а аккуратно проверить все варианты в setup_bot_connection().
bot = Bot(token=BOT_TOKEN, session=AiohttpSession())