Add source

bot/modules/media_loader/__init__.py

@@ -0,0 +1,4 @@
+"""
+Media loader module
+"""
+

bot/modules/media_loader/direct.py

@@ -0,0 +1,68 @@
+"""
+Direct link downloads
+"""
+import aiohttp
+import aiofiles
+from pathlib import Path
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+async def download_file(url: str, output_path: str, chunk_size: int = 8192) -> bool:
+    """
+    Download file from direct link
+    
+    Args:
+        url: File URL
+        output_path: Path to save
+        chunk_size: Chunk size for download
+    
+    Returns:
+        True if successful, False otherwise
+    """
+    try:
+        # Create directory if it doesn't exist
+        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+        
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                if response.status != 200:
+                    logger.error(f"Download error: HTTP {response.status}")
+                    return False
+                
+                async with aiofiles.open(output_path, 'wb') as f:
+                    async for chunk in response.content.iter_chunked(chunk_size):
+                        await f.write(chunk)
+        
+        logger.info(f"File downloaded: {output_path}")
+        return True
+    
+    except Exception as e:
+        logger.error(f"Error downloading file: {e}")
+        return False
+
+
+async def get_file_size(url: str) -> Optional[int]:
+    """
+    Get file size from URL
+    
+    Args:
+        url: File URL
+    
+    Returns:
+        File size in bytes or None
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.head(url) as response:
+                if response.status == 200:
+                    content_length = response.headers.get('Content-Length')
+                    if content_length:
+                        return int(content_length)
+    except Exception as e:
+        logger.error(f"Error getting file size: {e}")
+    
+    return None
+

bot/modules/media_loader/sender.py

@@ -0,0 +1,270 @@
+"""
+Sending files to users
+"""
+from pathlib import Path
+from pyrogram import Client
+from pyrogram.types import Message
+from typing import Optional
+import aiohttp
+import aiofiles
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+async def download_thumbnail(url: str, output_path: str) -> Optional[str]:
+    """
+    Download thumbnail from URL
+    
+    Args:
+        url: Thumbnail URL
+        output_path: Path to save
+    
+    Returns:
+        Path to downloaded file or None
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                if response.status == 200:
+                    async with aiofiles.open(output_path, 'wb') as f:
+                        async for chunk in response.content.iter_chunked(8192):
+                            await f.write(chunk)
+                    return output_path
+    except Exception as e:
+        logger.warning(f"Failed to download thumbnail: {e}")
+        return None
+
+
+async def send_file_to_user(
+    client: Client,
+    chat_id: int,
+    file_path: str,
+    caption: Optional[str] = None,
+    thumbnail: Optional[str] = None
+) -> bool:
+    """
+    Send file to user
+    
+    Args:
+        client: Pyrogram client
+        chat_id: Chat ID
+        file_path: Path to file
+        caption: File caption
+        thumbnail: Path to thumbnail or URL
+    
+    Returns:
+        True if successful, False otherwise
+    """
+    thumbnail_path = None
+    try:
+        file = Path(file_path)
+        
+        if not file.exists():
+            logger.error(f"File not found: {file_path}")
+            return False
+        
+        # Maximum file size for Telegram (2GB)
+        max_size = 2 * 1024 * 1024 * 1024
+        file_size = file.stat().st_size
+        
+        # If file is larger than 2GB, split into parts
+        if file_size > max_size:
+            logger.info(f"File too large ({file_size / (1024*1024*1024):.2f} GB), splitting into parts...")
+            return await send_large_file_in_parts(
+                client, chat_id, file_path, caption, thumbnail
+            )
+        
+        # Process thumbnail (can be URL or file path)
+        if thumbnail:
+            if thumbnail.startswith(('http://', 'https://')):
+                # This is a URL - download thumbnail
+                thumbnail_path = f"downloads/thumb_{file.stem}.jpg"
+                downloaded = await download_thumbnail(thumbnail, thumbnail_path)
+                if downloaded:
+                    thumbnail_path = downloaded
+                else:
+                    thumbnail_path = None  # Don't use thumbnail if download failed
+            else:
+                # This is a file path
+                thumb_file = Path(thumbnail)
+                if thumb_file.exists():
+                    thumbnail_path = thumbnail
+                else:
+                    thumbnail_path = None
+        
+        # Determine file type
+        if file.suffix.lower() in ['.mp4', '.avi', '.mov', '.mkv', '.webm']:
+            # Video - if no thumbnail, try to generate one
+            if not thumbnail_path:
+                from bot.utils.file_processor import generate_thumbnail
+                thumbnail_path_temp = f"downloads/thumb_{file.stem}.jpg"
+                if await generate_thumbnail(str(file), thumbnail_path_temp):
+                    thumbnail_path = thumbnail_path_temp
+            
+            await client.send_video(
+                chat_id=chat_id,
+                video=str(file),
+                caption=caption,
+                thumb=thumbnail_path
+            )
+        elif file.suffix.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.webp']:
+            # Image
+            await client.send_photo(
+                chat_id=chat_id,
+                photo=str(file),
+                caption=caption
+            )
+        elif file.suffix.lower() in ['.mp3', '.wav', '.ogg', '.m4a', '.flac']:
+            # Audio
+            await client.send_audio(
+                chat_id=chat_id,
+                audio=str(file),
+                caption=caption,
+                thumb=thumbnail_path
+            )
+        else:
+            # Document
+            await client.send_document(
+                chat_id=chat_id,
+                document=str(file),
+                caption=caption,
+                thumb=thumbnail_path
+            )
+        
+        logger.info(f"File sent to user {chat_id}: {file_path}")
+        return True
+    
+    except Exception as e:
+        logger.error(f"Error sending file: {e}", exc_info=True)
+        return False
+    
+    finally:
+        # Delete temporary thumbnail if it was downloaded
+        if thumbnail_path and thumbnail_path.startswith("downloads/thumb_"):
+            try:
+                thumb_file = Path(thumbnail_path)
+                if thumb_file.exists():
+                    thumb_file.unlink()
+                    logger.debug(f"Temporary thumbnail deleted: {thumbnail_path}")
+            except Exception as e:
+                logger.warning(f"Failed to delete temporary thumbnail: {e}")
+
+
+async def send_large_file_in_parts(
+    client: Client,
+    chat_id: int,
+    file_path: str,
+    caption: Optional[str] = None,
+    thumbnail: Optional[str] = None
+) -> bool:
+    """
+    Send large file in parts
+    
+    Args:
+        client: Pyrogram client
+        chat_id: Chat ID
+        file_path: Path to file
+        caption: File caption
+        thumbnail: Path to thumbnail or URL
+    
+    Returns:
+        True if successful, False otherwise
+    """
+    from bot.utils.file_splitter import split_file, delete_file_parts, get_part_info
+    
+    parts = []
+    try:
+        # Split file into parts
+        parts = await split_file(file_path)
+        part_info = get_part_info(parts)
+        total_parts = part_info["total_parts"]
+        
+        logger.info(f"Sending file in parts: {total_parts} parts")
+        
+        # Send each part
+        for part_num, part_path in enumerate(parts, 1):
+            part_caption = None
+            if caption:
+                part_caption = f"{caption}\n\n📦 Part {part_num} of {total_parts}"
+            else:
+                part_caption = f"📦 Part {part_num} of {total_parts}"
+            
+            # Send thumbnail only with first part
+            part_thumbnail = thumbnail if part_num == 1 else None
+            
+            try:
+                await client.send_document(
+                    chat_id=chat_id,
+                    document=str(part_path),
+                    caption=part_caption,
+                    thumb=part_thumbnail
+                )
+                logger.info(f"Sent part {part_num}/{total_parts}")
+            except Exception as e:
+                logger.error(f"Error sending part {part_num}: {e}", exc_info=True)
+                # Continue sending other parts
+                continue
+        
+        logger.info(f"File sent in parts to user {chat_id}")
+        return True
+    
+    except Exception as e:
+        logger.error(f"Error sending large file in parts: {e}", exc_info=True)
+        return False
+    
+    finally:
+        # Delete file parts after sending
+        if parts:
+            await delete_file_parts(parts)
+            logger.debug("File parts deleted")
+
+
+async def delete_file(file_path: str, max_retries: int = 3) -> bool:
+    """
+    Delete file with retries
+    
+    Args:
+        file_path: Path to file
+        max_retries: Maximum number of retries
+    
+    Returns:
+        True if successful
+    """
+    import asyncio
+    
+    file = Path(file_path)
+    if not file.exists():
+        return True  # File already doesn't exist
+    
+    for attempt in range(max_retries):
+        try:
+            file.unlink()
+            logger.info(f"File deleted: {file_path}")
+            return True
+        except PermissionError as e:
+            # File may be locked by another process
+            if attempt < max_retries - 1:
+                wait_time = (attempt + 1) * 0.5  # Exponential backoff
+                logger.warning(f"File locked, retrying in {wait_time}s: {file_path}")
+                await asyncio.sleep(wait_time)
+            else:
+                logger.error(f"Failed to delete file after {max_retries} attempts (locked): {file_path}")
+                # Add to cleanup queue for background cleanup task
+                from bot.utils.file_cleanup import add_file_to_cleanup_queue
+                add_file_to_cleanup_queue(str(file_path))
+                return False
+        except Exception as e:
+            if attempt < max_retries - 1:
+                wait_time = (attempt + 1) * 0.5
+                logger.warning(f"Error deleting file, retrying in {wait_time}s: {e}")
+                await asyncio.sleep(wait_time)
+            else:
+                logger.error(f"Error deleting file after {max_retries} attempts: {e}", exc_info=True)
+                # Add to cleanup queue for background cleanup task
+                from bot.utils.file_cleanup import add_file_to_cleanup_queue
+                add_file_to_cleanup_queue(str(file_path))
+                return False
+    
+    return False
+

bot/modules/media_loader/ytdlp.py

@@ -0,0 +1,358 @@
+"""
+Downloads via yt-dlp
+"""
+import yt_dlp
+from pathlib import Path
+from typing import Optional, Dict, Callable
+import asyncio
+import threading
+import logging
+import time
+
+logger = logging.getLogger(__name__)
+
+
+def create_progress_hook(progress_callback: Optional[Callable] = None, event_loop=None, cancel_event: Optional[threading.Event] = None, last_update_time: list = None):
+    """
+    Create progress hook for tracking download progress
+    
+    Args:
+        progress_callback: Async callback function for updating progress
+        event_loop: Event loop from main thread (for calling from executor)
+        cancel_event: Event for checking download cancellation
+        last_update_time: List to store last update time (for rate limiting)
+    
+    Returns:
+        Hook function for yt-dlp
+    """
+    if last_update_time is None:
+        last_update_time = [0]
+    
+    def progress_hook(d: dict):
+        # Check for cancellation
+        if cancel_event and cancel_event.is_set():
+            raise KeyboardInterrupt("Download cancelled")
+        
+        if d.get('status') == 'downloading':
+            percent = 0
+            if 'total_bytes' in d and d['total_bytes']:
+                percent = (d.get('downloaded_bytes', 0) / d['total_bytes']) * 100
+            elif 'total_bytes_estimate' in d and d['total_bytes_estimate']:
+                percent = (d.get('downloaded_bytes', 0) / d['total_bytes_estimate']) * 100
+            
+            # Limit update frequency (no more than once per second)
+            current_time = time.time()
+            if progress_callback and percent > 0 and event_loop and (current_time - last_update_time[0] >= 1.0):
+                try:
+                    last_update_time[0] = current_time
+                    # Use provided event loop for safe call from another thread
+                    # run_coroutine_threadsafe doesn't block current thread and doesn't block event loop
+                    future = asyncio.run_coroutine_threadsafe(
+                        progress_callback(int(percent)), 
+                        event_loop
+                    )
+                    # Don't wait for completion (future.result()) to avoid blocking download
+                except Exception as e:
+                    logger.debug(f"Error updating progress: {e}")
+    
+    return progress_hook
+
+
+async def download_media(
+    url: str, 
+    output_dir: str = "downloads", 
+    quality: str = "best",
+    progress_callback: Optional[Callable] = None,
+    cookies_file: Optional[str] = None,
+    cancel_event: Optional[threading.Event] = None,
+    task_id: Optional[int] = None
+) -> Optional[Dict]:
+    """
+    Download media via yt-dlp
+    
+    Args:
+        url: Video/media URL
+        output_dir: Directory for saving
+        quality: Video quality (best, worst, 720p, etc.)
+        progress_callback: Function for updating progress (accepts int 0-100)
+        cookies_file: Path to cookies file (optional)
+        cancel_event: Event for cancellation check (optional)
+        task_id: Task ID for unique file naming (optional)
+    
+    Returns:
+        Dictionary with downloaded file information or None
+    """
+    try:
+        # URL validation
+        from bot.utils.helpers import is_valid_url
+        if not is_valid_url(url):
+            logger.error(f"Invalid or unsafe URL: {url}")
+            return None
+        
+        # Create directory
+        Path(output_dir).mkdir(parents=True, exist_ok=True)
+        
+        # Check free disk space (minimum 1GB)
+        import shutil
+        try:
+            disk_usage = shutil.disk_usage(output_dir)
+            free_space_gb = disk_usage.free / (1024 ** 3)
+            min_free_space_gb = 1.0  # Minimum 1GB free space
+            if free_space_gb < min_free_space_gb:
+                logger.error(f"Insufficient free disk space: {free_space_gb:.2f} GB (minimum {min_free_space_gb} GB required)")
+                return None
+        except Exception as e:
+            logger.warning(f"Failed to check free disk space: {e}")
+        
+        # Get event loop BEFORE starting executor to pass it to progress hook
+        # Use get_running_loop() for explicit check that we're in async context
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # If no running loop, try to get current one (for backward compatibility)
+            loop = asyncio.get_event_loop()
+        
+        # List to store last progress update time
+        last_update_time = [0]
+        
+        # Configure yt-dlp with progress hook that uses correct event loop
+        progress_hook_func = create_progress_hook(
+            progress_callback, 
+            event_loop=loop,
+            cancel_event=cancel_event,
+            last_update_time=last_update_time
+        )
+        
+        # Form unique filename with task_id to prevent conflicts
+        if task_id:
+            outtmpl = str(Path(output_dir) / f'%(title)s_[task_{task_id}].%(ext)s')
+        else:
+            outtmpl = str(Path(output_dir) / '%(title)s.%(ext)s')
+        
+        # Configure format selector for maximum quality with correct aspect ratio
+        # Priority: best video + best audio, or best combined format
+        # This ensures we get the highest quality available while maintaining original proportions
+        if quality == "best":
+            # Format selector for maximum quality:
+            # 1. bestvideo (highest quality video) + bestaudio (highest quality audio) - best quality
+            # 2. best (best combined format if separate streams not available) - fallback
+            # This selector maintains original aspect ratio and resolution
+            format_selector = 'bestvideo+bestaudio/best'
+        else:
+            # Use custom quality if specified
+            format_selector = quality
+        
+        ydl_opts = {
+            'format': format_selector,
+            'outtmpl': outtmpl,
+            'quiet': False,
+            'no_warnings': False,
+            'progress_hooks': [progress_hook_func],
+            # Merge video and audio into single file (if separate streams)
+            'merge_output_format': 'mp4',
+            # Don't prefer free formats (they may be lower quality)
+            'prefer_free_formats': False,
+            # Additional options for better quality
+            'writesubtitles': False,
+            'writeautomaticsub': False,
+            'ignoreerrors': False,
+        }
+        
+        # Add cookies if specified (for Instagram and other sites)
+        if cookies_file:
+            # Resolve cookies file path (support relative and absolute paths)
+            cookies_path = Path(cookies_file)
+            if not cookies_path.is_absolute():
+                # If path is relative, search relative to project root
+                project_root = Path(__file__).parent.parent.parent.parent
+                cookies_path = project_root / cookies_file
+                # Also check current working directory
+                if not cookies_path.exists():
+                    cookies_path = Path(cookies_file).resolve()
+            
+            if cookies_path.exists():
+                ydl_opts['cookiefile'] = str(cookies_path)
+                logger.info(f"Using cookies from file: {cookies_path}")
+            else:
+                logger.warning(
+                    f"Cookies file not found: {cookies_path} (original path: {cookies_file}). "
+                    f"Continuing without cookies. Check COOKIES_FILE path in configuration."
+                )
+        
+        def run_download():
+            """Synchronous function to execute in separate thread"""
+            # This function runs in a separate thread (ThreadPoolExecutor)
+            # progress hook will be called from this thread and use
+            # run_coroutine_threadsafe for safe call in main event loop
+            try:
+                with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                    # Check for cancellation before start
+                    if cancel_event and cancel_event.is_set():
+                        raise KeyboardInterrupt("Download cancelled")
+                    
+                    # Get video information
+                    info = ydl.extract_info(url, download=False)
+                    
+                    # Check for cancellation after getting info
+                    if cancel_event and cancel_event.is_set():
+                        raise KeyboardInterrupt("Download cancelled")
+                    
+                    # Download (progress hook will be called from this thread)
+                    ydl.download([url])
+                    
+                    return info
+            except KeyboardInterrupt:
+                # Interrupt download on cancellation
+                logger.info("Download interrupted")
+                raise
+        
+        # Execute in executor for non-blocking download
+        # None uses ThreadPoolExecutor by default
+        # This ensures download doesn't block message processing
+        # Event loop continues processing messages in parallel with download
+        info = await loop.run_in_executor(None, run_download)
+        
+        # Search for downloaded file
+        title = info.get('title', 'video')
+        # Clean title from invalid characters
+        title = "".join(c for c in title if c.isalnum() or c in (' ', '-', '_')).strip()
+        ext = info.get('ext', 'mp4')
+        
+        logger.info(f"Searching for downloaded file. Title: {title}, ext: {ext}, task_id: {task_id}")
+        
+        # Form filename with task_id
+        if task_id:
+            filename = f"{title}_[task_{task_id}].{ext}"
+        else:
+            filename = f"{title}.{ext}"
+        
+        file_path = Path(output_dir) / filename
+        logger.debug(f"Expected file path: {file_path}")
+        
+        # If file not found at expected path, search in directory
+        if not file_path.exists():
+            logger.info(f"File not found at expected path {file_path}, starting search...")
+            
+            # If task_id exists, search for file with this task_id
+            if task_id:
+                # Pattern 1: exact match with task_id
+                pattern = f"*[task_{task_id}].{ext}"
+                files = list(Path(output_dir).glob(pattern))
+                logger.debug(f"Search by pattern '{pattern}': found {len(files)} files")
+                
+                if not files:
+                    # Pattern 2: search files containing task_id (in case format differs slightly)
+                    pattern2 = f"*task_{task_id}*.{ext}"
+                    files = list(Path(output_dir).glob(pattern2))
+                    logger.debug(f"Search by pattern '{pattern2}': found {len(files)} files")
+                
+                if files:
+                    # Take newest file from found ones
+                    file_path = max(files, key=lambda p: p.stat().st_mtime)
+                    logger.info(f"Found file by task_id: {file_path}")
+                else:
+                    # If not found by task_id, search newest file with this extension
+                    logger.info(f"File with task_id {task_id} not found, searching newest .{ext} file")
+                    files = list(Path(output_dir).glob(f"*.{ext}"))
+                    if files:
+                        # Filter files created recently (last 5 minutes)
+                        import time
+                        current_time = time.time()
+                        recent_files = [
+                            f for f in files 
+                            if (current_time - f.stat().st_mtime) < 300  # 5 minutes
+                        ]
+                        if recent_files:
+                            file_path = max(recent_files, key=lambda p: p.stat().st_mtime)
+                            logger.info(f"Found recently created file: {file_path}")
+                        else:
+                            file_path = max(files, key=lambda p: p.stat().st_mtime)
+                            logger.warning(f"No recent files found, taking newest: {file_path}")
+            else:
+                # Search file by extension
+                files = list(Path(output_dir).glob(f"*.{ext}"))
+                if files:
+                    # Take newest file
+                    file_path = max(files, key=lambda p: p.stat().st_mtime)
+                    logger.info(f"Found file by time: {file_path}")
+        
+        if file_path.exists():
+            file_size = file_path.stat().st_size
+            logger.info(f"File found: {file_path}, size: {file_size / (1024*1024):.2f} MB")
+            return {
+                'file_path': str(file_path),
+                'title': title,
+                'duration': info.get('duration'),
+                'thumbnail': info.get('thumbnail'),
+                'size': file_size
+            }
+        else:
+            # Output list of all files in directory for debugging
+            all_files = list(Path(output_dir).glob("*"))
+            logger.error(
+                f"File not found after download: {file_path}\n"
+                f"Files in downloads directory: {[str(f.name) for f in all_files[:10]]}"
+            )
+            return None
+    
+    except Exception as e:
+        logger.error(f"Error downloading via yt-dlp: {e}", exc_info=True)
+        return None
+
+
+async def get_media_info(url: str, cookies_file: Optional[str] = None) -> Optional[Dict]:
+    """
+    Get media information without downloading
+    
+    Args:
+        url: Media URL
+        cookies_file: Path to cookies file (optional)
+    
+    Returns:
+        Dictionary with information or None
+    """
+    try:
+        loop = asyncio.get_running_loop()
+        ydl_opts = {
+            'quiet': True,
+            'no_warnings': True,
+        }
+        
+        # Add cookies if specified
+        if cookies_file:
+            # Resolve cookies file path (support relative and absolute paths)
+            cookies_path = Path(cookies_file)
+            if not cookies_path.is_absolute():
+                # If path is relative, search relative to project root
+                project_root = Path(__file__).parent.parent.parent.parent
+                cookies_path = project_root / cookies_file
+                # Also check current working directory
+                if not cookies_path.exists():
+                    cookies_path = Path(cookies_file).resolve()
+            
+            if cookies_path.exists():
+                ydl_opts['cookiefile'] = str(cookies_path)
+                logger.debug(f"Using cookies to get info: {cookies_path}")
+            else:
+                logger.warning(f"Cookies file not found for get_media_info: {cookies_path}")
+        
+        def extract_info_sync():
+            """Synchronous function for extracting information"""
+            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                return ydl.extract_info(url, download=False)
+        
+        # Run synchronous yt-dlp in executor to avoid blocking event loop
+        info = await loop.run_in_executor(None, extract_info_sync)
+        
+        return {
+            'title': info.get('title'),
+            'duration': info.get('duration'),
+            'thumbnail': info.get('thumbnail'),
+            'uploader': info.get('uploader'),
+            'view_count': info.get('view_count'),
+        }
+    
+    except Exception as e:
+        logger.error(f"Error getting media info: {e}", exc_info=True)
+        return None
+