qbittorent_bot/bot.py

import logging
import os
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    filters,
    CallbackQueryHandler,
    ContextTypes,
)
from qbittorrent import Client, APIError

# Настройка логирования
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger(__name__)

# --- Конфигурация (будет загружаться из переменных окружения Docker) ---
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
QBT_HOST = os.getenv("QBT_HOST")
QBT_PORT = os.getenv("QBT_PORT", "8080")
QBT_USERNAME = os.getenv("QBT_USERNAME")
QBT_PASSWORD = os.getenv("QBT_PASSWORD")

# Определение доступных директорий (предполагаем, что они известны)
# В реальной системе можно было бы получить список из qBittorrent API
# или настроить их через переменные окружения.
DOWNLOAD_DIRECTORIES = {
    "Фильмы": "/share/Data/Felms",
    "Сериалы": "/share/Data/Serials",
    "Музыка": "/share/Music",
    "Другое": "/share/Data/torrents",
}

# --- Инициализация qBittorrent клиента ---
qb = None

def init_qbittorrent_client():
    global qb
    if not all([QBT_HOST, QBT_USERNAME, QBT_PASSWORD]):
        logger.error("QBittorrent credentials are not fully set in environment variables.")
        return False
    try:
        qb = Client(f"http://{QBT_HOST}:{QBT_PORT}/")
        qb.login(QBT_USERNAME, QBT_PASSWORD)
        logger.info(f"Successfully connected to qBittorrent at {QBT_HOST}:{QBT_PORT}")
        return True
    except APIError as e:
        logger.error(f"Failed to connect or login to qBittorrent: {e}")
        qb = None
        return False
    except Exception as e:
        logger.error(f"An unexpected error occurred during qBittorrent connection: {e}")
        qb = None
        return False

# --- Обработчики команд ---

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not init_qbittorrent_client():
        await update.message.reply_text(
            "Не удалось подключиться к qBittorrent. Проверьте переменные окружения и доступность сервера."
        )
        return
    await update.message.reply_text(
        "Привет! Я бот для управления qBittorrent.\n"
        "Отправь мне magnet-ссылку или URL torrent-файла, "
        "чтобы добавить загрузку.\n"
        "Используй /status для просмотра текущих загрузок."
    )

async def status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not qb:
        if not init_qbittorrent_client():
            await update.message.reply_text(
                "Не удалось подключиться к qBittorrent. Попробуйте еще раз или проверьте настройки."
            )
            return

    try:
        torrents = qb.torrents()
        if not torrents:
            await update.message.reply_text("Нет активных загрузок.")
            return

        message = "Текущие загрузки:\n\n"
        for t in torrents:
            progress = f"{t.progress:.2%}"
            size = f"{t.size / (1024*1024*1024):.2f} GB" if t.size else "N/A"
            download_speed = f"{t.dlspeed / (1024*1024):.2f} MB/s"
            upload_speed = f"{t.upspeed / (1024*1024):.2f} MB/s"

            message += (
                f"📝 Имя: {t.name}\n"
                f"📊 Прогресс: {progress}\n"
                f"📦 Размер: {size}\n"
                f"⬇️ Скорость загрузки: {download_speed}\n"
                f"⬆️ Скорость отдачи: {upload_speed}\n"
                f"🚦 Статус: {t.state}\n"
                f"--- \n"
            )
        await update.message.reply_text(message)

    except APIError as e:
        logger.error(f"Error fetching torrents: {e}")
        await update.message.reply_text(f"Ошибка при получении списка загрузок: {e}")
    except Exception as e:
        logger.error(f"An unexpected error occurred: {e}")
        await update.message.reply_text(f"Произошла непредвиденная ошибка: {e}")


async def handle_url(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    if not qb:
        if not init_qbittorrent_client():
            await update.message.reply_text(
                "Не удалось подключиться к qBittorrent. Попробуйте еще раз или проверьте настройки."
            )
            return

    url = update.message.text
    context.user_data['download_url'] = url

    keyboard = []
    for name, path in DOWNLOAD_DIRECTORIES.items():
        keyboard.append([InlineKeyboardButton(name, callback_data=f"dir_{path}")])
    reply_markup = InlineKeyboardMarkup(keyboard)

    await update.message.reply_text(
        f"Вы хотите загрузить: `{url}`\n"
        "Выберите директорию для загрузки:",
        reply_markup=reply_markup,
        parse_mode='Markdown'
    )

async def button_callback(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()

    if query.data.startswith("dir_"):
        selected_directory = query.data.replace("dir_", "")
        download_url = context.user_data.get('download_url')

        if not download_url:
            await query.edit_message_text("Ошибка: URL для загрузки не найден. Пожалуйста, отправьте ссылку снова.")
            return

        try:
            qb.download_from_link(download_url, save_path=selected_directory)
            await query.edit_message_text(
                f"Загрузка '{download_url}' добавлена в '{selected_directory}'."
            )
            logger.info(f"Added download '{download_url}' to '{selected_directory}'")
        except APIError as e:
            logger.error(f"Error adding download: {e}")
            await query.edit_message_text(f"Ошибка при добавлении загрузки: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred while adding download: {e}")
            await query.edit_message_text(f"Произошла непредвиденная ошибка при добавлении загрузки: {e}")


def main() -> None:
    application = Application.builder().token(TELEGRAM_BOT_TOKEN).build()

    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("status", status))
    application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_url))
    application.add_handler(CallbackQueryHandler(button_callback))

    logger.info("Bot started polling...")
    application.run_polling(allowed_updates=Update.ALL_TYPES)

if __name__ == "__main__":
    if not TELEGRAM_BOT_TOKEN:
        logger.error("TELEGRAM_BOT_TOKEN is not set. Please set the environment variable.")
        exit(1)
    if not all([QBT_HOST, QBT_USERNAME, QBT_PASSWORD]):
        logger.warning("QBittorrent connection details (QBT_HOST, QBT_USERNAME, QBT_PASSWORD) are not fully set. Bot will attempt to connect on first use of qBittorrent related commands.")
    main()