338 lines
13 KiB
Python
338 lines
13 KiB
Python
"""Main bot module."""
|
|
|
|
import logging
|
|
|
|
from telebot.async_telebot import AsyncTeleBot
|
|
from telebot import types, asyncio_helper
|
|
|
|
from config import TelegramConfig
|
|
from handlers import (
|
|
BotState,
|
|
create_main_keyboard,
|
|
create_qbittorrent_keyboard,
|
|
create_jellyfin_keyboard,
|
|
create_category_keyboard,
|
|
)
|
|
from services import QbittorrentService, JellyfinService
|
|
from utils import escape_markdown
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ServerManagerBot:
|
|
"""Main bot class that orchestrates all services and handlers."""
|
|
|
|
def __init__(
|
|
self,
|
|
telegram_config: TelegramConfig,
|
|
qbittorrent_service: QbittorrentService,
|
|
jellyfin_service: JellyfinService,
|
|
):
|
|
"""Initialize the bot."""
|
|
self.telegram_config = telegram_config
|
|
self.qbittorrent = qbittorrent_service
|
|
self.jellyfin = jellyfin_service
|
|
if self.telegram_config.proxy_server is not None:
|
|
asyncio_helper.proxy = self.telegram_config.proxy_server
|
|
self.bot = AsyncTeleBot(
|
|
telegram_config.bot_token
|
|
)
|
|
self.state = BotState()
|
|
self._setup_handlers()
|
|
|
|
def _setup_handlers(self) -> None:
|
|
"""Register all bot handlers."""
|
|
self.bot.message_handler(commands=["help", "start"])(self.send_welcome)
|
|
self.bot.message_handler(commands=["chatid"])(self.handle_chat_id)
|
|
self.bot.message_handler(func=lambda message: message.text == "Go back")(self.go_back)
|
|
self.bot.message_handler(func=lambda message: message.text == "qBittorrent")(self.handle_qbittorrent)
|
|
self.bot.message_handler(func=lambda message: message.text == "Jellyfin")(self.handle_jellyfin)
|
|
self.bot.message_handler(func=lambda message: message.text == "Add torrent")(self.handle_add_torrent)
|
|
self.bot.message_handler(func=lambda message: message.text == "Delete torrent")(self.handle_delete_torrent_menu)
|
|
self.bot.message_handler(func=lambda message: message.text == "Get list of active torrents")(self.handle_active_torrents)
|
|
self.bot.message_handler(func=lambda message: message.text == "Get total stats")(self.handle_stats)
|
|
self.bot.message_handler(func=lambda message: message.text == "Refresh Jellyfin")(self.handle_jellyfin_refresh)
|
|
self.bot.message_handler(func=lambda message: message.text == "Upload video")(self.handle_upload_video)
|
|
self.bot.message_handler(content_types=["document"])(self.handle_document)
|
|
self.bot.message_handler(content_types=["video"])(self.handle_video)
|
|
self.bot.message_handler(content_types=["text"])(self.handle_text_message)
|
|
|
|
async def send_welcome(self, message: types.Message) -> None:
|
|
"""Send welcome message with main menu."""
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Hello! Choose an option below:",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
async def handle_chat_id(self, message: types.Message) -> None:
|
|
"""Handle chat ID request for debugging."""
|
|
chat_id = message.chat.id
|
|
thread_id = message.message_thread_id
|
|
|
|
if message.chat.is_forum and thread_id:
|
|
logger.info(f"Group: {chat_id} | Topic: {thread_id}")
|
|
else:
|
|
logger.info(f"Group: {chat_id} | (Not a forum or no topic)")
|
|
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Captured chat_id and thread_id, check developer console",
|
|
)
|
|
|
|
async def go_back(self, message: types.Message) -> None:
|
|
"""Return to main menu."""
|
|
await self.send_welcome(message)
|
|
|
|
async def handle_qbittorrent(self, message: types.Message) -> None:
|
|
"""Show qBittorrent menu."""
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Select action.",
|
|
reply_markup=create_qbittorrent_keyboard(),
|
|
)
|
|
|
|
async def handle_jellyfin(self, message: types.Message) -> None:
|
|
"""Show Jellyfin menu."""
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Select action.",
|
|
reply_markup=create_jellyfin_keyboard(),
|
|
)
|
|
|
|
async def handle_add_torrent(self, message: types.Message) -> None:
|
|
"""Start torrent addition flow."""
|
|
self.state.set_state(message.chat.id, States.AWAITING_FOR_TORRENT)
|
|
await self.bot.reply_to(message, "Please send a magnet link or a torrent file")
|
|
|
|
async def handle_active_torrents(self, message: types.Message) -> None:
|
|
"""Show active downloading torrents."""
|
|
torrents = self.qbittorrent.get_downloading_torrents()
|
|
if not torrents:
|
|
msg = "No torrents are being downloaded at the moment."
|
|
else:
|
|
msg = "\n".join(
|
|
f"- {torrent.name}, ETA: {torrent.eta}s"
|
|
for torrent in torrents
|
|
)
|
|
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
msg,
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
async def handle_stats(self, message: types.Message) -> None:
|
|
"""Send qBittorrent statistics."""
|
|
try:
|
|
stats = self.qbittorrent.get_stats()
|
|
stats_message = (
|
|
"📊 <b>qBittorrent Metrics</b>\n\n"
|
|
"🔄 <b>Current Session</b>\n"
|
|
f"⬇️ Downloaded: {stats['session_download']}\n"
|
|
f"⬆️ Uploaded: {stats['session_upload']}\n\n"
|
|
"⚡ <b>Live Speeds</b>\n"
|
|
f"⏬ Down: {stats['current_download']}\n"
|
|
f"⏫ Up: {stats['current_upload']}\n\n"
|
|
"🧩 <b>Torrent Summary</b>\n"
|
|
f"✅ Total: {stats['total_torrents']}\n"
|
|
f"⏬ Active DL: {stats['downloading_count']}\n"
|
|
f"⏸️ Paused: {stats['paused_count']}"
|
|
)
|
|
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
stats_message,
|
|
parse_mode="HTML",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Error fetching stats: {e}")
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"❌ <b>Error fetching metrics:</b>",
|
|
parse_mode="HTML",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
async def handle_delete_torrent_menu(self, message: types.Message) -> None:
|
|
"""Show completed torrents for deletion."""
|
|
torrents = self.qbittorrent.get_completed_torrents()
|
|
if not torrents:
|
|
msg = "No completed torrents available."
|
|
else:
|
|
msg = "\n".join(
|
|
f"**{i}**. {escape_markdown(torrent.name)}"
|
|
for i, torrent in enumerate(torrents, start=1)
|
|
)
|
|
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
msg,
|
|
parse_mode="MarkdownV2",
|
|
)
|
|
self.state.set_state(message.chat.id, States.AWAITING_FOR_TORRENT_NUMBER)
|
|
|
|
async def handle_jellyfin_refresh(self, message: types.Message) -> None:
|
|
"""Trigger Jellyfin library refresh."""
|
|
try:
|
|
result = await self.jellyfin.refresh_library()
|
|
if result:
|
|
await self.bot.send_message(message.chat.id, result)
|
|
else:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Library refresh started. Please wait a few minutes and then restart your client.",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Jellyfin refresh failed: {e}")
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
f"Failed to refresh library: {e}",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
async def handle_upload_video(self, message: types.Message) -> None:
|
|
"""Start video upload flow."""
|
|
self.state.set_state(message.chat.id, States.AWAITING_FOR_VIDEO)
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Please, send video file to upload to Jellyfin.",
|
|
)
|
|
|
|
async def _download_and_save_video(self, message: types.Message, file_id: str, file_name: str) -> None:
|
|
"""Download and save a video file."""
|
|
await self.bot.send_message(message.chat.id, "Downloading...")
|
|
file_info = await self.bot.get_file(file_id)
|
|
downloaded_file = await self.bot.download_file(file_info.file_path)
|
|
|
|
try:
|
|
await self.jellyfin.upload_video(file_id, file_name, lambda fid: self.bot.download_file(file_info.file_path))
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"File downloaded. Triggering Jellyfin library update, please wait a few minutes...",
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Video upload failed: {e}")
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
f"Failed to upload video: {e}",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
finally:
|
|
self.state.clear_state(message.chat.id)
|
|
|
|
async def handle_document(self, message: types.Message) -> None:
|
|
"""Handle document (video file) messages."""
|
|
if self.state.get_state(message.chat.id) != States.AWAITING_FOR_VIDEO:
|
|
return
|
|
|
|
file_id = message.document.file_id
|
|
file_name = message.document.file_name
|
|
await self._download_and_save_video(message, file_id, file_name)
|
|
|
|
async def handle_video(self, message: types.Message) -> None:
|
|
"""Handle video messages."""
|
|
if self.state.get_state(message.chat.id) != States.AWAITING_FOR_VIDEO:
|
|
return
|
|
|
|
file_id = message.video.file_id
|
|
file_name = message.video.file_name
|
|
await self._download_and_save_video(message, file_id, file_name)
|
|
|
|
async def handle_text_message(self, message: types.Message) -> None:
|
|
"""Handle text messages based on current state."""
|
|
state = self.state.get_state(message.chat.id)
|
|
|
|
if state == States.AWAITING_FOR_TORRENT:
|
|
self.state.set_awaiting_torrent(message.chat.id, message.text.strip())
|
|
self.state.set_state(message.chat.id, States.AWAITING_FOR_CATEGORY)
|
|
categories = self.qbittorrent.get_categories()
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Select torrent category.",
|
|
reply_markup=create_category_keyboard(categories),
|
|
)
|
|
|
|
elif state == States.AWAITING_FOR_CATEGORY:
|
|
category = message.text.strip()
|
|
categories = self.qbittorrent.get_categories()
|
|
|
|
if category not in categories:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Select torrent category.",
|
|
reply_markup=create_category_keyboard(categories),
|
|
)
|
|
return
|
|
|
|
magnet = self.state.get_awaiting_torrent(message.chat.id)
|
|
if not magnet:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Error: No torrent found. Please start over.",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
return
|
|
|
|
try:
|
|
result = self.qbittorrent.add_torrent(magnet, category)
|
|
if result == "Ok.":
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Torrent added successfully! Track download progress in qBitController",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
else:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
f"Failed to add torrent. Error: '{result}'",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Failed to add torrent: {e}")
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Torrent already has been added",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
elif state == States.AWAITING_FOR_TORRENT_NUMBER:
|
|
try:
|
|
number = int(message.text)
|
|
except ValueError:
|
|
await self.bot.send_message(message.chat.id, "Please, send a number")
|
|
return
|
|
|
|
torrents = self.qbittorrent.get_completed_torrents()
|
|
if number < 1 or number > len(torrents):
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Invalid number. Please try again.",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
return
|
|
|
|
torrent_hash = torrents[number - 1].hash.lower()
|
|
success = self.qbittorrent.delete_torrent(torrent_hash)
|
|
|
|
if success:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Torrent deleted successfully.",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
else:
|
|
await self.bot.send_message(
|
|
message.chat.id,
|
|
"Failed to delete the torrent!",
|
|
reply_markup=create_main_keyboard(),
|
|
)
|
|
|
|
self.state.clear_state(message.chat.id)
|
|
|
|
async def run(self) -> None:
|
|
"""Start the bot."""
|
|
await self.bot.infinity_polling()
|