Fixed token logic, refactored byte streamer, client selection and workload handling

- **New Features**
- Added explicit token activation tracking and validation for improved access control.
- Enhanced error messages for invalid or expired activation links and token requirements.

- **Refactor**
- Simplified and improved performance of media streaming and file handling.
- Streamlined token generation and management with clearer activation logic.
- Optimised client selection and workload handling for streaming requests.
- Refined command handling with improved error logging and message parsing.

- **Bug Fixes**
- Improved error handling and user feedback for invalid commands and payloads.

- **Style**
- Updated user-facing messages for clarity and consistency.

Thunder/bot/plugins/common.py

@@ -1,431 +1,426 @@
-# Thunder/bot/plugins/common.py
-
-import time
-import asyncio
-import uuid
-from typing import Tuple, Optional
-from urllib.parse import quote_plus
-from pyrogram import filters, StopPropagation
-from pyrogram.client import Client
-from pyrogram.errors import RPCError, MessageNotModified
-from pyrogram.types import (
-InlineKeyboardButton,
-InlineKeyboardMarkup,
-Message,
-User,
-CallbackQuery,
-LinkPreviewOptions
-)
-from Thunder.bot import StreamBot
-from Thunder.vars import Var
-from Thunder.utils.database import db
-from Thunder.utils.messages import *
-from Thunder.utils.human_readable import humanbytes
-from Thunder.utils.file_properties import get_media_file_size, get_name
-from Thunder.utils.force_channel import force_channel_check
-from Thunder.utils.logger import logger
-from Thunder.utils.tokens import get, check, validate_activation_token
-from Thunder.utils.decorators import check_banned
-from Thunder.utils.bot_utils import (
-notify_channel,
-log_new_user,
-generate_media_links,
-handle_user_error,
-generate_dc_text,
-get_user_safely
-)
-
-def has_media(message):
-    return any(
-        hasattr(message, attr) and getattr(message, attr)
-        for attr in ["document", "photo", "video", "audio", "voice",
-                    "sticker", "animation", "video_note"]
-    )
-
-@check_banned
-@StreamBot.on_message(filters.command("start") & filters.private)
-async def start_command(bot: Client, message: Message):
-    try:
-        user_id_str = str(message.from_user.id) if message.from_user else 'unknown'
-        if message.from_user:
-            await log_new_user(bot, message.from_user.id, message.from_user.first_name)
-        
-        if len(message.command) == 2:
-            token = message.command[1]
-            try:
-                if token.startswith("token"):
-                    validation_result = await validate_activation_token(token)
-                    if validation_result["valid"]:
-                        await db.update_user_token(message.from_user.id, token)
-                        await message.reply_text(
-                            MSG_TOKEN_ACTIVATED.format(
-                                expiry_date=validation_result['expiry_date'],
-                                description=validation_result['description']
-                            ),
-                            link_preview_options=LinkPreviewOptions(is_disabled=True),
-                            quote=True
-                        )
-                        return
-                    else:
-                        await message.reply_text(
-                            MSG_TOKEN_FAILED.format(
-                                reason=validation_result['reason'],
-                                error_id=uuid.uuid4().hex[:8]
-                            ),
-                            link_preview_options=LinkPreviewOptions(is_disabled=True),
-                            quote=True
-                        )
-                        return
-                else:
-                    record = await get(token)
-                    if record and record["user_id"] == message.from_user.id and await check(message.from_user.id):
-                        await message.reply_text(MSG_TOKEN_VERIFIED, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
-                        return
-                    else:
-                        await message.reply_text(MSG_TOKEN_INVALID, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
-                        return
-            except Exception as e:
-                logger.error(f"Token activation error: {e}")
-                await message.reply_text(MSG_TOKEN_ERROR.format(error_id=uuid.uuid4().hex[:8]), quote=True)
-                return
-
-        parts = message.text.strip().split(maxsplit=1)
-        if len(parts) == 1 or (len(parts) > 1 and parts[1].lower() == "start"):
-            welcome_text = MSG_WELCOME.format(user_name=message.from_user.first_name)
-            
-            if Var.FORCE_CHANNEL_ID:
-                error_id_context = uuid.uuid4().hex[:8]
-                try:
-                    chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
-                    invite_link = getattr(chat, 'invite_link', None)
-                    chat_username = getattr(chat, 'username', None)
-                    chat_title = getattr(chat, 'title', 'Channel')
-                    
-                    if not invite_link and chat_username:
-                        invite_link = f"https://t.me/{chat_username}"
-                    
-                    if invite_link:
-                        welcome_text += f"\n\n{MSG_COMMUNITY_CHANNEL.format(channel_title=chat_title)}"
-                    else:
-                        logger.warning(f"(ID: {error_id_context}) Could not retrieve invite link for FORCE_CHANNEL_ID {Var.FORCE_CHANNEL_ID} for /start message text (Channel: {chat_title}, User: {user_id_str}).")
-                        welcome_text += f"\n\n{MSG_COMMUNITY_CHANNEL.format(channel_title=chat_title)}"
-                except asyncio.TimeoutError as e_timeout:
-                    logger.error(f"(ID: {error_id_context}) Timeout fetching Force Channel details for /start text (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_timeout}")
-                    welcome_text += f"\n\n{MSG_DEFAULT_WELCOME}"
-                except RPCError as e_rpc:
-                    logger.error(f"(ID: {error_id_context}) RPCError fetching Force Channel details for /start text (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_rpc}")
-                    welcome_text += f"\n\n{MSG_DEFAULT_WELCOME}"
-                except Exception as e:
-                    logger.error(f"(ID: {error_id_context}) Error adding force channel link to /start message text (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e}")
-                    welcome_text += f"\n\n{MSG_DEFAULT_WELCOME}"
-
-            reply_markup = InlineKeyboardMarkup([
-                [
-                    InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command"),
-                    InlineKeyboardButton(MSG_BUTTON_ABOUT, callback_data="about_command"),
-                ],
-                [
-                    InlineKeyboardButton(MSG_BUTTON_GITHUB, url="https://github.com/fyaz05/FileToLink/"),
-                    InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")
-                ]
-            ])
-
-            if Var.FORCE_CHANNEL_ID:
-                error_id_context_button = uuid.uuid4().hex[:8]
-                try:
-                    chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
-                    invite_link = getattr(chat, 'invite_link', None)
-                    chat_username = getattr(chat, 'username', None)
-                    chat_title = getattr(chat, 'title', 'Channel')
-                    
-                    if not invite_link and chat_username:
-                        invite_link = f"https://t.me/{chat_username}"
-                    
-                    if invite_link:
-                        reply_markup.inline_keyboard.append([InlineKeyboardButton(MSG_BUTTON_JOIN_CHANNEL.format(channel_title=chat_title), url=invite_link)])
-                    else:
-                        logger.warning(f"(ID: {error_id_context_button}) Could not retrieve invite link for FORCE_CHANNEL_ID {Var.FORCE_CHANNEL_ID} for /start message button (Channel: {chat_title}, User: {user_id_str}).")
-                except asyncio.TimeoutError as e_timeout_btn:
-                    logger.error(f"(ID: {error_id_context_button}) Timeout fetching Force Channel details for /start button (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_timeout_btn}")
-                except RPCError as e_rpc_btn:
-                    logger.error(f"(ID: {error_id_context_button}) RPCError fetching Force Channel details for /start button (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_rpc_btn}")
-                except Exception as e_btn:
-                    logger.error(f"(ID: {error_id_context_button}) Error adding force channel button to /start message (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_btn}")
-
-            await message.reply_text(text=welcome_text, reply_markup=reply_markup, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
-            return
-
-        payload = parts[1]
-        try:
-            msg_id = int(payload)
-            retrieved_messages = await bot.get_messages(chat_id=Var.BIN_CHANNEL, message_ids=msg_id)
-            if not retrieved_messages:
-                await handle_user_error(message, MSG_ERROR_FILE_INVALID)
-                return
-            
-            file_msg = retrieved_messages[0] if isinstance(retrieved_messages, list) else retrieved_messages
-            if not file_msg:
-                await handle_user_error(message, MSG_ERROR_FILE_INVALID)
-                return
-
-            stream_link, download_link, file_name, file_size = await generate_media_links(file_msg)
-            reply_text = MSG_LINKS.format(
-                file_name=file_name,
-                file_size=file_size,
-                download_link=download_link,
-                stream_link=stream_link
-            )
-
-            await message.reply_text(
-                text=reply_text,
-                reply_markup=InlineKeyboardMarkup([
-                    [
-                        InlineKeyboardButton(MSG_BUTTON_STREAM_NOW, url=stream_link),
-                        InlineKeyboardButton(MSG_BUTTON_DOWNLOAD, url=download_link)
-                    ]
-                ]),
-                quote=True,
-                link_preview_options=LinkPreviewOptions(is_disabled=True)
-            )
-        except ValueError:
-            await handle_user_error(message, MSG_ERROR_FILE_INVALID_ID)
-        except Exception as e:
-            logger.error(f"Error processing /start with payload: {e}")
-            await handle_user_error(message, MSG_FILE_ACCESS_ERROR)
-    except Exception as e:
-        logger.error(f"Error in start_command: {e}")
-        await handle_user_error(message, MSG_ERROR_GENERIC)
-
-@check_banned
-@StreamBot.on_message(filters.command("help") & filters.private)
-async def help_command(bot: Client, message: Message):
-    try:
-        user_id_str = str(message.from_user.id) if message.from_user else 'unknown'
-        if message.from_user:
-            await log_new_user(bot, message.from_user.id, message.from_user.first_name)
-
-        help_text = MSG_HELP
-        buttons = [
-            [InlineKeyboardButton(MSG_BUTTON_ABOUT, callback_data="about_command")]
-        ]
-
-        if Var.FORCE_CHANNEL_ID:
-            error_id_context_help = uuid.uuid4().hex[:8]
-            try:
-                chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
-                invite_link = getattr(chat, 'invite_link', None)
-                chat_username = getattr(chat, 'username', None)
-                chat_title = getattr(chat, 'title', 'Channel')
-                
-                if not invite_link and chat_username:
-                    invite_link = f"https://t.me/{chat_username}"
-                
-                if invite_link:
-                    buttons.append([
-                        InlineKeyboardButton(MSG_BUTTON_JOIN_CHANNEL.format(channel_title=chat_title), url=invite_link)
-                    ])
-                else:
-                    logger.warning(f"(ID: {error_id_context_help}) Could not retrieve or construct invite link for FORCE_CHANNEL_ID {Var.FORCE_CHANNEL_ID} for /help command (User: {user_id_str}).")
-            except Exception as e_help:
-                logger.error(f"(ID: {error_id_context_help}) Error adding force channel button to /help command (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_help}")
-
-        buttons.append([InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")])
-        reply_markup = InlineKeyboardMarkup(buttons)
-
-        await message.reply_text(
-            text=help_text,
-            reply_markup=reply_markup,
-            quote=True,
-            link_preview_options=LinkPreviewOptions(is_disabled=True)
-        )
-    except Exception as e:
-        logger.error(f"Error in help_command: {e}")
-        await handle_user_error(message, MSG_ERROR_GENERIC)
-
-@check_banned
-@StreamBot.on_message(filters.command("about") & filters.private)
-async def about_command(bot: Client, message: Message):
-    try:
-        if message.from_user:
-            await log_new_user(bot, message.from_user.id, message.from_user.first_name)
-
-        about_text = MSG_ABOUT
-        buttons = [
-            [InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command")],
-            [InlineKeyboardButton(MSG_BUTTON_GITHUB, url="https://github.com/fyaz05/FileToLink"),
-             InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")]
-        ]
-        reply_markup = InlineKeyboardMarkup(buttons)
-
-        await message.reply_text(
-            text=about_text,
-            reply_markup=reply_markup,
-            quote=True,
-            link_preview_options=LinkPreviewOptions(is_disabled=True)
-        )
-    except Exception as e:
-        logger.error(f"Error in about_command: {e}")
-        await handle_user_error(message, MSG_ERROR_GENERIC)
-
-@check_banned
-@force_channel_check
-@StreamBot.on_message(filters.command("dc"))
-async def dc_command(bot: Client, message: Message):
-    try:
-        if not message.from_user and not message.reply_to_message:
-            await handle_user_error(message, MSG_DC_ANON_ERROR)
-            return
-
-        async def process_dc_info(user: User):
-            dc_text = await generate_dc_text(user)
-            buttons = []
-            
-            if user.username:
-                profile_url = f"https://t.me/{user.username}"
-            else:
-                profile_url = f"tg://user?id={user.id}"
-            
-            buttons.append([
-                InlineKeyboardButton(MSG_BUTTON_VIEW_PROFILE, url=profile_url)
-            ])
-            buttons.append([
-                InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")
-            ])
-            
-            dc_keyboard = InlineKeyboardMarkup(buttons)
-            await message.reply_text(
-                dc_text,
-                reply_markup=dc_keyboard,
-                quote=True,
-                link_preview_options=LinkPreviewOptions(is_disabled=True)
-            )
-
-        async def process_file_dc_info(file_msg: Message):
-            try:
-                file_name = get_name(file_msg) or "Untitled File"
-                file_size = humanbytes(get_media_file_size(file_msg))
-                
-                file_type_map = {
-                    "document": MSG_FILE_TYPE_DOCUMENT,
-                    "photo": MSG_FILE_TYPE_PHOTO,
-                    "video": MSG_FILE_TYPE_VIDEO,
-                    "audio": MSG_FILE_TYPE_AUDIO,
-                    "voice": MSG_FILE_TYPE_VOICE,
-                    "sticker": MSG_FILE_TYPE_STICKER,
-                    "animation": MSG_FILE_TYPE_ANIMATION,
-                    "video_note": MSG_FILE_TYPE_VIDEO_NOTE
-                }
-                
-                file_type_attr = next((attr for attr in file_type_map if getattr(file_msg, attr, None)), "unknown")
-                file_type_display = file_type_map.get(file_type_attr, MSG_FILE_TYPE_UNKNOWN)
-                
-                actual_media = None
-                if file_type_attr == "photo" and file_msg.photo and file_msg.photo.sizes:
-                    actual_media = file_msg.photo.sizes[-1]
-                elif file_type_attr != "unknown":
-                    potential = getattr(file_msg, file_type_attr, None)
-                    if potential:
-                        actual_media = potential
-                
-                dc_id = MSG_DC_UNKNOWN
-                if hasattr(file_msg, 'raw') and hasattr(file_msg.raw, 'media'):
-                    if hasattr(file_msg.raw.media, 'document') and hasattr(file_msg.raw.media.document, 'dc_id'):
-                        dc_id = file_msg.raw.media.document.dc_id
-                
-                dc_text = MSG_DC_FILE_INFO.format(
-                    file_name=file_name,
-                    file_size=file_size,
-                    file_type=file_type_display,
-                    dc_id=dc_id
-                )
-                
-                await message.reply_text(
-                    dc_text,
-                    quote=True,
-                    link_preview_options=LinkPreviewOptions(is_disabled=True)
-                )
-            except Exception as e:
-                logger.error(f"Error processing file info for DC command: {e}")
-                await handle_user_error(message, MSG_DC_FILE_ERROR)
-
-        args = message.text.strip().split(maxsplit=1)
-        if len(args) > 1:
-            query = args[1].strip()
-            user = await get_user_safely(bot, query)
-            if user:
-                await process_dc_info(user)
-            else:
-                await handle_user_error(message, MSG_ERROR_USER_INFO)
-            return
-
-        if message.reply_to_message:
-            if has_media(message.reply_to_message):
-                await process_file_dc_info(message.reply_to_message)
-                return
-            elif message.reply_to_message.from_user:
-                await process_dc_info(message.reply_to_message.from_user)
-                return
-            else:
-                await handle_user_error(message, MSG_DC_INVALID_USAGE)
-                return
-
-        if message.from_user:
-            await process_dc_info(message.from_user)
-        else:
-            await handle_user_error(message, MSG_DC_ANON_ERROR)
-    except Exception as e:
-        logger.error(f"Error in dc_command: {e}")
-        await handle_user_error(message, MSG_ERROR_GENERIC)
-
-@check_banned
-@force_channel_check
-@StreamBot.on_message(filters.command("ping") & filters.private)
-async def ping_command(bot: Client, message: Message):
-    try:
-        start_time = time.time()
-        sent_msg = await message.reply_text(MSG_PING_START, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
-        end_time = time.time()
-        time_taken_ms = (end_time - start_time) * 1000
-        
-        buttons = [
-            [InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command"),
-             InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")]
-        ]
-        
-        await sent_msg.edit_text(
-            MSG_PING_RESPONSE.format(time_taken_ms=time_taken_ms),
-            reply_markup=InlineKeyboardMarkup(buttons),
-            link_preview_options=LinkPreviewOptions(is_disabled=True)
-        )
-    except Exception as e:
-        logger.error(f"Error in ping_command: {e}")
-        await handle_user_error(message, MSG_ERROR_GENERIC)
-
-async def handle_callback_error(callback_query, error, operation="callback"):
-    error_id = uuid.uuid4().hex[:8]
-    logger.error(f"Error in {operation}: {error}")
-    try:
-        await callback_query.answer(
-            MSG_ERROR_GENERIC_CALLBACK.format(error_id=error_id),
-            show_alert=True
-        )
-    except Exception as e:
-        logger.error(f"Failed to send error callback: {e}")
-
-@StreamBot.on_callback_query(filters.regex(r"^close_panel$"))
-async def close_panel_callback(client: Client, callback_query: CallbackQuery):
-    try:
-        await callback_query.answer()
-        await callback_query.message.delete()
-        if callback_query.message.reply_to_message:
-            try:
-                ctx = callback_query.message.reply_to_message
-                await ctx.delete()
-                if ctx.reply_to_message:
-                    await ctx.reply_to_message.delete()
-            except Exception as e:
-                logger.warning(f"Error deleting command messages: {e}")
-    except Exception as e:
-        await handle_callback_error(callback_query, e, "close_panel_callback")
-    finally:
-        raise StopPropagation
+# Thunder/bot/plugins/common.py
+
+import time
+import asyncio
+import uuid
+from datetime import datetime, timedelta
+from typing import Tuple, Optional
+from urllib.parse import quote_plus
+from pyrogram import filters, StopPropagation
+from pyrogram.client import Client
+from pyrogram.errors import RPCError, MessageNotModified
+from pyrogram.types import (
+    InlineKeyboardButton,
+    InlineKeyboardMarkup,
+    Message,
+    User,
+    CallbackQuery,
+    LinkPreviewOptions
+)
+from Thunder.bot import StreamBot
+from Thunder.vars import Var
+from Thunder.utils.database import db
+from Thunder.utils.messages import *
+from Thunder.utils.human_readable import humanbytes
+from Thunder.utils.file_properties import get_media_file_size, get_name
+from Thunder.utils.force_channel import force_channel_check
+from Thunder.utils.logger import logger
+from Thunder.utils.tokens import check
+from Thunder.utils.decorators import check_banned
+from Thunder.utils.bot_utils import (
+    notify_channel,
+    log_new_user,
+    generate_media_links,
+    handle_user_error,
+    generate_dc_text,
+    get_user_safely
+)
+
+def has_media(message):
+    return any(
+        hasattr(message, attr) and getattr(message, attr)
+        for attr in ["document", "photo", "video", "audio", "voice",
+                    "sticker", "animation", "video_note"]
+    )
+
+@check_banned
+@StreamBot.on_message(filters.command("start") & filters.private)
+async def start_command(bot: Client, message: Message):
+    user_id = message.from_user.id if message.from_user else None
+    user_id_str = str(user_id) if user_id else 'unknown'
+
+    try:
+        if message.from_user:
+            await log_new_user(bot, user_id, message.from_user.first_name)
+
+        if len(message.command) == 2:
+            payload = message.command[1]
+
+            token_doc = await db.token_col.find_one({"token": payload})
+
+            if token_doc:
+                if token_doc["user_id"] == user_id:
+                    if not token_doc.get("activated", False):
+                        token_expires_at = token_doc.get("expires_at")
+                        if isinstance(token_expires_at, datetime) and token_expires_at > datetime.utcnow():
+                            await db.token_col.update_one(
+                                {"token": payload, "user_id": user_id},
+                                {"$set": {"activated": True}}
+                            )
+                            duration_hours = getattr(Var, "TOKEN_TTL_HOURS", 24)
+                            await message.reply_text(
+                                MSG_TOKEN_ACTIVATED.format(duration_hours=duration_hours),
+                                link_preview_options=LinkPreviewOptions(is_disabled=True),
+                                quote=True
+                            )
+                        else:
+                            await message.reply_text(
+                                MSG_TOKEN_FAILED.format(reason="This activation link has expired.", error_id=uuid.uuid4().hex[:8]),
+                                link_preview_options=LinkPreviewOptions(is_disabled=True),
+                                quote=True
+                            )
+                    else:
+                        await message.reply_text(
+                            MSG_TOKEN_FAILED.format(reason="Token has already been activated.", error_id=uuid.uuid4().hex[:8]),
+                            link_preview_options=LinkPreviewOptions(is_disabled=True),
+                            quote=True
+                        )
+                else:
+                    logger.warning(f"User {user_id} attempted to activate token {payload} belonging to user {token_doc.get('user_id')}")
+                    await message.reply_text(
+                        MSG_TOKEN_FAILED.format(reason="This activation link is not for your account.", error_id=uuid.uuid4().hex[:8]),
+                        link_preview_options=LinkPreviewOptions(is_disabled=True),
+                        quote=True
+                    )
+                return
+
+            try:
+                msg_id = int(payload)
+                retrieved_messages = await bot.get_messages(chat_id=Var.BIN_CHANNEL, message_ids=msg_id)
+                if not retrieved_messages:
+                    await handle_user_error(message, MSG_ERROR_FILE_INVALID)
+                    return
+
+                file_msg = retrieved_messages[0] if isinstance(retrieved_messages, list) else retrieved_messages
+                if not file_msg:
+                    await handle_user_error(message, MSG_ERROR_FILE_INVALID)
+                    return
+
+                stream_link, download_link, file_name, file_size = await generate_media_links(file_msg)
+                reply_text = MSG_LINKS.format(
+                    file_name=file_name,
+                    file_size=file_size,
+                    download_link=download_link,
+                    stream_link=stream_link
+                )
+
+                await message.reply_text(
+                    text=reply_text,
+                    reply_markup=InlineKeyboardMarkup([
+                        [
+                            InlineKeyboardButton(MSG_BUTTON_STREAM_NOW, url=stream_link),
+                            InlineKeyboardButton(MSG_BUTTON_DOWNLOAD, url=download_link)
+                        ]
+                    ]),
+                    quote=True,
+                    link_preview_options=LinkPreviewOptions(is_disabled=True)
+                )
+                return
+            except ValueError:
+                logger.warning(f"Invalid /start payload from user {user_id_str}: {payload}")
+                await message.reply_text(
+                    MSG_START_INVALID_PAYLOAD.format(error_id=uuid.uuid4().hex[:8]),
+                    quote=True,
+                    link_preview_options=LinkPreviewOptions(is_disabled=True)
+                )
+                return
+            except Exception as e:
+                logger.error(f"Error processing /start payload '{payload}' for user {user_id_str}: {e}")
+                await handle_user_error(message, MSG_FILE_ACCESS_ERROR)
+                return
+
+        parts = message.text.strip().split(maxsplit=1)
+        if len(message.command) == 1 or (len(parts) > 1 and parts[1].lower() == "start"):
+            welcome_text = MSG_WELCOME.format(user_name=message.from_user.first_name if message.from_user else "Guest")
+
+            if Var.FORCE_CHANNEL_ID:
+                error_id_context = uuid.uuid4().hex[:8]
+                try:
+                    chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
+                    invite_link = getattr(chat, 'invite_link', None)
+                    chat_username = getattr(chat, 'username', None)
+                    chat_title = getattr(chat, 'title', 'Channel')
+
+                    if not invite_link and chat_username:
+                        invite_link = f"https://t.me/{chat_username}"
+
+                    if invite_link:
+                        welcome_text += f"\n\n{MSG_COMMUNITY_CHANNEL.format(channel_title=chat_title)}"
+                    else:
+                        logger.warning(f"(ID: {error_id_context}) Could not retrieve invite link for FORCE_CHANNEL_ID {Var.FORCE_CHANNEL_ID} for /start message text (Channel: {chat_title}, User: {user_id_str}).")
+                except Exception as e:
+                    logger.error(f"(ID: {error_id_context}) Error adding force channel link to /start message text (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e}")
+
+            reply_markup_buttons = [
+                [
+                    InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command"),
+                    InlineKeyboardButton(MSG_BUTTON_ABOUT, callback_data="about_command"),
+                ],
+                [
+                    InlineKeyboardButton(MSG_BUTTON_GITHUB, url="https://github.com/fyaz05/FileToLink/"),
+                    InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")
+                ]
+            ]
+
+            if Var.FORCE_CHANNEL_ID:
+                try:
+                    chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
+                    invite_link = getattr(chat, 'invite_link', None)
+                    chat_username = getattr(chat, 'username', None)
+                    if not invite_link and chat_username: invite_link = f"https://t.me/{chat_username}"
+                    if invite_link:
+                        reply_markup_buttons.append([InlineKeyboardButton(MSG_BUTTON_JOIN_CHANNEL.format(channel_title=getattr(chat, 'title', 'Channel')), url=invite_link)])
+                except Exception as e:
+                    logger.error(f"Error adding force channel button to /start message (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e}")
+
+            reply_markup = InlineKeyboardMarkup(reply_markup_buttons)
+            await message.reply_text(text=welcome_text, reply_markup=reply_markup, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
+            return
+
+    except Exception as e:
+        logger.error(f"Error in start_command for user {user_id_str}: {e}")
+        await handle_user_error(message, MSG_ERROR_GENERIC)
+
+@check_banned
+@StreamBot.on_message(filters.command("help") & filters.private)
+async def help_command(bot: Client, message: Message):
+    try:
+        user_id_str = str(message.from_user.id) if message.from_user else 'unknown'
+        if message.from_user:
+            await log_new_user(bot, message.from_user.id, message.from_user.first_name)
+
+        help_text = MSG_HELP
+        buttons = [
+            [InlineKeyboardButton(MSG_BUTTON_ABOUT, callback_data="about_command")]
+        ]
+
+        if Var.FORCE_CHANNEL_ID:
+            error_id_context_help = uuid.uuid4().hex[:8]
+            try:
+                chat = await bot.get_chat(Var.FORCE_CHANNEL_ID)
+                invite_link = getattr(chat, 'invite_link', None)
+                chat_username = getattr(chat, 'username', None)
+                chat_title = getattr(chat, 'title', 'Channel')
+
+                if not invite_link and chat_username:
+                    invite_link = f"https://t.me/{chat_username}"
+
+                if invite_link:
+                    buttons.append([
+                        InlineKeyboardButton(MSG_BUTTON_JOIN_CHANNEL.format(channel_title=chat_title), url=invite_link)
+                    ])
+                else:
+                    logger.warning(f"(ID: {error_id_context_help}) Could not retrieve or construct invite link for FORCE_CHANNEL_ID {Var.FORCE_CHANNEL_ID} for /help command (User: {user_id_str}).")
+            except Exception as e_help:
+                logger.error(f"(ID: {error_id_context_help}) Error adding force channel button to /help command (User: {user_id_str}, FChannel: {Var.FORCE_CHANNEL_ID}): {e_help}")
+
+        buttons.append([InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")])
+        reply_markup = InlineKeyboardMarkup(buttons)
+
+        await message.reply_text(
+            text=help_text,
+            reply_markup=reply_markup,
+            quote=True,
+            link_preview_options=LinkPreviewOptions(is_disabled=True)
+        )
+    except Exception as e:
+        logger.error(f"Error in help_command: {e}")
+        await handle_user_error(message, MSG_ERROR_GENERIC)
+
+@check_banned
+@StreamBot.on_message(filters.command("about") & filters.private)
+async def about_command(bot: Client, message: Message):
+    try:
+        if message.from_user:
+            await log_new_user(bot, message.from_user.id, message.from_user.first_name)
+
+        about_text = MSG_ABOUT
+        buttons = [
+            [InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command")],
+            [InlineKeyboardButton(MSG_BUTTON_GITHUB, url="https://github.com/fyaz05/FileToLink/"),
+             InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")]
+        ]
+        reply_markup = InlineKeyboardMarkup(buttons)
+
+        await message.reply_text(
+            text=about_text,
+            reply_markup=reply_markup,
+            quote=True,
+            link_preview_options=LinkPreviewOptions(is_disabled=True)
+        )
+    except Exception as e:
+        logger.error(f"Error in about_command: {e}")
+        await handle_user_error(message, MSG_ERROR_GENERIC)
+
+@check_banned
+@force_channel_check
+@StreamBot.on_message(filters.command("dc"))
+async def dc_command(bot: Client, message: Message):
+    try:
+        if not message.from_user and not message.reply_to_message:
+            await handle_user_error(message, MSG_DC_ANON_ERROR)
+            return
+
+        async def process_dc_info(user: User):
+            dc_text = await generate_dc_text(user)
+            buttons = []
+
+            if user.username:
+                profile_url = f"https://t.me/{user.username}"
+            else:
+                profile_url = f"tg://user?id={user.id}"
+
+            buttons.append([
+                InlineKeyboardButton(MSG_BUTTON_VIEW_PROFILE, url=profile_url)
+            ])
+            buttons.append([
+                InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")
+            ])
+
+            dc_keyboard = InlineKeyboardMarkup(buttons)
+            await message.reply_text(
+                dc_text,
+                reply_markup=dc_keyboard,
+                quote=True,
+                link_preview_options=LinkPreviewOptions(is_disabled=True)
+            )
+
+        async def process_file_dc_info(file_msg: Message):
+            try:
+                file_name = get_name(file_msg) or "Untitled File"
+                file_size = humanbytes(get_media_file_size(file_msg))
+
+                file_type_map = {
+                    "document": MSG_FILE_TYPE_DOCUMENT,
+                    "photo": MSG_FILE_TYPE_PHOTO,
+                    "video": MSG_FILE_TYPE_VIDEO,
+                    "audio": MSG_FILE_TYPE_AUDIO,
+                    "voice": MSG_FILE_TYPE_VOICE,
+                    "sticker": MSG_FILE_TYPE_STICKER,
+                    "animation": MSG_FILE_TYPE_ANIMATION,
+                    "video_note": MSG_FILE_TYPE_VIDEO_NOTE
+                }
+
+                file_type_attr = next((attr for attr in file_type_map if getattr(file_msg, attr, None)), "unknown")
+                file_type_display = file_type_map.get(file_type_attr, MSG_FILE_TYPE_UNKNOWN)
+
+                actual_media = None
+                if file_type_attr == "photo" and file_msg.photo and file_msg.photo.sizes:
+                    actual_media = file_msg.photo.sizes[-1]
+                elif file_type_attr != "unknown":
+                    potential = getattr(file_msg, file_type_attr, None)
+                    if potential:
+                        actual_media = potential
+
+                dc_id = MSG_DC_UNKNOWN
+                if hasattr(file_msg, 'raw') and hasattr(file_msg.raw, 'media'):
+                    if hasattr(file_msg.raw.media, 'document') and hasattr(file_msg.raw.media.document, 'dc_id'):
+                        dc_id = file_msg.raw.media.document.dc_id
+
+                dc_text = MSG_DC_FILE_INFO.format(
+                    file_name=file_name,
+                    file_size=file_size,
+                    file_type=file_type_display,
+                    dc_id=dc_id
+                )
+
+                await message.reply_text(
+                    dc_text,
+                    quote=True,
+                    link_preview_options=LinkPreviewOptions(is_disabled=True)
+                )
+            except Exception as e:
+                logger.error(f"Error processing file info for DC command: {e}")
+                await handle_user_error(message, MSG_DC_FILE_ERROR)
+
+        args = message.text.strip().split(maxsplit=1)
+        if len(args) > 1:
+            query = args[1].strip()
+            user = await get_user_safely(bot, query)
+            if user:
+                await process_dc_info(user)
+            else:
+                await handle_user_error(message, MSG_ERROR_USER_INFO)
+            return
+
+        if message.reply_to_message:
+            if has_media(message.reply_to_message):
+                await process_file_dc_info(message.reply_to_message)
+                return
+            elif message.reply_to_message.from_user:
+                await process_dc_info(message.reply_to_message.from_user)
+                return
+            else:
+                await handle_user_error(message, MSG_DC_INVALID_USAGE)
+                return
+
+        if message.from_user:
+            await process_dc_info(message.from_user)
+        else:
+            await handle_user_error(message, MSG_DC_ANON_ERROR)
+    except Exception as e:
+        logger.error(f"Error in dc_command: {e}")
+        await handle_user_error(message, MSG_ERROR_GENERIC)
+
+@check_banned
+@force_channel_check
+@StreamBot.on_message(filters.command("ping") & filters.private)
+async def ping_command(bot: Client, message: Message):
+    try:
+        start_time = time.time()
+        sent_msg = await message.reply_text(MSG_PING_START, quote=True, link_preview_options=LinkPreviewOptions(is_disabled=True))
+        end_time = time.time()
+        time_taken_ms = (end_time - start_time) * 1000
+
+        buttons = [
+            [InlineKeyboardButton(MSG_BUTTON_GET_HELP, callback_data="help_command"),
+             InlineKeyboardButton(MSG_BUTTON_CLOSE, callback_data="close_panel")]
+        ]
+
+        await sent_msg.edit_text(
+            MSG_PING_RESPONSE.format(time_taken_ms=time_taken_ms),
+            reply_markup=InlineKeyboardMarkup(buttons),
+            link_preview_options=LinkPreviewOptions(is_disabled=True)
+        )
+    except Exception as e:
+        logger.error(f"Error in ping_command: {e}")
+        await handle_user_error(message, MSG_ERROR_GENERIC)
+
+async def handle_callback_error(callback_query, error, operation="callback"):
+    error_id = uuid.uuid4().hex[:8]
+    logger.error(f"Error in {operation}: {error}")
+    try:
+        await callback_query.answer(
+            MSG_ERROR_GENERIC_CALLBACK.format(error_id=error_id),
+            show_alert=True
+        )
+    except Exception as e:
+        logger.error(f"Failed to send error callback: {e}")
+
+@StreamBot.on_callback_query(filters.regex(r"^close_panel$"))
+async def close_panel_callback(client: Client, callback_query: CallbackQuery):
+    try:
+        await callback_query.answer()
+        await callback_query.message.delete()
+        if callback_query.message.reply_to_message:
+            try:
+                ctx = callback_query.message.reply_to_message
+                await ctx.delete()
+                if ctx.reply_to_message:
+                    await ctx.reply_to_message.delete()
+            except Exception as e:
+                logger.warning(f"Error deleting command messages: {e}")
+    except Exception as e:
+        await handle_callback_error(callback_query, e, "close_panel_callback")
+    finally:
+        raise StopPropagation

Thunder/server/stream_routes.py

@@ -1,21 +1,10 @@
 # Thunder/server/stream_routes.py
 
-import asyncio
-import inspect
-import mimetypes
 import re
 import secrets
 import time
-import json
-import traceback
-from concurrent.futures import ThreadPoolExecutor
 from aiohttp import web
-from aiohttp.client_exceptions import ClientConnectionError
-from cachetools import LRUCache
-from functools import wraps
 from urllib.parse import unquote, quote
-from contextlib import asynccontextmanager
-
 from Thunder import StartTime, __version__
 from Thunder.bot import multi_clients, StreamBot, work_loads
 from Thunder.server.exceptions import FileNotFound, InvalidHash
@@ -23,149 +12,96 @@ from Thunder.utils.custom_dl import ByteStreamer
 from Thunder.utils.logger import logger
 from Thunder.utils.render_template import render_page
 from Thunder.utils.time_format import get_readable_time
-from Thunder.vars import Var
 
 routes = web.RouteTableDef()
 
-# Configuration constants.
 SECURE_HASH_LENGTH = 6
-CHUNK_SIZE = 1 * 1024 * 1024
-MAX_CLIENTS = 50
-THREADPOOL_MAX_WORKERS = 10
-
-# Updated regex: Correctly capture start and end values with named groups.
+CHUNK_SIZE = 1024 * 1024
+MAX_CONCURRENT_PER_CLIENT = 8
+RANGE_REGEX = re.compile(r"bytes=(?P<start>\d*)-(?P<end>\d*)")
 PATTERN_HASH_FIRST = re.compile(rf"^([a-zA-Z0-9_-]{{{SECURE_HASH_LENGTH}}})(\d+)(?:/.*)?$")
 PATTERN_ID_FIRST = re.compile(r"^(\d+)(?:/.*)?$")
-RANGE_REGEX = re.compile(r"bytes=(?P<start>\d*)-(?P<end>\d*)")
-
-# Global LRU cache for ByteStreamer instances.
-class_cache = LRUCache(maxsize=Var.CACHE_SIZE)
-cache_lock = asyncio.Lock()
+VALID_HASH_REGEX = re.compile(r'^[a-zA-Z0-9_-]+$')
 
-executor = ThreadPoolExecutor(max_workers=THREADPOOL_MAX_WORKERS)
+streamers = {}
 
-def json_error(status, message):
-    return json.dumps({"error": message})
-
-def exception_handler(func):
-    @wraps(func)
-    async def wrapper(request):
-        try:
-            return await func(request)
-        except InvalidHash as e:
-            logger.debug(f"InvalidHash exception: {e}")
-            raise web.HTTPForbidden(
-                text=json_error(403, "Invalid security credentials"),
-                content_type="application/json"
-            )
-        except FileNotFound as e:
-            logger.debug(f"FileNotFound exception: {e}")
-            raise web.HTTPNotFound(
-                text=json_error(404, "File not found"),
-                content_type="application/json"
-            )
-        except (ClientConnectionError, asyncio.CancelledError):
-            return web.Response(status=499)
-        except web.HTTPException:
-            raise
-        except Exception as e:
-            error_id = secrets.token_hex(6)
-            logger.error(f"Unhandled exception (ID: {error_id}): {str(e)}")
-            logger.error(f"Stack trace for error {error_id}:\n{traceback.format_exc()}")
-            
-            raise web.HTTPInternalServerError(
-                text=json_error(500, f"Internal server error (Reference ID: {error_id})"),
-                content_type="application/json"
-            )
-    return wrapper
+def get_streamer(client_id: int) -> ByteStreamer:
+    if client_id not in streamers:
+        streamers[client_id] = ByteStreamer(multi_clients[client_id])
+    return streamers[client_id]
 
 def parse_media_request(path: str, query: dict) -> tuple[int, str]:
     clean_path = unquote(path).strip('/')
+    
     match = PATTERN_HASH_FIRST.match(clean_path)
     if match:
         try:
             message_id = int(match.group(2))
             secure_hash = match.group(1)
-            if not re.match(r'^[a-zA-Z0-9_-]+$', secure_hash):
-                raise InvalidHash("Security token contains invalid characters")
-            return message_id, secure_hash
-        except ValueError:
-            raise InvalidHash("Invalid message ID format")
+            if len(secure_hash) == SECURE_HASH_LENGTH and VALID_HASH_REGEX.match(secure_hash):
+                return message_id, secure_hash
+        except ValueError as e:
+            raise InvalidHash("Invalid message ID format") from e
+    
     match = PATTERN_ID_FIRST.match(clean_path)
     if match:
         try:
             message_id = int(match.group(1))
             secure_hash = query.get("hash", "").strip()
-            if len(secure_hash) != SECURE_HASH_LENGTH:
-                raise InvalidHash("Security token length mismatch")
-            if not re.match(r'^[a-zA-Z0-9_-]+$', secure_hash):
-                raise InvalidHash("Security token contains invalid characters")
-            return message_id, secure_hash
-        except ValueError:
-            raise InvalidHash("Invalid message ID format")
+            if len(secure_hash) == SECURE_HASH_LENGTH and VALID_HASH_REGEX.match(secure_hash):
+                return message_id, secure_hash
+        except ValueError as e:
+            raise InvalidHash("Invalid message ID format") from e
+    
     raise InvalidHash("Invalid URL structure")
 
-def optimal_client_selection():
+def select_optimal_client() -> tuple[int, ByteStreamer]:
     if not work_loads:
-        raise web.HTTPInternalServerError(
-            text=json_error(500, "No available clients"),
-            content_type="application/json"
-        )
-    client_id, _ = min(work_loads.items(), key=lambda item: item[1])
-    return client_id, multi_clients[client_id]
-
-async def get_cached_streamer(client) -> ByteStreamer:
-    streamer = class_cache.get(client)
-    if streamer is None:
-        async with cache_lock:
-            streamer = class_cache.get(client)
-            if streamer is None:
-                streamer = ByteStreamer(client)
-                class_cache[client] = streamer
-    return streamer
-
-def sanitize_filename(filename: str) -> str:
-    return re.sub(r'[\n\r";]', '', filename)
+        raise web.HTTPInternalServerError(text="No available clients")
+    
+    available_clients = [(cid, load) for cid, load in work_loads.items() if load < MAX_CONCURRENT_PER_CLIENT]
+    
+    if available_clients:
+        client_id = min(available_clients, key=lambda x: x[1])[0]
+    else:
+        client_id = min(work_loads, key=work_loads.get)
+    
+    return client_id, get_streamer(client_id)
 
-@asynccontextmanager
-async def track_workload(client_id):
-    work_loads[client_id] += 1
-    try:
-        yield
-    finally:
-        work_loads[client_id] -= 1
+def parse_range_header(range_header: str, file_size: int) -> tuple[int, int]:
+    if not range_header:
+        return 0, file_size - 1
+    
+    match = RANGE_REGEX.match(range_header)
+    if not match:
+        raise web.HTTPBadRequest(text="Invalid range")
+    
+    start = int(match.group("start")) if match.group("start") else 0
+    end = int(match.group("end")) if match.group("end") else file_size - 1
+    
+    if start < 0 or end >= file_size or start > end:
+        raise web.HTTPRequestRangeNotSatisfiable(
+            headers={"Content-Range": f"bytes */{file_size}"}
+        )
+    
+    return start, end
 
 @routes.get("/", allow_head=True)
-@exception_handler
 async def root_redirect(request):
     raise web.HTTPFound("https://github.com/fyaz05/FileToLink")
 
 @routes.get("/status", allow_head=True)
 async def status_endpoint(request):
-    uptime_seconds = time.time() - StartTime
-    uptime_formatted = get_readable_time(uptime_seconds)
-    
-    sorted_workloads = dict(sorted(
-        {f"client_{k}": v for k, v in work_loads.items()}.items(),
-        key=lambda item: item[1]
-    ))
-    
+    uptime = time.time() - StartTime
     total_load = sum(work_loads.values())
-    if total_load < MAX_CLIENTS * 0.5:
-        status_level = "optimal"
-    elif total_load < MAX_CLIENTS * 0.8:
-        status_level = "high"
-    else:
-        status_level = "critical"
     
-    status_data = {
+    workload_distribution = {str(k): v for k, v in sorted(work_loads.items())}
+    
+    return web.json_response({
         "server": {
             "status": "operational",
-            "status_level": status_level,
             "version": __version__,
-            "uptime": uptime_formatted,
-            "uptime_seconds": int(uptime_seconds),
-            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())
+            "uptime": get_readable_time(uptime)
         },
         "telegram_bot": {
             "username": f"@{StreamBot.username}",
@@ -173,131 +109,115 @@ async def status_endpoint(request):
         },
         "resources": {
             "total_workload": total_load,
-            "max_clients": MAX_CLIENTS,
-            "workload_distribution": sorted_workloads,
-            "cache": {
-                "current": len(class_cache),
-                "maximum": Var.CACHE_SIZE,
-                "utilization_percent": round((len(class_cache) / Var.CACHE_SIZE) * 100, 2)
-            }
-        },
-        "system": {
-            "chunk_size": f"{CHUNK_SIZE / (1024 * 1024):.1f} MB",
-            "thread_pool_workers": THREADPOOL_MAX_WORKERS
+            "workload_distribution": workload_distribution
         }
-    }
-    
-    return web.json_response(
-        status_data,
-        dumps=lambda obj: json.dumps(obj, indent=2, sort_keys=False)
-    )
+    })
 
 @routes.get(r"/watch/{path:.+}", allow_head=True)
-@exception_handler
 async def media_preview(request: web.Request):
-    path = request.match_info["path"]
-    message_id, secure_hash = parse_media_request(path, request.query)
-    rendered_page = await render_page(message_id, secure_hash, requested_action='stream')
-    return web.Response(
-        text=rendered_page,
-        content_type='text/html',
-        headers={
-            "Cache-Control": "no-cache, must-revalidate",
-            "Content-Security-Policy": "default-src 'self' https:; script-src 'self' 'unsafe-inline' https://cdn.jsdelivr.net https://cdnjs.cloudflare.com https://cdn.plyr.io; style-src 'self' 'unsafe-inline' https://fonts.googleapis.com https://cdn.jsdelivr.net https://cdnjs.cloudflare.com https://cdn.plyr.io; img-src 'self' data: https:; font-src 'self' https://fonts.gstatic.com https://cdnjs.cloudflare.com; connect-src 'self' https:; media-src 'self' https:;",
-            "X-Content-Type-Options": "nosniff",
-            "X-Frame-Options": "SAMEORIGIN",  
-            "X-XSS-Protection": "1; mode=block",
-            "Referrer-Policy": "strict-origin-when-cross-origin"
-        }
-    )
+    try:
+        path = request.match_info["path"]
+        message_id, secure_hash = parse_media_request(path, request.query)
+        
+        rendered_page = await render_page(message_id, secure_hash, requested_action='stream')
+        return web.Response(text=rendered_page, content_type='text/html')
+        
+    except (InvalidHash, FileNotFound) as e:
+        logger.debug(f"Client error in preview: {type(e).__name__}")
+        raise web.HTTPNotFound(text="Resource not found")
+    except Exception as e:
+        error_id = secrets.token_hex(6)
+        logger.error(f"Preview error {error_id}: {e}")
+        raise web.HTTPInternalServerError(text="Server error") from e
 
 @routes.get(r"/{path:.+}", allow_head=True)
-@exception_handler
 async def media_delivery(request: web.Request):
-    client_id, client = optimal_client_selection()
-    async with track_workload(client_id):
+    try:
         path = request.match_info["path"]
         message_id, secure_hash = parse_media_request(path, request.query)
-        return await handle_media_stream(request, message_id, secure_hash, client_id, client)
-
-async def handle_media_stream(request, message_id, secure_hash, client_id, client):
-    streamer = await get_cached_streamer(client)
-    file_meta = await streamer.get_file_properties(message_id)
-    if file_meta.unique_id[:SECURE_HASH_LENGTH] != secure_hash:
-        raise InvalidHash("Security token mismatch")
-    file_size = file_meta.file_size
-    range_header = request.headers.get("Range", "")
-    if range_header:
-        range_match = RANGE_REGEX.fullmatch(range_header)
-        if not range_match:
-            logger.debug(f"Malformed range header received: {range_header}")
-            raise web.HTTPBadRequest(
-                text=json_error(400, "Malformed range header"),
-                content_type="application/json"
-            )
-        start = int(range_match.group("start")) if range_match.group("start") else 0
-        end = int(range_match.group("end")) if range_match.group("end") else file_size - 1
-    else:
+        
+        client_id, streamer = select_optimal_client()
+        
+        work_loads[client_id] += 1
+        
         try:
-            if hasattr(request, "http_range") and request.http_range:
-                start = request.http_range.start or 0
-                end = (request.http_range.stop or file_size) - 1
-            else:
-                start, end = 0, file_size - 1
-        except Exception:
-            start, end = 0, file_size - 1
-    if start < 0 or end >= file_size or start > end:
-        logger.debug(f"Invalid range request: start={start}, end={end}, file_size={file_size}")
-        raise web.HTTPRequestRangeNotSatisfiable(
-            headers={"Content-Range": f"bytes */{file_size}"}
-        )
-    offset = start - (start % CHUNK_SIZE)
-    first_chunk_cut = start - offset
-    last_chunk_cut = (end % CHUNK_SIZE) + 1
-    total_chunks = ((end - offset) // CHUNK_SIZE) + 1
-    mime_type = file_meta.mime_type or mimetypes.guess_type(file_meta.file_name)[0] or "application/octet-stream"
-    disposition = "inline" if mime_type.startswith(("video/", "audio/")) else "attachment"
-    original_filename = unquote(file_meta.file_name) if file_meta.file_name else f"file_{secrets.token_hex(4)}"
-    safe_filename = sanitize_filename(original_filename)
-    encoded_filename = quote(safe_filename)
-    headers = {
-        "Content-Type": mime_type,
-        "Content-Range": f"bytes {start}-{end}/{file_size}",
-        "Content-Length": str(end - start + 1),
-        "Content-Disposition": f"{disposition}; filename*=UTF-8''{encoded_filename}",
-        "Accept-Ranges": "bytes",
-        "Cache-Control": "public, max-age=31536000, immutable",
-        "Connection": "keep-alive",
-        "X-Content-Type-Options": "nosniff",
-        "Referrer-Policy": "strict-origin-when-cross-origin"
-    }
-    if hasattr(streamer, 'async_yield_file'):
-        stream_generator = streamer.async_yield_file(
-            file_meta, client_id, offset, first_chunk_cut, last_chunk_cut, total_chunks, CHUNK_SIZE
-        )
-    else:
-        stream_generator = streamer.yield_file(
-            file_meta, client_id, offset, first_chunk_cut, last_chunk_cut, total_chunks, CHUNK_SIZE
-        )
-    if inspect.isgenerator(stream_generator):
-        stream_generator = async_gen_wrapper(stream_generator)
-    return web.Response(
-        status=206 if range_header else 200,
-        body=stream_generator,
-        headers=headers
-    )
-
-async def async_gen_wrapper(sync_gen):
-    loop = asyncio.get_running_loop()
-    try:
-        while True:
-            try:
-                yield await loop.run_in_executor(executor, next, sync_gen)
-            except StopIteration:
-                break
+            file_info = await streamer.get_file_info(message_id)
+            if not file_info.get('unique_id'):
+                raise FileNotFound("File not found")
+            
+            if file_info['unique_id'][:SECURE_HASH_LENGTH] != secure_hash:
+                raise InvalidHash("Invalid hash")
+            
+            file_size = file_info.get('file_size', 0)
+            if file_size == 0:
+                raise FileNotFound("File size unavailable")
+            
+            range_header = request.headers.get("Range", "")
+            start, end = parse_range_header(range_header, file_size)
+            content_length = end - start + 1
+            
+            if start == 0 and end == file_size - 1:
+                range_header = ""
+            
+            mime_type = file_info.get('mime_type') or 'application/octet-stream'
+            filename = file_info.get('file_name') or f"file_{secrets.token_hex(4)}"
+            
+            headers = {
+                "Content-Type": mime_type,
+                "Content-Length": str(content_length),
+                "Content-Disposition": f"inline; filename*=UTF-8''{quote(filename)}",
+                "Accept-Ranges": "bytes",
+                "Cache-Control": "public, max-age=31536000",
+                "Connection": "keep-alive"
+            }
+            
+            if range_header:
+                headers["Content-Range"] = f"bytes {start}-{end}/{file_size}"
+            
+            async def stream_generator():
+                try:
+                    bytes_sent = 0
+                    bytes_to_skip = start % CHUNK_SIZE
+                    
+                    async for chunk in streamer.stream_file(message_id, offset=start, limit=content_length):
+                        if bytes_to_skip > 0:
+                            if len(chunk) <= bytes_to_skip:
+                                bytes_to_skip -= len(chunk)
+                                continue
+                            chunk = chunk[bytes_to_skip:]
+                            bytes_to_skip = 0
+                        
+                        remaining = content_length - bytes_sent
+                        if len(chunk) > remaining:
+                            chunk = chunk[:remaining]
+                        
+                        if chunk:
+                            yield chunk
+                            bytes_sent += len(chunk)
+                        
+                        if bytes_sent >= content_length:
+                            break
+                finally:
+                    work_loads[client_id] -= 1
+            return web.Response(
+                status=206 if range_header else 200,
+                body=stream_generator(),
+                headers=headers
+            )
+            
+        except (FileNotFound, InvalidHash):
+            work_loads[client_id] -= 1
+            raise
+        except Exception as e:
+            work_loads[client_id] -= 1
+            error_id = secrets.token_hex(6)
+            logger.error(f"Stream error {error_id}: {e}")
+            raise web.HTTPInternalServerError(text="Server error") from e
+        
+    except (InvalidHash, FileNotFound) as e:
+        logger.debug(f"Client error: {type(e).__name__}")
+        raise web.HTTPNotFound(text="Resource not found")
     except Exception as e:
-        logger.error(f"Error in async generator wrapper: {e}")
-        try:
-            sync_gen.close()
-        except Exception:
-            pass
+        error_id = secrets.token_hex(6)
+        logger.error(f"Server error {error_id}: {e}")
+        raise web.HTTPInternalServerError(text="Server error") from e

Thunder/utils/custom_dl.py

@@ -1,371 +1,60 @@
 # Thunder/utils/custom_dl.py
 
-import math
-import asyncio
-import time
-from typing import Dict, Union, Optional, AsyncGenerator
-from pyrogram import Client, utils, raw
-from pyrogram.session import Session, Auth
-from pyrogram.errors import AuthBytesInvalid, RPCError, FloodWait, FileReferenceExpired, FileReferenceInvalid
-from pyrogram.file_id import FileId, FileType, ThumbnailSource
+from typing import Dict, Any, AsyncGenerator
+from pyrogram import Client
+from pyrogram.types import Message
 from Thunder.vars import Var
-from Thunder.bot import work_loads
 from Thunder.server.exceptions import FileNotFound
-from .file_properties import get_file_ids
 from Thunder.utils.logger import logger
 
 class ByteStreamer:
+    __slots__ = ('client', 'chat_id')
     
     def __init__(self, client: Client):
         self.client = client
-        self.clean_timer = 30 * 60
-        self.cached_file_ids: Dict[int, FileId] = {}
-        self.file_references_cache: Dict[int, Dict[str, Union[FileId, float]]] = {}
-        self.cache_lock = asyncio.Lock()
-        self.media_sessions_lock = asyncio.Lock()
-        
-        self.cache_cleaner_task = asyncio.create_task(self.clean_cache())
-        self.session_cleaner_task = asyncio.create_task(self.cleanup_media_sessions())
-    
-    async def shutdown(self):
-        if hasattr(self, 'cache_cleaner_task'):
-            self.cache_cleaner_task.cancel()
-            try:
-                await self.cache_cleaner_task
-            except asyncio.CancelledError:
-                pass
-        
-        if hasattr(self, 'session_cleaner_task'):
-            self.session_cleaner_task.cancel()
-            try:
-                await self.session_cleaner_task
-            except asyncio.CancelledError:
-                pass
-        
-        for dc_id, session in list(self.client.media_sessions.items()):
-            try:
-                await session.stop()
-            except Exception as e:
-                logger.debug(f"Error stopping media session for DC {dc_id}: {e}")
-        
-        self.client.media_sessions.clear()
-    
-    async def get_file_properties(self, message_id: int) -> FileId:
-        async with self.cache_lock:
-            file_id = self.cached_file_ids.get(message_id)
-        
-        if not file_id:
-            file_id = await self.generate_file_properties(message_id)
-            async with self.cache_lock:
-                self.cached_file_ids[message_id] = file_id
-        
-        return file_id
-    
-    async def generate_file_properties(self, message_id: int) -> FileId:
-        file_id = await get_file_ids(self.client, Var.BIN_CHANNEL, message_id)
-        
-        if not file_id:
-            logger.debug(f"Message ID {message_id} not found in the channel.")
-            raise FileNotFound(f"File with message ID {message_id} not found.")
-        
-        async with self.cache_lock:
-            self.cached_file_ids[message_id] = file_id
-        
-        return file_id
-    
-    async def refresh_file_reference(self, file_id: FileId) -> Optional[FileId]:
-        logger.debug(f"Refreshing file reference for message ID {file_id.message_id}")
-        
-        async with self.cache_lock:
-            cached_ref = self.file_references_cache.get(file_id.message_id)
-            if cached_ref and time.time() - cached_ref["timestamp"] < 3600:
-                return cached_ref["file_id"]
-        
+        self.chat_id = int(Var.BIN_CHANNEL)
+
+    async def get_message(self, message_id: int) -> Message:
         try:
-            message = await self.client.get_messages(
-                chat_id=Var.BIN_CHANNEL,
-                message_ids=file_id.message_id
-            )
-            
+            message = await self.client.get_messages(self.chat_id, message_id)
             if not message or not message.media:
-                logger.debug(f"Media message {file_id.message_id} not found when refreshing file reference")
-                return None
-            
-            new_file_id = await get_file_ids(self.client, Var.BIN_CHANNEL, file_id.message_id)
-            if not new_file_id:
-                logger.error(f"Failed to extract file ID from message {file_id.message_id}")
-                return None
-            
-            async with self.cache_lock:
-                self.file_references_cache[file_id.message_id] = {
-                    "file_id": new_file_id,
-                    "timestamp": time.time()
-                }
-            
-            return new_file_id
+                raise FileNotFound(f"Message {message_id} not found")
+            return message
         except Exception as e:
-            logger.error(f"Error refreshing file reference: {e}")
-            return None
-    
-    async def generate_media_session(self, client: Client, file_id: FileId) -> Session:
-        async with self.media_sessions_lock:
-            media_session = client.media_sessions.get(file_id.dc_id, None)
-            
-            if media_session is None:
-                if file_id.dc_id != await client.storage.dc_id():
-                    media_session = Session(
-                        client,
-                        file_id.dc_id,
-                        await Auth(
-                            client, file_id.dc_id, await client.storage.test_mode()
-                        ).create(),
-                        await client.storage.test_mode(),
-                        is_media=True,
-                    )
-                    await media_session.start()
-                    
-                    for attempt in range(6):
-                        if attempt > 0:
-                            await asyncio.sleep(1 * (2 ** (attempt - 1)))
-                        
-                        try:
-                            exported_auth = await client.invoke(
-                                raw.functions.auth.ExportAuthorization(dc_id=file_id.dc_id)
-                            )
-                            
-                            try:
-                                await media_session.send(
-                                    raw.functions.auth.ImportAuthorization(
-                                        id=exported_auth.id, bytes=exported_auth.bytes
-                                    )
-                                )
-                                break
-                            except AuthBytesInvalid:
-                                continue
-                        except FloodWait as e:
-                            logger.warning(f"FloodWait during auth export: {e.value} seconds")
-                            await asyncio.sleep(e.value + 1)
-                            continue
-                    else:
-                        await media_session.stop()
-                        raise AuthBytesInvalid
-                else:
-                    media_session = Session(
-                        client,
-                        file_id.dc_id,
-                        await client.storage.auth_key(),
-                        await client.storage.test_mode(),
-                        is_media=True,
-                    )
-                    await media_session.start()
-                
-                setattr(media_session, 'last_used', time.time())
-                client.media_sessions[file_id.dc_id] = media_session
-            else:
-                setattr(media_session, 'last_used', time.time())
-        
-        return media_session
-    
-    @staticmethod
-    async def get_location(file_id: FileId) -> Union[
-        raw.types.InputPhotoFileLocation,
-        raw.types.InputDocumentFileLocation,
-        raw.types.InputPeerPhotoFileLocation,
-        None
-    ]:
-        file_type = file_id.file_type
-        
-        if file_type == FileType.CHAT_PHOTO:
-            if file_id.chat_id > 0:
-                peer = raw.types.InputPeerUser(
-                    user_id=file_id.chat_id, access_hash=file_id.chat_access_hash
-                )
-            else:
-                if file_id.chat_access_hash == 0:
-                    peer = raw.types.InputPeerChat(chat_id=-file_id.chat_id)
-                else:
-                    peer = raw.types.InputPeerChannel(
-                        channel_id=utils.get_channel_id(file_id.chat_id),
-                        access_hash=file_id.chat_access_hash,
-                    )
-            location = raw.types.InputPeerPhotoFileLocation(
-                peer=peer,
-                volume_id=file_id.volume_id,
-                local_id=file_id.local_id,
-                big=file_id.thumbnail_source == ThumbnailSource.CHAT_PHOTO_BIG,
-            )
-        elif file_type == FileType.PHOTO:
-            if not file_id.file_reference:
-                logger.debug("Missing file_reference for PHOTO type")
-                return None
-            location = raw.types.InputPhotoFileLocation(
-                id=file_id.media_id,
-                access_hash=file_id.access_hash,
-                file_reference=file_id.file_reference,
-                thumb_size=file_id.thumbnail_size,
-            )
-        else:
-            if not file_id.file_reference:
-                logger.debug("Missing file_reference for DOCUMENT type")
-                return None
-            location = raw.types.InputDocumentFileLocation(
-                id=file_id.media_id,
-                access_hash=file_id.access_hash,
-                file_reference=file_id.file_reference,
-                thumb_size=file_id.thumbnail_size,
-            )
-        return location
-    
-    async def yield_file(
-        self,
-        file_id: FileId,
-        index: int,
-        offset: int,
-        first_part_cut: int,
-        last_part_cut: int,
-        part_count: int,
-        chunk_size: int,
-    ) -> AsyncGenerator[bytes, None]:
-        client = self.client
-        work_loads[index] += 1
+            logger.debug(f"Error fetching message {message_id}: {e}")
+            raise FileNotFound(f"Message {message_id} not found") from e
+
+    async def stream_file(self, message_id: int, offset: int = 0, limit: int = 0) -> AsyncGenerator[bytes, None]:
+        message = await self.get_message(message_id)
         
-        try:
-            media_session = await self.generate_media_session(client, file_id)
-            current_part = 1
-            
-            # Get file location
-            location = await self.get_location(file_id)
-            
-            # Handle missing file reference
-            if location is None:
-                logger.warning(f"Missing file reference, attempting to refresh")
-                refreshed_file_id = await self.refresh_file_reference(file_id)
-                if refreshed_file_id:
-                    location = await self.get_location(refreshed_file_id)
-                    file_id = refreshed_file_id
-                
-                if location is None:
-                    logger.error("Failed to get location even after refreshing file reference")
-                    return
+        if limit > 0:
+            chunk_offset = offset // (1024 * 1024)
+            chunk_limit = (limit + 1024 * 1024 - 1) // (1024 * 1024)
             
-            max_retries = 5
-            retry_count = 0
-            
-            while current_part <= part_count:
-                try:
-                    r = await media_session.send(
-                        raw.functions.upload.GetFile(
-                            location=location, offset=offset, limit=chunk_size
-                        ),
-                    )
-                except FloodWait as e:
-                    logger.warning(f"FloodWait encountered: waiting for {e.value} seconds.")
-                    await asyncio.sleep(e.value + 1)
-                    continue
-                except (FileReferenceExpired, FileReferenceInvalid) as e:
-                    logger.debug(f"File reference expired: {e}")
-                    refreshed_file_id = await self.refresh_file_reference(file_id)
-                    if refreshed_file_id:
-                        file_id = refreshed_file_id
-                        location = await self.get_location(file_id)
-                        if location is None:
-                            logger.error("Failed to get location after refresh")
-                            return
-                        retry_count = 0
-                        continue
-                    else:
-                        logger.error("Failed to refresh file reference")
-                        return
-                except RPCError as e:
-                    retry_count += 1
-                    logger.warning(f"RPC Error during file fetch: {e}")
-                    if retry_count >= max_retries:
-                        return
-                    await asyncio.sleep(1 * (2 ** (retry_count - 1)))
-                    continue
-                except ConnectionError as e:
-                    retry_count += 1
-                    logger.error(f"Connection error during file fetch: {e}")
-                    if retry_count >= max_retries:
-                        return
-                    await asyncio.sleep(1)
-                    continue
-                except Exception as e:
-                    logger.error(f"Unexpected error during file fetch: {str(e)}", exc_info=True)
-                    return
-                
-                if isinstance(r, raw.types.upload.File):
-                    chunk = r.bytes
-                    if not chunk:
-                        break
-                    elif part_count == 1:
-                        yield chunk[first_part_cut:last_part_cut]
-                    elif current_part == 1:
-                        yield chunk[first_part_cut:]
-                    elif current_part == part_count:
-                        yield chunk[:last_part_cut]
-                    else:
-                        yield chunk
+            async for chunk in self.client.stream_media(message, offset=chunk_offset, limit=chunk_limit):
+                yield chunk
+        else:
+            async for chunk in self.client.stream_media(message):
+                yield chunk
 
-                    current_part += 1
-                    offset += chunk_size
-                    retry_count = 0
+    def get_file_info_sync(self, message: Message) -> Dict[str, Any]:
+        media = message.document or message.video or message.audio or message.photo
+        if not media:
+            return {"message_id": message.id, "error": "No media"}
+        
+        return {
+            "message_id": message.id,
+            "file_size": getattr(media, 'file_size', 0) or 0,
+            "file_name": getattr(media, 'file_name', None),
+            "mime_type": getattr(media, 'mime_type', None),
+            "unique_id": getattr(media, 'file_unique_id', None),
+            "media_type": type(media).__name__.lower()
+        }
 
-                    if current_part > part_count:
-                        break
-        except (TimeoutError, AttributeError) as e:
-            logger.error(f"Error while yielding file: {e}")
-        finally:
-            work_loads[index] -= 1
-    
-    async def clean_cache(self) -> None:
-        while True:
-            try:
-                await asyncio.sleep(self.clean_timer)
-                
-                async with self.cache_lock:
-                    self.cached_file_ids.clear()
-                    
-                    now = time.time()
-                    expired_refs = [
-                        key for key, value in self.file_references_cache.items()
-                        if now - value["timestamp"] > 3600
-                    ]
-                    for key in expired_refs:
-                        self.file_references_cache.pop(key, None)
-                
-                if expired_refs:
-                    logger.debug(f"Cache cleaned. Removed {len(expired_refs)} expired file references.")
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                logger.error(f"Error in cache cleaning task: {e}")
-                await asyncio.sleep(60)
-    
-    async def cleanup_media_sessions(self) -> None:
-        while True:
-            try:
-                await asyncio.sleep(1800)
-                
-                async with self.media_sessions_lock:
-                    now = time.time()
-                    sessions_to_remove = []
-                    
-                    for dc_id, session in list(self.client.media_sessions.items()):
-                        if hasattr(session, 'last_used') and now - session.last_used > 1800:  # 30 minutes
-                            sessions_to_remove.append(dc_id)
-                    
-                    for dc_id in sessions_to_remove:
-                        try:
-                            await self.client.media_sessions[dc_id].stop()
-                            self.client.media_sessions.pop(dc_id, None)
-                        except Exception as e:
-                            logger.error(f"Error stopping media session for DC {dc_id}: {e}")
-                
-                if sessions_to_remove:
-                    logger.debug(f"Cleaned up {len(sessions_to_remove)} inactive media sessions")
-            except asyncio.CancelledError:
-                break
-            except Exception as e:
-                logger.error(f"Error in media session cleaning task: {e}")
-                await asyncio.sleep(60)
+    async def get_file_info(self, message_id: int) -> Dict[str, Any]:
+        try:
+            message = await self.get_message(message_id)
+            return self.get_file_info_sync(message)
+        except Exception as e:
+            logger.debug(f"Error getting file info for {message_id}: {e}")
+            return {"message_id": message_id, "error": str(e)}

Thunder/utils/database.py

@@ -1,190 +1,183 @@
-# Thunder/utils/database.py
-
-import datetime
-from typing import Optional, Dict, Any
-from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection
-from Thunder.vars import Var
-from Thunder.utils.logger import logger
-from Thunder.utils.error_handling import log_errors
-
-class Database:
-    def __init__(self, uri: str, database_name: str, *args, **kwargs):
-        self._client = AsyncIOMotorClient(uri, *args, **kwargs)
-        self.db = self._client[database_name]
-        self.col: AsyncIOMotorCollection = self.db.users
-        self.banned_users_col: AsyncIOMotorCollection = self.db.banned_users
-        self.token_col: AsyncIOMotorCollection = self.db.tokens
-        self.authorized_users_col: AsyncIOMotorCollection = self.db.authorized_users
-        self.restart_message_col: AsyncIOMotorCollection = self.db.restart_message
-
-    @log_errors
-    async def ensure_indexes(self):
-        await self.banned_users_col.create_index("user_id", unique=True)
-        await self.token_col.create_index("token", unique=True)
-        await self.authorized_users_col.create_index("user_id", unique=True)
-        await self.col.create_index("id", unique=True)
-        await self.token_col.create_index("expires_at", expireAfterSeconds=0)
-        await self.restart_message_col.create_index("message_id", unique=True)
-        await self.restart_message_col.create_index("timestamp", expireAfterSeconds=3600)
-        logger.info("Database indexes ensured.")
-
-    @log_errors
-    def new_user(self, user_id: int) -> dict:
-        return {
-            'id': user_id,
-            'join_date': datetime.datetime.utcnow()
-        }
-
-    @log_errors
-    async def add_user(self, user_id: int):
-        if not await self.is_user_exist(user_id):
-            await self.col.insert_one(self.new_user(user_id))
-            logger.debug(f"Added new user {user_id} to database.")
-
-    @log_errors
-    async def add_user_pass(self, user_id: int, ag_pass: str):
-        await self.add_user(user_id)
-        await self.col.update_one({'id': user_id}, {'$set': {'ag_p': ag_pass}})
-        logger.debug(f"Updated password for user {user_id}.")
-
-    @log_errors
-    async def get_user_pass(self, user_id: int) -> Optional[str]:
-        user_data = await self.col.find_one({'id': user_id}, {'ag_p': 1})
-        return user_data.get('ag_p') if user_data else None
-
-    @log_errors
-    async def is_user_exist(self, user_id: int) -> bool:
-        user = await self.col.find_one({'id': user_id}, {'_id': 1})
-        return bool(user)
-
-    @log_errors
-    async def total_users_count(self) -> int:
-        return await self.col.count_documents({})
-
-    @log_errors
-    async def get_all_users(self):
-        return self.col.find({})
-
-    @log_errors
-    async def delete_user(self, user_id: int):
-        await self.col.delete_one({'id': user_id})
-        logger.info(f"Deleted user {user_id}.")
-
-    @log_errors
-    async def create_index(self):
-        await self.col.create_index("id", unique=True)
-        logger.info("Created index for 'id' on users collection.")
-
-    @log_errors
-    async def get_active_users(self, days: int = 7):
-        cutoff = datetime.datetime.utcnow() - datetime.timedelta(days=days)
-        return self.col.find({'join_date': {'$gte': cutoff}})
-
-    @log_errors
-    async def add_banned_user(
-        self, user_id: int, banned_by: Optional[int] = None,
-        reason: Optional[str] = None, ban_time: Optional[str] = None
-    ):
-        ban_data = {
-            "user_id": user_id,
-            "banned_at": datetime.datetime.utcnow(),
-            "banned_by": banned_by,
-            "reason": reason
-        }
-        await self.banned_users_col.update_one(
-            {"user_id": user_id},
-            {"$set": ban_data},
-            upsert=True
-        )
-        logger.info(f"Added/Updated banned user {user_id}. Reason: {reason}")
-
-    @log_errors
-    async def remove_banned_user(self, user_id: int) -> bool:
-        result = await self.banned_users_col.delete_one({"user_id": user_id})
-        if result.deleted_count > 0:
-            logger.info(f"Removed banned user {user_id}.")
-            return True
-        return False
-
-    @log_errors
-    async def is_user_banned(self, user_id: int) -> Optional[Dict[str, Any]]:
-        return await self.banned_users_col.find_one({"user_id": user_id})
-
-    @log_errors
-    async def check_user_token(self, user_id: int) -> Dict[str, bool]:
-        token_data = await self.token_col.find_one({
-            "user_id": user_id,
-            "expires_at": {"$gt": datetime.datetime.utcnow()}
-        })
-        return {"has_token": bool(token_data)}
-
-    @log_errors
-    async def update_user_token(self, user_id: int, token: str) -> None:
-        token_data = await self.token_col.find_one({"token": token})
-        expires_at = datetime.datetime.utcnow() + datetime.timedelta(hours=Var.TOKEN_TTL_HOURS)
-        if token_data:
-            await self.token_col.update_one(
-                {"token": token},
-                {"$set": {"user_id": user_id, "expires_at": expires_at}}
-            )
-            logger.info(f"Updated token for user {user_id}.")
-        else:
-            await self.token_col.insert_one({
-                "user_id": user_id,
-                "token": token,
-                "created_at": datetime.datetime.utcnow(),
-                "expires_at": expires_at
-            })
-            logger.info(f"Inserted new token for user {user_id}.")
-
-    @log_errors
-    async def save_broadcast_state(self, broadcast_id, state_data):
-        await self.db.broadcasts.update_one(
-            {"_id": broadcast_id},
-            {"$set": state_data},
-            upsert=True
-        )
-        logger.debug(f"Saved broadcast state for ID {broadcast_id}.")
-
-    @log_errors
-    async def get_broadcast_state(self, broadcast_id):
-        return await self.db.broadcasts.find_one({"_id": broadcast_id})
-
-    @log_errors
-    async def list_active_broadcasts(self):
-        cursor = self.db.broadcasts.find({"is_cancelled": False})
-        return await cursor.to_list(length=None)
-
-    @log_errors
-    async def add_restart_message(self, message_id: int, chat_id: int) -> None:
-        try:
-            await self.restart_message_col.insert_one({
-                "message_id": message_id, 
-                "chat_id": chat_id, 
-                "timestamp": datetime.datetime.utcnow()
-            })
-            logger.info(f"Added restart message {message_id} for chat {chat_id}.")
-        except Exception as e:
-            logger.error(f"Error adding restart message {message_id}: {e}")
-
-    @log_errors
-    async def get_restart_message(self) -> Optional[Dict[str, Any]]:
-        try:
-            return await self.restart_message_col.find_one(sort=[("timestamp", -1)])
-        except Exception as e:
-            logger.error(f"Error getting restart message: {e}")
-            return None
-
-    @log_errors
-    async def delete_restart_message(self, message_id: int) -> None:
-        try:
-            await self.restart_message_col.delete_one({"message_id": message_id})
-            logger.info(f"Deleted restart message {message_id}.")
-        except Exception as e:
-            logger.error(f"Error deleting restart message {message_id}: {e}")
-
-    async def close(self):
-        if self._client:
-            self._client.close()
-
-db = Database(Var.DATABASE_URL, Var.NAME)
+# Thunder/utils/database.py
+
+import datetime
+from typing import Optional, Dict, Any
+from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection
+from Thunder.vars import Var
+from Thunder.utils.logger import logger
+from Thunder.utils.error_handling import log_errors
+
+class Database:
+    def __init__(self, uri: str, database_name: str, *args, **kwargs):
+        self._client = AsyncIOMotorClient(uri, *args, **kwargs)
+        self.db = self._client[database_name]
+        self.col: AsyncIOMotorCollection = self.db.users
+        self.banned_users_col: AsyncIOMotorCollection = self.db.banned_users
+        self.token_col: AsyncIOMotorCollection = self.db.tokens
+        self.authorized_users_col: AsyncIOMotorCollection = self.db.authorized_users
+        self.restart_message_col: AsyncIOMotorCollection = self.db.restart_message
+
+    @log_errors
+    async def ensure_indexes(self):
+        await self.banned_users_col.create_index("user_id", unique=True)
+        await self.token_col.create_index("token", unique=True)
+        await self.authorized_users_col.create_index("user_id", unique=True)
+        await self.col.create_index("id", unique=True)
+        await self.token_col.create_index("expires_at", expireAfterSeconds=0)
+        await self.token_col.create_index("activated")
+        await self.restart_message_col.create_index("message_id", unique=True)
+        await self.restart_message_col.create_index("timestamp", expireAfterSeconds=3600)
+
+        logger.info("Database indexes ensured.")
+
+    @log_errors
+    def new_user(self, user_id: int) -> dict:
+        return {
+            'id': user_id,
+            'join_date': datetime.datetime.utcnow()
+        }
+
+    @log_errors
+    async def add_user(self, user_id: int):
+        if not await self.is_user_exist(user_id):
+            await self.col.insert_one(self.new_user(user_id))
+            logger.debug(f"Added new user {user_id} to database.")
+
+    @log_errors
+    async def add_user_pass(self, user_id: int, ag_pass: str):
+        await self.add_user(user_id)
+        await self.col.update_one({'id': user_id}, {'$set': {'ag_p': ag_pass}})
+        logger.debug(f"Updated password for user {user_id}.")
+
+    @log_errors
+    async def get_user_pass(self, user_id: int) -> Optional[str]:
+        user_data = await self.col.find_one({'id': user_id}, {'ag_p': 1})
+        return user_data.get('ag_p') if user_data else None
+
+    @log_errors
+    async def is_user_exist(self, user_id: int) -> bool:
+        user = await self.col.find_one({'id': user_id}, {'_id': 1})
+        return bool(user)
+
+    @log_errors
+    async def total_users_count(self) -> int:
+        return await self.col.count_documents({})
+
+    @log_errors
+    async def get_all_users(self):
+        return self.col.find({})
+
+    @log_errors
+    async def delete_user(self, user_id: int):
+        await self.col.delete_one({'id': user_id})
+        logger.info(f"Deleted user {user_id}.")
+
+    @log_errors
+    async def create_index(self):
+        await self.col.create_index("id", unique=True)
+        logger.info("Created index for 'id' on users collection.")
+
+    @log_errors
+    async def get_active_users(self, days: int = 7):
+        cutoff = datetime.datetime.utcnow() - datetime.timedelta(days=days)
+        return self.col.find({'join_date': {'$gte': cutoff}})
+
+    @log_errors
+    async def add_banned_user(
+        self, user_id: int, banned_by: Optional[int] = None,
+        reason: Optional[str] = None, ban_time: Optional[str] = None
+    ):
+        ban_data = {
+            "user_id": user_id,
+            "banned_at": datetime.datetime.utcnow(),
+            "banned_by": banned_by,
+            "reason": reason
+        }
+        await self.banned_users_col.update_one(
+            {"user_id": user_id},
+            {"$set": ban_data},
+            upsert=True
+        )
+        logger.info(f"Added/Updated banned user {user_id}. Reason: {reason}")
+
+    @log_errors
+    async def remove_banned_user(self, user_id: int) -> bool:
+        result = await self.banned_users_col.delete_one({"user_id": user_id})
+        if result.deleted_count > 0:
+            logger.info(f"Removed banned user {user_id}.")
+            return True
+        return False
+
+    @log_errors
+    async def is_user_banned(self, user_id: int) -> Optional[Dict[str, Any]]:
+        return await self.banned_users_col.find_one({"user_id": user_id})
+
+    @log_errors
+    async def save_main_token(self, user_id: int, token_value: str, expires_at: datetime.datetime, created_at: datetime.datetime, activated: bool) -> None:
+        try:
+            await self.token_col.update_one(
+                {"user_id": user_id, "token": token_value},
+                {"$set": {
+                    "expires_at": expires_at,
+                    "created_at": created_at,
+                    "activated": activated
+                    }
+                },
+                upsert=True
+            )
+            logger.info(f"Saved main token {token_value} for user {user_id} with activated status {activated}.")
+        except Exception as e:
+            logger.error(f"Error saving main token for user {user_id}: {e}")
+            raise
+
+    @log_errors
+    async def save_broadcast_state(self, broadcast_id, state_data):
+        await self.db.broadcasts.update_one(
+            {"_id": broadcast_id},
+            {"$set": state_data},
+            upsert=True
+        )
+        logger.debug(f"Saved broadcast state for ID {broadcast_id}.")
+
+    @log_errors
+    async def get_broadcast_state(self, broadcast_id):
+        return await self.db.broadcasts.find_one({"_id": broadcast_id})
+
+    @log_errors
+    async def list_active_broadcasts(self):
+        cursor = self.db.broadcasts.find({"is_cancelled": False})
+        return await cursor.to_list(length=None)
+
+    @log_errors
+    async def add_restart_message(self, message_id: int, chat_id: int) -> None:
+        try:
+            await self.restart_message_col.insert_one({
+                "message_id": message_id,
+                "chat_id": chat_id,
+                "timestamp": datetime.datetime.utcnow()
+            })
+            logger.info(f"Added restart message {message_id} for chat {chat_id}.")
+        except Exception as e:
+            logger.error(f"Error adding restart message {message_id}: {e}")
+
+    @log_errors
+    async def get_restart_message(self) -> Optional[Dict[str, Any]]:
+        try:
+            return await self.restart_message_col.find_one(sort=[("timestamp", -1)])
+        except Exception as e:
+            logger.error(f"Error getting restart message: {e}")
+            return None
+
+    @log_errors
+    async def delete_restart_message(self, message_id: int) -> None:
+        try:
+            await self.restart_message_col.delete_one({"message_id": message_id})
+            logger.info(f"Deleted restart message {message_id}.")
+        except Exception as e:
+            logger.error(f"Error deleting restart message {message_id}: {e}")
+
+    async def close(self):
+        if self._client:
+            self._client.close()
+
+db = Database(Var.DATABASE_URL, Var.NAME)

Thunder/utils/decorators.py

@@ -1,147 +1,162 @@
-# Thunder/utils/decorators.py
-
-from functools import wraps
-from pyrogram.types import Message, InlineKeyboardButton, InlineKeyboardMarkup
-from Thunder.vars import Var
-from Thunder.utils.logger import logger
-from Thunder.utils.messages import (
-    MSG_DECORATOR_BANNED,
-    MSG_TOKEN_INVALID,
-    MSG_ERROR_UNAUTHORIZED
-)
-from Thunder.utils.database import db
-from Thunder.utils.tokens import check, allowed, generate
-from Thunder.utils.shortener import shorten
-
-
-def check_banned(func):
-    @wraps(func)
-    async def wrapper(client, message: Message, *args, **kwargs):
-        try:
-            if not message.from_user:
-                return await func(client, message, *args, **kwargs)
-            user_id = message.from_user.id
-            if isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID:
-                 return await func(client, message, *args, **kwargs)
-            elif isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID:
-                 return await func(client, message, *args, **kwargs)
-            
-            ban_details = await db.is_user_banned(user_id)
-            if ban_details:
-                banned_at = ban_details.get('banned_at')
-                ban_time = (
-                    banned_at.strftime('%B %d, %Y, %I:%M %p UTC') 
-                    if banned_at and hasattr(banned_at, 'strftime') 
-                    else str(banned_at) if banned_at else 'N/A'
-                )
-                await message.reply_text(
-                    MSG_DECORATOR_BANNED.format(
-                        reason=ban_details.get('reason', 'Not specified'),
-                        ban_time=ban_time
-                    ),
-                    quote=True
-                )
-                logger.info(f"Blocked banned user {user_id} from accessing {func.__name__}.")
-                return
-            return await func(client, message, *args, **kwargs)
-        except Exception as e:
-            logger.error(f"Error in check_banned decorator for {func.__name__}: {e}")
-            return await func(client, message, *args, **kwargs)
-    return wrapper
-
-def require_token(func):
-    @wraps(func)
-    async def wrapper(client, message: Message, *args, **kwargs):
-        try:
-            if not message.from_user:
-                return await func(client, message, *args, **kwargs)
-            if not getattr(Var, "TOKEN_ENABLED", False):
-                return await func(client, message, *args, **kwargs)
-            
-            user_id = message.from_user.id
-            is_owner = (isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID) or \
-                       (isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID)
-            
-            if is_owner or await allowed(user_id) or await check(user_id):
-                return await func(client, message, *args, **kwargs)
-            
-            token_data = await generate(user_id)
-            if not token_data or 'token' not in token_data:
-                logger.error(f"Failed to generate token for user {user_id} in require_token.")
-                await message.reply_text("Sorry, could not generate an access token. Please try again later.", quote=True)
-                return
-
-            me = await client.get_me()
-            deep_link = f"https://t.me/{me.username}?start={token_data['token']}"
-            short_url = deep_link
-            try:
-                short_url_result = await shorten(deep_link)
-                if short_url_result:
-                    short_url = short_url_result
-            except Exception as e:
-                logger.warning(f"Failed to shorten token link for user {user_id}: {e}. Using full link.")
-            
-            await message.reply_text(
-                MSG_TOKEN_INVALID,
-                reply_markup=InlineKeyboardMarkup([
-                    [InlineKeyboardButton("Get Token", url=short_url)]
-                ]),
-                quote=True
-            )
-            logger.info(f"Sent token request to user {user_id} for {func.__name__}.")
-            return
-        except Exception as e:
-            logger.error(f"Error in require_token decorator for {func.__name__}: {e}")
-            try:
-                await message.reply_text("An error occurred while checking your authorization. Please try again.", quote=True)
-            except Exception as inner_e:
-                logger.error(f"Failed to send error message to user in require_token: {inner_e}")
-            return
-    return wrapper
-
-def shorten_link(func):
-    @wraps(func)
-    async def wrapper(client, message: Message, *args, **kwargs):
-        try:
-            user_id = message.from_user.id if message.from_user else None
-            use_shortener = getattr(Var, "SHORTEN_MEDIA_LINKS", False)
-            if user_id:
-                try:
-                    is_owner = (isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID) or \
-                               (isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID)
-                    if is_owner or await allowed(user_id):
-                        use_shortener = False
-                except Exception as e:
-                    logger.warning(f"Error checking allowed status for user {user_id} in shorten_link: {e}. Defaulting shortener behavior.")
-            
-            kwargs['shortener'] = use_shortener
-            return await func(client, message, *args, **kwargs)
-        except Exception as e:
-            logger.error(f"Error in shorten_link decorator for {func.__name__}: {e}")
-            kwargs['shortener'] = getattr(Var, "SHORTEN_MEDIA_LINKS", False) 
-            return await func(client, message, *args, **kwargs)
-    return wrapper
-
-def owner_only(func):
-    _cached_owner_ids = None
-    async def wrapper(client, callback_query):
-        nonlocal _cached_owner_ids
-        try:
-            if _cached_owner_ids is None:
-                owner_ids_config = getattr(Var, 'OWNER_ID', [])
-                _cached_owner_ids = [owner_ids_config] if isinstance(owner_ids_config, int) else list(owner_ids_config)
-            
-            if callback_query.from_user.id not in _cached_owner_ids:
-                await callback_query.answer(MSG_ERROR_UNAUTHORIZED, show_alert=True)
-                logger.warning(f"Unauthorized access attempt by {callback_query.from_user.id} to owner_only function {func.__name__}.")
-                return
-            await func(client, callback_query)
-        except Exception as e:
-            logger.error(f"Error in owner_only decorator for {func.__name__}: {e}")
-            try:
-                await callback_query.answer("An error occurred. Please try again.", show_alert=True)
-            except Exception as inner_e:
-                logger.error(f"Failed to send error answer in owner_only: {inner_e}")
-            return
-    return wrapper
-
+# Thunder/utils/decorators.py
+
+from functools import wraps
+from pyrogram.types import Message, InlineKeyboardButton, InlineKeyboardMarkup
+from Thunder.vars import Var
+from Thunder.utils.logger import logger
+from Thunder.utils.messages import (
+    MSG_DECORATOR_BANNED,
+    MSG_TOKEN_INVALID,
+    MSG_ERROR_UNAUTHORIZED
+)
+from Thunder.utils.database import db
+from Thunder.utils.tokens import check, allowed, generate
+from Thunder.utils.shortener import shorten
+
+
+def check_banned(func):
+    @wraps(func)
+    async def wrapper(client, message: Message, *args, **kwargs):
+        try:
+            if not message.from_user:
+                return await func(client, message, *args, **kwargs)
+            user_id = message.from_user.id
+            if isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID:
+                 return await func(client, message, *args, **kwargs)
+            elif isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID:
+                 return await func(client, message, *args, **kwargs)
+
+            ban_details = await db.is_user_banned(user_id)
+            if ban_details:
+                banned_at = ban_details.get('banned_at')
+                ban_time = (
+                    banned_at.strftime('%B %d, %Y, %I:%M %p UTC')
+                    if banned_at and hasattr(banned_at, 'strftime')
+                    else str(banned_at) if banned_at else 'N/A'
+                )
+                await message.reply_text(
+                    MSG_DECORATOR_BANNED.format(
+                        reason=ban_details.get('reason', 'Not specified'),
+                        ban_time=ban_time
+                    ),
+                    quote=True
+                )
+                logger.info(f"Blocked banned user {user_id} from accessing {func.__name__}.")
+                return
+            return await func(client, message, *args, **kwargs)
+        except Exception as e:
+            logger.error(f"Error in check_banned decorator for {func.__name__}: {e}")
+            return await func(client, message, *args, **kwargs)
+    return wrapper
+
+def require_token(func):
+    @wraps(func)
+    async def wrapper(client, message: Message, *args, **kwargs):
+        try:
+            if not message.from_user:
+                return await func(client, message, *args, **kwargs)
+
+            if not getattr(Var, "TOKEN_ENABLED", False):
+                return await func(client, message, *args, **kwargs)
+
+            user_id = message.from_user.id
+            is_owner = (isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID) or \
+                       (isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID)
+
+            if is_owner or await allowed(user_id) or await check(user_id):
+                return await func(client, message, *args, **kwargs)
+
+            temp_token_string = None
+            try:
+                temp_token_string = await generate(user_id)
+            except Exception as e:
+                logger.error(f"Failed to generate temporary token for user {user_id} in require_token: {e}")
+                await message.reply_text("Sorry, could not generate an access token link. Please try again later.", quote=True)
+                return
+
+            if not temp_token_string:
+                logger.error(f"Temporary token generation returned empty for user {user_id} in require_token.")
+                await message.reply_text("Sorry, could not generate an access token link. Please try again later.", quote=True)
+                return
+
+            me = await client.get_me()
+            deep_link = f"https://t.me/{me.username}?start={temp_token_string}"
+            short_url = deep_link
+
+            try:
+                short_url_result = await shorten(deep_link)
+                if short_url_result:
+                    short_url = short_url_result
+            except Exception as e:
+                logger.warning(f"Failed to shorten token link for user {user_id}: {e}. Using full link.")
+
+            await message.reply_text(
+                MSG_TOKEN_INVALID,
+                reply_markup=InlineKeyboardMarkup([
+                    [InlineKeyboardButton("Activate Access", url=short_url)]
+                ]),
+                quote=True
+            )
+            logger.debug(f"Sent temporary token activation link to user {user_id} for {func.__name__}.")
+            return
+        except Exception as e:
+            logger.error(f"Error in require_token decorator for {func.__name__}: {e}")
+            try:
+                await message.reply_text("An error occurred while checking your authorization. Please try again.", quote=True)
+            except Exception as inner_e:
+                logger.error(f"Failed to send error message to user in require_token: {inner_e}")
+            return
+    return wrapper
+
+def shorten_link(func):
+    @wraps(func)
+    async def wrapper(client, message: Message, *args, **kwargs):
+        try:
+            user_id = message.from_user.id if message.from_user else None
+            use_shortener = getattr(Var, "SHORTEN_MEDIA_LINKS", False)
+            if user_id:
+                try:
+                    is_owner = (isinstance(Var.OWNER_ID, int) and user_id == Var.OWNER_ID) or \
+                               (isinstance(Var.OWNER_ID, (list, tuple, set)) and user_id in Var.OWNER_ID)
+                    if is_owner or await allowed(user_id):
+                        use_shortener = False
+                except Exception as e:
+                    logger.warning(f"Error checking allowed status for user {user_id} in shorten_link: {e}. Defaulting shortener behavior.")
+
+            kwargs['shortener'] = use_shortener
+            return await func(client, message, *args, **kwargs)
+        except Exception as e:
+            logger.error(f"Error in shorten_link decorator for {func.__name__}: {e}")
+            kwargs['shortener'] = getattr(Var, "SHORTEN_MEDIA_LINKS", False)
+            return await func(client, message, *args, **kwargs)
+    return wrapper
+
+def owner_only(func):
+    _cached_owner_ids = None
+    @wraps(func)
+    async def wrapper(client, callback_query: Message):
+        nonlocal _cached_owner_ids
+        try:
+            if _cached_owner_ids is None:
+                owner_ids_config = getattr(Var, 'OWNER_ID', [])
+                if isinstance(owner_ids_config, int):
+                    _cached_owner_ids = {owner_ids_config}
+                elif isinstance(owner_ids_config, (list, tuple, set)):
+                    _cached_owner_ids = set(owner_ids_config)
+                else:
+                    _cached_owner_ids = set()
+
+            if callback_query.from_user.id not in _cached_owner_ids:
+                await callback_query.answer(MSG_ERROR_UNAUTHORIZED, show_alert=True)
+                logger.warning(f"Unauthorized access attempt by {callback_query.from_user.id} to owner_only function {func.__name__}.")
+                return
+
+            return await func(client, callback_query)
+        except Exception as e:
+            logger.error(f"Error in owner_only decorator for {func.__name__}: {e}")
+            try:
+                await callback_query.answer("An error occurred. Please try again.", show_alert=True)
+            except Exception as inner_e:
+                logger.error(f"Failed to send error answer in owner_only: {inner_e}")
+            return
+    return wrapper

Thunder/utils/messages.py

@@ -48,6 +48,7 @@ MSG_TOKEN_FAILED = (
     "🔑 Please check your token or contact support."
 )
 MSG_TOKEN_ERROR = "⚙️ **Token Activation Error:** Something went wrong. Please try again."
+MSG_START_INVALID_PAYLOAD = "Invalid command format or expired/invalid link. Please use a valid command or activation link. Error ID: {error_id}"
 MSG_SHELL_ERROR = """**❌ Shell Command Error ❌**
 <pre>{error}</pre>"""
 MSG_SHELL_LARGE_OUTPUT = """Output is too large, sending as a file.
@@ -101,14 +102,9 @@ MSG_DEAUTHORIZE_SUCCESS = (
     "> 👤 User ID: `{user_id}`\n"
     "> 🔒 Access: Revoked"
 )
-MSG_TOKEN_ACTIVATED = (
-    "✅ **Token Activated Successfully!**\n\n"
-    "> 🗓️ Access granted until: **{expiry_date}**\n"
-    "> 📝 Details: _{description}_\n\n"
-    "✨ Enjoy all the features!"
-)
+MSG_TOKEN_ACTIVATED = "✅ Token successfully activated!\n\n⏳ This token is valid for {duration_hours} hours."
 MSG_TOKEN_VERIFIED = "🎉 **Token Verified!** You're all set to use the bot's features."
-MSG_TOKEN_INVALID = "🚫 **Expired or Invalid Token.** Please get a new token to continue."
+MSG_TOKEN_INVALID = "Access to this feature requires an active token. Please click the button below to activate your access token."
 MSG_NO_AUTH_USERS = "ℹ️ **No Authorized Users Found:** The list is currently empty."
 MSG_AUTH_USER_INFO = """{i}. 👤 User ID: `{user_id}`
    • Authorized by: `{authorized_by}`

Thunder/utils/tokens.py

@@ -1,139 +1,112 @@
-import secrets
-from datetime import datetime, timedelta
-from typing import Optional, Dict, Any, List
-from Thunder.utils.database import db
-from Thunder.vars import Var
-from Thunder.utils.error_handling import log_errors
-
-_OWNER_IDS_CACHE = None
-
-def _get_owner_ids():
-    global _OWNER_IDS_CACHE
-    if _OWNER_IDS_CACHE is None:
-        owner_id = Var.OWNER_ID
-        _OWNER_IDS_CACHE = set(owner_id if isinstance(owner_id, (list, tuple, set)) else [owner_id])
-    return _OWNER_IDS_CACHE
-
-@log_errors
-async def check(user_id: int) -> bool:
-    if not getattr(Var, "TOKEN_ENABLED", False):
-        return True
-    if user_id in _get_owner_ids():
-        return True
-    current_time = datetime.utcnow()
-    auth_result = await db.authorized_users_col.find_one(
-        {"user_id": user_id},
-        {"_id": 1}
-    )
-    if auth_result:
-        return True
-    token_result = await db.token_col.find_one(
-        {"user_id": user_id, "expires_at": {"$gt": current_time}},
-        {"_id": 1}
-    )
-    return bool(token_result)
-
-@log_errors
-async def generate(user_id: int) -> Dict[str, Any]:
-    duration = getattr(Var, "TOKEN_TTL_HOURS", 24)
-    token = secrets.token_urlsafe(32)
-    created_at = datetime.utcnow()
-    expires_at = created_at + timedelta(hours=duration)
-    token_data = {
-        "user_id": user_id,
-        "token": token,
-        "created_at": created_at,
-        "expires_at": expires_at
-    }
-    await db.token_col.update_one(
-        {"user_id": user_id},
-        {"$set": token_data},
-        upsert=True
-    )
-    return token_data
-
-@log_errors
-async def allowed(user_id: int) -> bool:
-    result = await db.authorized_users_col.find_one(
-        {"user_id": user_id},
-        {"_id": 1}
-    )
-    return bool(result)
-
-@log_errors
-async def authorize(user_id: int, authorized_by: int) -> bool:
-    auth_data = {
-        "user_id": user_id,
-        "authorized_by": authorized_by,
-        "authorized_at": datetime.utcnow()
-    }
-    await db.authorized_users_col.update_one(
-        {"user_id": user_id},
-        {"$set": auth_data},
-        upsert=True
-    )
-    return True
-
-@log_errors
-async def deauthorize(user_id: int) -> bool:
-    result = await db.authorized_users_col.delete_one({"user_id": user_id})
-    return result.deleted_count > 0
-
-@log_errors
-async def get_user(user_id: int) -> Optional[Dict[str, Any]]:
-    return await db.token_col.find_one({"user_id": user_id})
-
-@log_errors
-async def get(token: str) -> Optional[Dict[str, Any]]:
-    return await db.token_col.find_one({"token": token})
-
-@log_errors
-async def list_allowed() -> List[Dict[str, Any]]:
-    cursor = db.authorized_users_col.find(
-        {},
-        {"user_id": 1, "authorized_by": 1, "authorized_at": 1}
-    )
-    return await cursor.to_list(length=None)
-
-@log_errors
-async def list_tokens() -> List[Dict[str, Any]]:
-    current_time = datetime.utcnow()
-    cursor = db.token_col.find(
-        {"expires_at": {"$gt": current_time}},
-        {"user_id": 1, "expires_at": 1, "created_at": 1}
-    )
-    return await cursor.to_list(length=None)
-
-@log_errors
-async def validate_activation_token(token: str) -> Dict[str, Any]:
-    clean_token = token
-    if token.startswith("token"):
-        clean_token = token[6:] if len(token) > 5 and token[5] in ('-', '_') else token[5:]
-    current_time = datetime.utcnow()
-    token_record = await db.token_col.find_one({"token": clean_token})
-    if not token_record:
-        return {
-            "valid": False,
-            "reason": "Token does not exist or has been revoked",
-            "token": clean_token
-        }
-    if token_record["expires_at"] <= current_time:
-        return {
-            "valid": False,
-            "reason": "Token has expired",
-            "token": clean_token,
-            "expiry_date": token_record["expires_at"].strftime("%Y-%m-%d %H:%M:%S UTC")
-        }
-    return {
-        "valid": True,
-        "token": clean_token,
-        "user_id": token_record["user_id"],
-        "expiry_date": token_record["expires_at"].strftime("%Y-%m-%d %H:%M:%S UTC"),
-        "created_at": token_record["created_at"].strftime("%Y-%m-%d %H:%M:%S UTC"),
-        "description": "Access token"
-    }
-
-@log_errors
-async def cleanup_expired_tokens() -> int:
-    result = await db.token_col.delete_many({"expires_at": {"$lte": datetime.utcnow()}})
-    return result.deleted_count
+import secrets
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Any, List
+from Thunder.utils.database import db
+from Thunder.vars import Var
+from Thunder.utils.error_handling import log_errors
+from Thunder.utils.logger import logger
+
+_OWNER_IDS_CACHE = None
+
+def _get_owner_ids():
+    global _OWNER_IDS_CACHE
+    if _OWNER_IDS_CACHE is None:
+        owner_id = Var.OWNER_ID
+        _OWNER_IDS_CACHE = set(owner_id if isinstance(owner_id, (list, tuple, set)) else [owner_id])
+    return _OWNER_IDS_CACHE
+
+@log_errors
+async def check(user_id: int) -> bool:
+    if not getattr(Var, "TOKEN_ENABLED", False):
+        return True
+    if user_id in _get_owner_ids():
+        return True
+    current_time = datetime.utcnow()
+    auth_result = await db.authorized_users_col.find_one(
+        {"user_id": user_id},
+        {"_id": 1}
+    )
+    if auth_result:
+        return True
+    token_result = await db.token_col.find_one(
+        {"user_id": user_id, "expires_at": {"$gt": current_time}, "activated": True},
+        {"_id": 1}
+    )
+    return bool(token_result)
+
+@log_errors
+async def generate(user_id: int) -> str:
+    token_str = secrets.token_urlsafe(32)
+    ttl_hours = getattr(Var, "TOKEN_TTL_HOURS", 24)
+    created_at = datetime.utcnow()
+    expires_at = created_at + timedelta(hours=ttl_hours)
+
+    try:
+        await db.save_main_token(
+            user_id=user_id,
+            token_value=token_str,
+            created_at=created_at,
+            expires_at=expires_at,
+            activated=False
+        )
+        return token_str
+    except Exception as e:
+        logger.error(f"Failed to generate and save unactivated token for user {user_id}: {e}")
+        raise
+
+@log_errors
+async def allowed(user_id: int) -> bool:
+    result = await db.authorized_users_col.find_one(
+        {"user_id": user_id},
+        {"_id": 1}
+    )
+    return bool(result)
+
+@log_errors
+async def authorize(user_id: int, authorized_by: int) -> bool:
+    auth_data = {
+        "user_id": user_id,
+        "authorized_by": authorized_by,
+        "authorized_at": datetime.utcnow()
+    }
+    await db.authorized_users_col.update_one(
+        {"user_id": user_id},
+        {"$set": auth_data},
+        upsert=True
+    )
+    return True
+
+@log_errors
+async def deauthorize(user_id: int) -> bool:
+    result = await db.authorized_users_col.delete_one({"user_id": user_id})
+    return result.deleted_count > 0
+
+@log_errors
+async def get_user(user_id: int) -> Optional[Dict[str, Any]]:
+    return await db.token_col.find_one({"user_id": user_id})
+
+@log_errors
+async def get(token: str) -> Optional[Dict[str, Any]]:
+    return await db.token_col.find_one({"token": token})
+
+@log_errors
+async def list_allowed() -> List[Dict[str, Any]]:
+    cursor = db.authorized_users_col.find(
+        {},
+        {"user_id": 1, "authorized_by": 1, "authorized_at": 1}
+    )
+    return await cursor.to_list(length=None)
+
+@log_errors
+async def list_tokens() -> List[Dict[str, Any]]:
+    current_time = datetime.utcnow()
+    cursor = db.token_col.find(
+        {"expires_at": {"$gt": current_time}},
+        {"user_id": 1, "expires_at": 1, "created_at": 1, "activated": 1}
+    )
+    return await cursor.to_list(length=None)
+
+@log_errors
+async def cleanup_expired_tokens() -> int:
+    result = await db.token_col.delete_many({"expires_at": {"$lte": datetime.utcnow()}})
+    return result.deleted_count