Add source
This commit is contained in:
747
bot/modules/message_handler/commands.py
Normal file
747
bot/modules/message_handler/commands.py
Normal file
@@ -0,0 +1,747 @@
|
||||
"""
|
||||
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, remove_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"
|
||||
"• `/removeuser <user_id или @username>` - Удалить пользователя\n"
|
||||
"• `/blockuser <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 removeuser_command(client: Client, message: Message):
|
||||
"""Remove 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("❌ Использование: /removeuser <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 remove_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 removeuser_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
|
||||
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 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()
|
||||
|
||||
await message.reply(
|
||||
f"✅ Ссылка получена!\n\n"
|
||||
f"🔗 {url}\n\n"
|
||||
f"📥 Загрузка добавлена в очередь. Я начну загрузку в ближайшее время.\n"
|
||||
f"⏳ Вы получите уведомление о статусе загрузки."
|
||||
)
|
||||
|
||||
|
||||
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(removeuser_command, filters=command("removeuser")))
|
||||
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")
|
||||
|
||||
Reference in New Issue
Block a user