786 lines
33 KiB
Python
786 lines
33 KiB
Python
"""
|
||
Command handling
|
||
"""
|
||
from pyrogram import Client
|
||
from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton
|
||
from pyrogram.filters import command
|
||
from pyrogram.handlers import MessageHandler
|
||
from bot.modules.access_control.permissions import require_permission, Permission
|
||
from bot.modules.access_control.user_manager import (
|
||
add_user, block_user, unblock_user,
|
||
add_admin, remove_admin
|
||
)
|
||
from bot.modules.message_handler.filters import is_url_message
|
||
from bot.utils.helpers import parse_user_id
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
async def get_start_keyboard(user_id: int) -> InlineKeyboardMarkup:
|
||
"""
|
||
Create keyboard for /start command
|
||
|
||
Args:
|
||
user_id: User ID
|
||
|
||
Returns:
|
||
InlineKeyboardMarkup with buttons
|
||
"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
|
||
# Base buttons for all users
|
||
buttons = [
|
||
[
|
||
InlineKeyboardButton("📥 Загрузить", callback_data="download"),
|
||
InlineKeyboardButton("📊 Статус", callback_data="status")
|
||
],
|
||
[
|
||
InlineKeyboardButton("❓ Помощь", callback_data="help")
|
||
]
|
||
]
|
||
|
||
# Additional buttons for administrators
|
||
if await is_admin(user_id):
|
||
buttons.append([
|
||
InlineKeyboardButton("📈 Статистика", callback_data="admin_stats")
|
||
])
|
||
|
||
return InlineKeyboardMarkup(buttons)
|
||
|
||
|
||
async def start_command(client: Client, message: Message):
|
||
"""Handle /start command"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
|
||
# Check authorization
|
||
if not await is_authorized(message.from_user.id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
welcome_text = (
|
||
"👋 **Привет! Я бот для загрузки медиа-файлов.**\n\n"
|
||
"📥 **Что я умею:**\n"
|
||
"• Загружать видео с YouTube, Instagram и других платформ\n"
|
||
"• Загружать файлы по прямым ссылкам\n"
|
||
"• Отправлять файлы вам в Telegram\n\n"
|
||
"**Как использовать:**\n"
|
||
"Просто отправьте мне ссылку на видео или файл, и я загружу его для вас!\n\n"
|
||
"Используйте кнопки ниже для управления:"
|
||
)
|
||
|
||
keyboard = await get_start_keyboard(message.from_user.id)
|
||
|
||
await message.reply(
|
||
welcome_text,
|
||
reply_markup=keyboard
|
||
)
|
||
|
||
|
||
async def help_command(client: Client, message: Message):
|
||
"""Handle /help command"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
|
||
# Check authorization
|
||
if not await is_authorized(message.from_user.id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
help_text = (
|
||
"👋 **Привет! Рад помочь!**\n\n"
|
||
|
||
"🎯 **Как начать работу:**\n"
|
||
"Это очень просто! Просто отправьте мне ссылку на видео или файл, и я сразу начну загрузку.\n\n"
|
||
|
||
"📥 **Что я умею загружать:**\n"
|
||
"• 🎬 Видео с YouTube, Instagram, TikTok и других платформ\n"
|
||
"• 📁 Файлы по прямым ссылкам\n"
|
||
"• 🎵 Аудио и музыку\n"
|
||
"• 📸 Изображения и фото\n\n"
|
||
|
||
"⌨️ **Основные команды:**\n"
|
||
"• `/start` - Открыть главное меню с кнопками\n"
|
||
"• `/help` - Показать эту справку\n"
|
||
"• `/status` - Посмотреть статус ваших загрузок\n"
|
||
"• `/cancel <task_id>` - Отменить задачу\n\n"
|
||
|
||
"💡 **Совет:** Используйте кнопки в главном меню для быстрого доступа к функциям!"
|
||
)
|
||
|
||
# Add information for administrators
|
||
from bot.modules.access_control.auth import is_admin
|
||
if await is_admin(message.from_user.id):
|
||
help_text += (
|
||
"\n\n"
|
||
"👑 **Команды для администраторов:**\n"
|
||
"• `/adduser <user_id или @username>` - Добавить нового пользователя\n"
|
||
"• `/blockuser <user_id или @username>` - Заблокировать пользователя\n"
|
||
"• `/unblockuser <user_id или @username>` - Разблокировать пользователя\n"
|
||
"• `/listusers` - Посмотреть список всех пользователей\n\n"
|
||
"💼 **Управление администраторами:**\n"
|
||
"• `/addadmin <user_id или @username>` - Назначить администратора\n"
|
||
"• `/removeadmin <user_id или @username>` - Снять права администратора\n"
|
||
"• `/listadmins` - Список всех администраторов"
|
||
)
|
||
await message.reply(help_text)
|
||
|
||
|
||
async def status_command(client: Client, message: Message):
|
||
"""Handle /status command"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
from bot.modules.task_scheduler.monitor import get_user_tasks_status
|
||
from bot.modules.task_scheduler.queue import TaskStatus
|
||
|
||
# Check authorization
|
||
if not await is_authorized(message.from_user.id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
tasks = await get_user_tasks_status(user_id)
|
||
|
||
if not tasks:
|
||
await message.reply("📊 У вас нет активных задач")
|
||
return
|
||
|
||
# Filter only active tasks (pending, processing)
|
||
active_tasks = [
|
||
t for t in tasks
|
||
if t.get('status') in ['pending', 'processing']
|
||
]
|
||
|
||
if not active_tasks:
|
||
await message.reply("📊 У вас нет активных задач")
|
||
return
|
||
|
||
status_text = "📊 **Ваши активные задачи:**\n\n"
|
||
|
||
for task in active_tasks[:10]: # Show maximum 10 tasks
|
||
task_id = task.get('id')
|
||
status = task.get('status', 'unknown')
|
||
progress = task.get('progress', 0)
|
||
url = task.get('url', 'N/A')
|
||
|
||
status_emoji = {
|
||
'pending': '⏳',
|
||
'processing': '🔄',
|
||
'completed': '✅',
|
||
'failed': '❌',
|
||
'cancelled': '🚫'
|
||
}.get(status, '❓')
|
||
|
||
status_text += (
|
||
f"{status_emoji} **Задача #{task_id}**\n"
|
||
f"🔗 {url[:50]}...\n"
|
||
f"📊 Прогресс: {progress}%\n"
|
||
f"📝 Статус: {status}\n\n"
|
||
)
|
||
|
||
if len(active_tasks) > 10:
|
||
status_text += f"... и еще {len(active_tasks) - 10} задач\n\n"
|
||
|
||
status_text += "💡 Используйте `/cancel <task_id>` для отмены задачи"
|
||
|
||
await message.reply(status_text)
|
||
|
||
|
||
async def cancel_command(client: Client, message: Message):
|
||
"""Handle /cancel command"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
from bot.modules.task_scheduler.monitor import cancel_user_task
|
||
from bot.modules.task_scheduler.queue import task_queue
|
||
|
||
# Check authorization
|
||
if not await is_authorized(message.from_user.id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /cancel <task_id>\n\nИспользуйте /status чтобы увидеть ID ваших задач")
|
||
return
|
||
|
||
try:
|
||
task_id = int(message.command[1])
|
||
except ValueError:
|
||
await message.reply("❌ Неверный формат task_id. Используйте число.")
|
||
return
|
||
|
||
user_id = message.from_user.id
|
||
|
||
# Cancel task
|
||
try:
|
||
success, message_text = await cancel_user_task(user_id, task_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in cancel_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка при отмене задачи: {str(e)}")
|
||
|
||
|
||
# User management commands (admin only)
|
||
async def adduser_command(client: Client, message: Message):
|
||
"""Add user"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from bot.utils.helpers import resolve_user_identifier
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /adduser <user_id или @username>")
|
||
return
|
||
|
||
identifier = message.command[1]
|
||
|
||
# Resolve identifier (user_id or username)
|
||
user_id, error_message = await resolve_user_identifier(identifier)
|
||
if not user_id:
|
||
await message.reply(f"❌ {error_message}")
|
||
return
|
||
|
||
try:
|
||
success, message_text = await add_user(user_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in adduser_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка: {str(e)}")
|
||
|
||
|
||
async def blockuser_command(client: Client, message: Message):
|
||
"""Block user"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from bot.utils.helpers import resolve_user_identifier
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /blockuser <user_id или @username>")
|
||
return
|
||
|
||
identifier = message.command[1]
|
||
|
||
# Resolve identifier (user_id or username)
|
||
user_id, error_message = await resolve_user_identifier(identifier)
|
||
if not user_id:
|
||
await message.reply(f"❌ {error_message}")
|
||
return
|
||
|
||
try:
|
||
success, message_text = await block_user(user_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in blockuser_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка: {str(e)}")
|
||
|
||
|
||
async def unblockuser_command(client: Client, message: Message):
|
||
"""Unblock user"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from bot.utils.helpers import resolve_user_identifier
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /unblockuser <user_id или @username>")
|
||
return
|
||
|
||
identifier = message.command[1]
|
||
|
||
# Resolve identifier (user_id or username)
|
||
user_id, error_message = await resolve_user_identifier(identifier)
|
||
if not user_id:
|
||
await message.reply(f"❌ {error_message}")
|
||
return
|
||
|
||
try:
|
||
success, message_text = await unblock_user(user_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in unblockuser_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка: {str(e)}")
|
||
|
||
|
||
async def listusers_command(client: Client, message: Message):
|
||
"""List users"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from shared.database.models import User
|
||
from shared.database.session import get_async_session_local
|
||
from sqlalchemy import select, func, desc
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
try:
|
||
async with get_async_session_local()() as session:
|
||
# Get total count
|
||
count_result = await session.execute(select(func.count(User.user_id)))
|
||
total_count = count_result.scalar() or 0
|
||
|
||
if total_count == 0:
|
||
await message.reply("📋 Пользователей в базе данных нет")
|
||
return
|
||
|
||
# Get users (limit to 50 for message length)
|
||
query = select(User).order_by(desc(User.created_at)).limit(50)
|
||
result = await session.execute(query)
|
||
users = list(result.scalars().all())
|
||
|
||
# Format message
|
||
text = f"📋 **Список пользователей** (всего: {total_count})\n\n"
|
||
|
||
for i, user in enumerate(users, 1):
|
||
username = f"@{user.username}" if user.username else "-"
|
||
name = f"{user.first_name or ''} {user.last_name or ''}".strip() or "-"
|
||
admin_badge = "👑" if user.is_admin else ""
|
||
blocked_badge = "🚫" if user.is_blocked else ""
|
||
|
||
text += (
|
||
f"{i}. {admin_badge} {blocked_badge} **ID:** `{user.user_id}`\n"
|
||
f" 👤 {username} ({name})\n"
|
||
f" 📅 Создан: {user.created_at.strftime('%Y-%m-%d %H:%M') if user.created_at else 'N/A'}\n\n"
|
||
)
|
||
|
||
if total_count > 50:
|
||
text += f"\n... и еще {total_count - 50} пользователей (показаны первые 50)"
|
||
|
||
# Split message if too long (Telegram limit is 4096 characters)
|
||
if len(text) > 4000:
|
||
# Send first part
|
||
await message.reply(text[:4000])
|
||
# Send remaining users count
|
||
await message.reply(f"... и еще {total_count - 50} пользователей")
|
||
else:
|
||
await message.reply(text)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error listing users: {e}", exc_info=True)
|
||
await message.reply("❌ Ошибка при получении списка пользователей")
|
||
|
||
|
||
async def addadmin_command(client: Client, message: Message):
|
||
"""Assign administrator"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from bot.utils.helpers import resolve_user_identifier
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /addadmin <user_id или @username>")
|
||
return
|
||
|
||
identifier = message.command[1]
|
||
requester_id = message.from_user.id
|
||
|
||
# Resolve identifier (user_id or username)
|
||
user_id, error_message = await resolve_user_identifier(identifier)
|
||
if not user_id:
|
||
await message.reply(f"❌ {error_message}")
|
||
return
|
||
|
||
try:
|
||
success, message_text = await add_admin(user_id, requester_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in addadmin_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка: {str(e)}")
|
||
|
||
|
||
async def removeadmin_command(client: Client, message: Message):
|
||
"""Remove administrator privileges"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from bot.utils.helpers import resolve_user_identifier
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
if not message.command or len(message.command) < 2:
|
||
await message.reply("❌ Использование: /removeadmin <user_id или @username>")
|
||
return
|
||
|
||
identifier = message.command[1]
|
||
requester_id = message.from_user.id
|
||
|
||
# Resolve identifier (user_id or username)
|
||
user_id, error_message = await resolve_user_identifier(identifier)
|
||
if not user_id:
|
||
await message.reply(f"❌ {error_message}")
|
||
return
|
||
|
||
try:
|
||
success, message_text = await remove_admin(user_id, requester_id)
|
||
if success:
|
||
await message.reply(f"✅ {message_text}")
|
||
else:
|
||
await message.reply(f"❌ {message_text}")
|
||
except Exception as e:
|
||
logger.error(f"Error in removeadmin_command: {e}", exc_info=True)
|
||
await message.reply(f"❌ Произошла ошибка: {str(e)}")
|
||
|
||
|
||
async def listadmins_command(client: Client, message: Message):
|
||
"""List administrators"""
|
||
from bot.modules.access_control.auth import is_admin
|
||
from shared.database.models import User
|
||
from shared.database.session import get_async_session_local
|
||
from sqlalchemy import select, func, desc
|
||
|
||
# Check access permissions
|
||
if not await is_admin(message.from_user.id):
|
||
await message.reply("❌ Эта команда доступна только администраторам")
|
||
return
|
||
|
||
try:
|
||
async with get_async_session_local()() as session:
|
||
# Get administrators
|
||
query = select(User).where(User.is_admin == True).order_by(desc(User.created_at))
|
||
result = await session.execute(query)
|
||
admins = list(result.scalars().all())
|
||
|
||
if not admins:
|
||
await message.reply("👑 Администраторов в базе данных нет")
|
||
return
|
||
|
||
# Format message
|
||
text = f"👑 **Список администраторов** (всего: {len(admins)})\n\n"
|
||
|
||
for i, admin in enumerate(admins, 1):
|
||
username = f"@{admin.username}" if admin.username else "-"
|
||
name = f"{admin.first_name or ''} {admin.last_name or ''}".strip() or "-"
|
||
blocked_badge = "🚫" if admin.is_blocked else ""
|
||
|
||
text += (
|
||
f"{i}. {blocked_badge} **ID:** `{admin.user_id}`\n"
|
||
f" 👤 {username} ({name})\n"
|
||
f" 📅 Создан: {admin.created_at.strftime('%Y-%m-%d %H:%M') if admin.created_at else 'N/A'}\n\n"
|
||
)
|
||
|
||
await message.reply(text)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error listing administrators: {e}", exc_info=True)
|
||
await message.reply("❌ Ошибка при получении списка администраторов")
|
||
|
||
|
||
async def login_command(client: Client, message: Message):
|
||
"""Handle /login command to get OTP code"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
from bot.modules.database.session import AsyncSessionLocal
|
||
from web.utils.otp import create_otp_code
|
||
from shared.config import settings
|
||
|
||
user_id = message.from_user.id
|
||
|
||
# Check authorization
|
||
if not await is_authorized(user_id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
try:
|
||
# Create OTP code
|
||
async with AsyncSessionLocal() as session:
|
||
code = await create_otp_code(user_id, session)
|
||
|
||
if code:
|
||
# Form URL for web interface
|
||
if settings.WEB_HOST == "0.0.0.0":
|
||
login_url = f"localhost:{settings.WEB_PORT}"
|
||
else:
|
||
login_url = f"{settings.WEB_HOST}:{settings.WEB_PORT}"
|
||
|
||
await message.reply(
|
||
f"🔐 **Ваш код для входа в веб-интерфейс:**\n\n"
|
||
f"**`{code}`**\n\n"
|
||
f"⏰ Код действителен 10 минут\n\n"
|
||
f"🌐 Перейдите на http://{login_url}/admin/login и введите этот код\n\n"
|
||
f"💡 Или используйте ваш User ID: `{user_id}`"
|
||
)
|
||
else:
|
||
await message.reply("❌ Не удалось создать код. Попробуйте позже.")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при создании OTP кода: {e}")
|
||
await message.reply("❌ Произошла ошибка при создании кода. Попробуйте позже.")
|
||
|
||
|
||
async def url_handler(client: Client, message: Message):
|
||
"""Handle URL messages"""
|
||
from bot.modules.access_control.auth import is_authorized
|
||
from bot.modules.task_scheduler.queue import task_queue, Task, TaskStatus
|
||
from bot.modules.task_scheduler.executor import task_executor, set_app_client
|
||
from bot.modules.message_handler.filters import is_youtube_url, is_instagram_url
|
||
from bot.modules.media_loader.ytdlp import get_videos_list
|
||
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
|
||
import time
|
||
|
||
# Check authorization
|
||
if not await is_authorized(message.from_user.id):
|
||
await message.reply("❌ У вас нет доступа к этому боту")
|
||
return
|
||
|
||
url = message.text.strip()
|
||
user_id = message.from_user.id
|
||
|
||
# URL validation
|
||
from bot.utils.helpers import is_valid_url
|
||
if not is_valid_url(url):
|
||
await message.reply(
|
||
"❌ Некорректный или небезопасный URL.\n\n"
|
||
"Пожалуйста, отправьте валидную ссылку (http:// или https://)"
|
||
)
|
||
return
|
||
|
||
# Check if URL is YouTube or Instagram (direct download)
|
||
is_youtube = is_youtube_url(url)
|
||
is_instagram = is_instagram_url(url)
|
||
|
||
# For non-YouTube/Instagram URLs, check if there are multiple videos
|
||
if not is_youtube and not is_instagram:
|
||
# Check if URL contains video selection callback data
|
||
if url.startswith('video_select:'):
|
||
# This is a callback from video selection - extract actual URL
|
||
actual_url = url.replace('video_select:', '', 1)
|
||
url = actual_url
|
||
else:
|
||
# Try to get list of videos from webpage
|
||
try:
|
||
status_msg = await message.reply("🔍 Анализирую страницу...")
|
||
from shared.config import settings
|
||
videos_info = await get_videos_list(url, cookies_file=settings.COOKIES_FILE)
|
||
|
||
if videos_info and videos_info.get('videos'):
|
||
videos = videos_info['videos']
|
||
|
||
if len(videos) > 1:
|
||
# Multiple videos found - show selection menu
|
||
await status_msg.delete()
|
||
|
||
playlist_title = videos_info.get('playlist_title', 'Найдено видео')
|
||
text = f"📹 **{playlist_title}**\n\n"
|
||
text += f"Найдено видео: **{len(videos)}**\n\n"
|
||
text += "Выберите видео для загрузки:\n\n"
|
||
|
||
# Create inline keyboard with video selection buttons
|
||
from bot.modules.message_handler.video_selection_cache import store_video_selection
|
||
|
||
buttons = []
|
||
for idx, video in enumerate(videos[:10], 1): # Limit to 10 videos
|
||
title = video.get('title', f'Видео {idx}')[:50] # Limit title length
|
||
duration = video.get('duration')
|
||
if duration:
|
||
from bot.utils.helpers import format_duration
|
||
duration_str = format_duration(duration)
|
||
title += f" ({duration_str})"
|
||
|
||
# Store video URL in cache and use short identifier in callback_data
|
||
# This avoids Telegram's 64-byte limit on callback_data
|
||
selection_id = store_video_selection(video['url'], user_id)
|
||
callback_data = f"video_select:{selection_id}"
|
||
buttons.append([InlineKeyboardButton(f"{idx}. {title}", callback_data=callback_data)])
|
||
|
||
keyboard = InlineKeyboardMarkup(buttons)
|
||
await message.reply(text, reply_markup=keyboard)
|
||
return
|
||
elif len(videos) == 1:
|
||
# Single video found - use its URL
|
||
url = videos[0]['url']
|
||
await status_msg.delete()
|
||
else:
|
||
# No videos found, but continue with original URL (might be direct video link)
|
||
await status_msg.delete()
|
||
except Exception as e:
|
||
logger.error(f"Error getting videos list: {e}", exc_info=True)
|
||
# Continue with original URL if error occurs
|
||
try:
|
||
await status_msg.delete()
|
||
except:
|
||
pass
|
||
|
||
# Check concurrent tasks count
|
||
from bot.config import settings
|
||
active_tasks_count = await task_queue.get_user_active_tasks_count(user_id)
|
||
if active_tasks_count >= settings.MAX_CONCURRENT_TASKS:
|
||
await message.reply(
|
||
f"❌ Превышен лимит одновременных задач ({settings.MAX_CONCURRENT_TASKS}).\n"
|
||
f"⏳ Дождитесь завершения текущих задач или отмените их через /cancel"
|
||
)
|
||
return
|
||
|
||
# Set client for task executor
|
||
set_app_client(client)
|
||
|
||
# Generate unique task_id using UUID
|
||
from bot.utils.helpers import generate_unique_task_id
|
||
task_id = generate_unique_task_id()
|
||
|
||
# Check that such ID doesn't exist yet (in case of collision, though probability is extremely low)
|
||
existing_task = await task_queue.get_task_by_id(task_id)
|
||
max_retries = 10
|
||
retries = 0
|
||
while existing_task and retries < max_retries:
|
||
task_id = generate_unique_task_id()
|
||
existing_task = await task_queue.get_task_by_id(task_id)
|
||
retries += 1
|
||
|
||
if existing_task:
|
||
# If after 10 attempts still collision (extremely unlikely), log error
|
||
logger.error(f"Failed to generate unique task_id after {max_retries} attempts")
|
||
await message.reply("❌ Ошибка при создании задачи. Попробуйте позже.")
|
||
return
|
||
|
||
# Duplicate URL check will be performed atomically in task_queue.add_task()
|
||
url_normalized = url.strip()
|
||
|
||
task = Task(
|
||
id=task_id,
|
||
user_id=user_id,
|
||
task_type="download",
|
||
url=url_normalized,
|
||
status=TaskStatus.PENDING
|
||
)
|
||
|
||
# Save task to database BEFORE adding to queue (race condition fix)
|
||
try:
|
||
from shared.database.models import Task as DBTask
|
||
from shared.database.session import get_async_session_local
|
||
from shared.database.user_helpers import ensure_user_exists
|
||
from datetime import datetime
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
async with get_async_session_local()() as session:
|
||
# Ensure user exists before creating task
|
||
await ensure_user_exists(user_id, session)
|
||
|
||
db_task = DBTask(
|
||
id=task_id,
|
||
user_id=user_id,
|
||
task_type=task.task_type,
|
||
status=task.status.value,
|
||
url=task.url,
|
||
progress=0,
|
||
created_at=datetime.utcnow(),
|
||
updated_at=datetime.utcnow()
|
||
)
|
||
session.add(db_task)
|
||
await session.commit()
|
||
logger.info(f"Task {task_id} saved to database from bot")
|
||
except IntegrityError as e:
|
||
logger.error(f"IntegrityError saving task {task_id} to database (possibly duplicate ID): {e}", exc_info=True)
|
||
# Generate new task_id and retry
|
||
from bot.utils.helpers import generate_unique_task_id
|
||
task_id = generate_unique_task_id()
|
||
task.id = task_id
|
||
try:
|
||
async with get_async_session_local()() as session:
|
||
from shared.database.user_helpers import ensure_user_exists
|
||
# Ensure user exists before creating task
|
||
await ensure_user_exists(user_id, session)
|
||
|
||
db_task = DBTask(
|
||
id=task_id,
|
||
user_id=user_id,
|
||
task_type=task.task_type,
|
||
status=task.status.value,
|
||
url=task.url,
|
||
progress=0,
|
||
created_at=datetime.utcnow(),
|
||
updated_at=datetime.utcnow()
|
||
)
|
||
session.add(db_task)
|
||
await session.commit()
|
||
logger.info(f"Task {task_id} saved to database from bot with new ID")
|
||
except Exception as e2:
|
||
logger.error(f"Error saving task {task_id} to database again: {e2}", exc_info=True)
|
||
await message.reply("❌ Ошибка при создании задачи. Попробуйте позже.")
|
||
return
|
||
except Exception as e:
|
||
logger.error(f"Error saving task {task_id} to database: {e}", exc_info=True)
|
||
await message.reply("❌ Ошибка при создании задачи. Попробуйте позже.")
|
||
return
|
||
|
||
# Add to queue (with duplicate URL check) AFTER saving to database
|
||
success = await task_queue.add_task(task, check_duplicate_url=True)
|
||
if not success:
|
||
# If failed to add to queue, remove from database
|
||
try:
|
||
async with get_async_session_local()() as session:
|
||
db_task = await session.get(DBTask, task_id)
|
||
if db_task:
|
||
await session.delete(db_task)
|
||
await session.commit()
|
||
except Exception as e:
|
||
logger.error(f"Error deleting task {task_id} from database after failed queue addition: {e}")
|
||
await message.reply(
|
||
f"⚠️ Задача с этим URL уже обрабатывается.\n"
|
||
f"Дождитесь завершения или отмените предыдущую задачу через /cancel"
|
||
)
|
||
return
|
||
|
||
# Start executor if not already started
|
||
if not task_executor._running:
|
||
await task_executor.start()
|
||
|
||
# Send initial status message and save message_id for updates
|
||
from bot.modules.task_scheduler.executor import set_task_message
|
||
status_message = await message.reply(
|
||
f"📥 **Загрузка начата**\n\n"
|
||
f"🔗 {url}\n\n"
|
||
f"📊 Прогресс: **0%**\n"
|
||
f"⏳ Ожидание начала загрузки..."
|
||
)
|
||
# Save message_id for task updates
|
||
set_task_message(task_id, status_message.id)
|
||
|
||
|
||
def register_commands(app: Client):
|
||
"""Register all commands"""
|
||
# Base commands (for all users)
|
||
app.add_handler(MessageHandler(start_command, filters=command("start")))
|
||
app.add_handler(MessageHandler(help_command, filters=command("help")))
|
||
app.add_handler(MessageHandler(status_command, filters=command("status")))
|
||
app.add_handler(MessageHandler(cancel_command, filters=command("cancel")))
|
||
app.add_handler(MessageHandler(login_command, filters=command("login")))
|
||
|
||
# User management commands (admin only)
|
||
app.add_handler(MessageHandler(adduser_command, filters=command("adduser")))
|
||
app.add_handler(MessageHandler(blockuser_command, filters=command("blockuser")))
|
||
app.add_handler(MessageHandler(unblockuser_command, filters=command("unblockuser")))
|
||
app.add_handler(MessageHandler(listusers_command, filters=command("listusers")))
|
||
|
||
# Administrator management commands (admin only)
|
||
app.add_handler(MessageHandler(addadmin_command, filters=command("addadmin")))
|
||
app.add_handler(MessageHandler(removeadmin_command, filters=command("removeadmin")))
|
||
app.add_handler(MessageHandler(listadmins_command, filters=command("listadmins")))
|
||
|
||
# URL message handling
|
||
app.add_handler(MessageHandler(url_handler, filters=is_url_message))
|
||
|
||
logger.info("Commands registered")
|