visual_enhancer.py

"""
Visual Enhancements - Beautiful messages, animations, and progress indicators
"""
import asyncio
from typing import Optional
from telegram import Update, Message
from telegram.constants import ChatAction
from loguru import logger

class VisualEnhancer:
    """Enhance bot messages with emojis, formatting, and animations"""
    
    @staticmethod
    def escape_markdown(text: str) -> str:
        """Escape markdown special characters for Telegram"""
        # Escape special markdown characters
        special_chars = ['_', '*', '[', ']', '(', ')', '~', '`', '>', '#', '+', '-', '=', '|', '{', '}', '.', '!']
        for char in special_chars:
            text = text.replace(char, f'\\{char}')
        return text
    
    # Emoji sets
    EMOJIS = {
        'loading': ['⏳', '⌛'],
        'progress': ['▱▱▱▱▱', '▰▱▱▱▱', '▰▰▱▱▱', '▰▰▰▱▱', '▰▰▰▰▱', '▰▰▰▰▰'],
        'scan': '🔍',
        'success': '✅',
        'error': '❌',
        'warning': '⚠️',
        'info': 'ℹ️',
        'rocket': '🚀',
        'fire': '🔥',
        'star': '⭐',
        'check': '✓',
        'cross': '✗',
        'arrow': '→',
        'bullet': '•',
        'shield': '🛡️',
        'lock': '🔒',
        'globe': '🌐',
        'speed': '⚡',
        'chart': '📊',
        'doc': '📄',
        'folder': '📁',
        'link': '🔗',
        'tools': '🔧',
        'clock': '⏰',
        'calendar': '📅',
        'location': '📍',
        'server': '🖥️',
        'database': '🗄️',
        'download': '⬇️',
        'upload': '⬆️',
        'heart': '❤️',
        'trophy': '🏆',
        'crown': '👑',
        'gem': '💎',
        'magic': '✨',
        'party': '🎉',
        'target': '🎯',
        'package': '📦',
        'bug': '🐛',
        'robot': '🤖'
    }
    
    @staticmethod
    def format_header(text: str, emoji: str = '🔍') -> str:
        """Format header with emojis"""
        line = '━' * 30
        return f"{line}\n{emoji} {text} {emoji}\n{line}"
    
    @staticmethod
    def format_section(title: str, emoji: str = '📊') -> str:
        """Format section title"""
        return f"\n{emoji} {title}\n{'─' * 25}"
    
    @staticmethod
    def format_item(key: str, value: str, emoji: str = '•') -> str:
        """Format key-value item"""
        return f"{emoji} {key}: {value}"
    
    @staticmethod
    def format_list(items: list, emoji: str = '▪️') -> str:
        """Format list items"""
        return '\n'.join([f"{emoji} {item}" for item in items])
    
    @staticmethod
    def create_progress_bar(current: int, total: int, width: int = 10) -> str:
        """Create visual progress bar"""
        if total == 0:
            return "▱" * width
        
        filled = int((current / total) * width)
        bar = "▰" * filled + "▱" * (width - filled)
        percentage = int((current / total) * 100)
        
        return f"{bar} {percentage}%"
    
    @staticmethod
    async def send_with_animation(message: Message, text: str, animation_type: str = 'typing'):
        """Send message with animation"""
        try:
            # Show action
            if animation_type == 'typing':
                await message.chat.send_action(ChatAction.TYPING)
            elif animation_type == 'upload':
                await message.chat.send_action(ChatAction.UPLOAD_DOCUMENT)
            
            await asyncio.sleep(0.5)
            
            # Send message
            return await message.reply_text(text, parse_mode='Markdown')
        except Exception as e:
            logger.error(f"Animation error: {e}")
            return await message.reply_text(text)
    
    @staticmethod
    async def show_progress(message: Message, current: int, total: int, item_name: str = "item"):
        """Update message with progress"""
        try:
            progress_bar = VisualEnhancer.create_progress_bar(current, total)
            text = f"🔄 Processing...\n\n{progress_bar}\n\n{current}/{total} {item_name}s completed"
            
            await message.edit_text(text)
        except Exception as e:
            logger.error(f"Progress update error: {e}")
    
    @staticmethod
    def format_scan_result(result: dict) -> str:
        """Format scan result beautifully"""
        try:
            lines = []
            
            # Header
            lines.append(VisualEnhancer.format_header("Scan Complete", "🎯"))
            lines.append("")
            
            # URL
            if 'url' in result:
                lines.append(f"🔗 URL: {result['url']}")
                lines.append("")
            
            # Status
            if 'status' in result:
                status_emoji = '✅' if result['status'] == 'success' else '❌'
                lines.append(f"{status_emoji} Status: {result['status']}")
                lines.append("")
            
            # Summary stats
            if 'summary' in result:
                lines.append(VisualEnhancer.format_section("Summary", "📊"))
                
                for key, value in result['summary'].items():
                    emoji = VisualEnhancer._get_summary_emoji(key)
                    lines.append(VisualEnhancer.format_item(key, str(value), emoji))
                
                lines.append("")
            
            # Scanner count
            if 'scanners' in result:
                scanner_count = len(result['scanners'])
                lines.append(f"🔍 Scanners Used: {scanner_count}")
                lines.append("")
            
            lines.append("━" * 30)
            
            return '\n'.join(lines)
        except Exception as e:
            logger.error(f"Format error: {e}")
            return "✅ Scan completed successfully"
    
    @staticmethod
    def format_batch_result(results: list, total_time: float) -> str:
        """Format batch scan results"""
        try:
            lines = []
            
            lines.append(VisualEnhancer.format_header("Batch Scan Complete", "🎉"))
            lines.append("")
            
            successful = sum(1 for r in results if r.get('status') == 'success')
            failed = len(results) - successful
            
            lines.append(f"📊 Total Scans: {len(results)}")
            lines.append(f"✅ Successful: {successful}")
            lines.append(f"❌ Failed: {failed}")
            lines.append(f"⏱️ Time: {total_time:.1f}s")
            lines.append("")
            
            # List URLs
            lines.append(VisualEnhancer.format_section("Scanned URLs", "🔗"))
            for i, result in enumerate(results[:10], 1):  # Show first 10
                status = "✓" if result.get('status') == 'success' else "✗"
                url = result.get('url', 'Unknown')
                safe_url = VisualEnhancer.escape_markdown(url[:50])
                lines.append(f"{i}. {status} {safe_url}...")
            
            if len(results) > 10:
                lines.append(f"\n... and {len(results) - 10} more")
            
            lines.append("")
            lines.append("━" * 30)
            
            return '\n'.join(lines)
        except Exception as e:
            logger.error(f"Format batch error: {e}")
            return f"✅ Batch scan completed: {len(results)} URLs"
    
    @staticmethod
    def format_deep_scan_result(result: dict) -> str:
        """Format deep scan results"""
        try:
            lines = []
            
            lines.append(VisualEnhancer.format_header("Deep Scan Complete", "🔬"))
            lines.append("")
            
            if 'base_url' in result:
                lines.append(f"🌐 Base URL: {result['base_url']}")
            
            if 'pages_scanned' in result:
                lines.append(f"📄 Pages Scanned: {result['pages_scanned']}")
            
            if 'max_depth_reached' in result:
                lines.append(f"📊 Max Depth: {result['max_depth_reached']}")
            
            lines.append("")
            
            # Visited URLs
            if 'visited_urls' in result:
                lines.append(VisualEnhancer.format_section("Visited Pages", "🔗"))
                for i, url in enumerate(result['visited_urls'][:8], 1):
                    safe_url = VisualEnhancer.escape_markdown(url[:60])
                    lines.append(f"{i}. {safe_url}...")
                
                if len(result['visited_urls']) > 8:
                    lines.append(f"\n... and {len(result['visited_urls']) - 8} more")
            
            lines.append("")
            lines.append("━" * 30)
            
            return '\n'.join(lines)
        except Exception as e:
            logger.error(f"Format deep scan error: {e}")
            return "✅ Deep scan completed"
    
    @staticmethod
    def format_queue_status(position: int, total: int, estimated_time: int) -> str:
        """Format queue status message"""
        try:
            lines = []
            
            lines.append(VisualEnhancer.format_header("Added to Queue", "⏳"))
            lines.append("")
            
            lines.append(f"📍 Your Position: #{position}")
            lines.append(f"👥 Queue Size: {total}")
            lines.append(f"⏱️ Estimated Wait: ~{estimated_time}s")
            lines.append("")
            
            # Progress visualization
            if total > 0:
                progress = VisualEnhancer.create_progress_bar(total - position, total)
                lines.append(f"Queue: {progress}")
            
            lines.append("")
            lines.append("💡 You'll be notified when your scan starts!")
            lines.append("━" * 30)
            
            return '\n'.join(lines)
        except Exception as e:
            logger.error(f"Format queue error: {e}")
            return f"⏳ Added to queue (position: {position})"
    
    @staticmethod
    def _get_summary_emoji(key: str) -> str:
        """Get appropriate emoji for summary key"""
        key_lower = key.lower()
        
        if 'time' in key_lower or 'duration' in key_lower:
            return '⏱️'
        elif 'success' in key_lower:
            return '✅'
        elif 'fail' in key_lower or 'error' in key_lower:
            return '❌'
        elif 'security' in key_lower:
            return '🛡️'
        elif 'speed' in key_lower or 'performance' in key_lower:
            return '⚡'
        elif 'size' in key_lower:
            return '📦'
        elif 'redirect' in key_lower:
            return '↪️'
        elif 'status' in key_lower:
            return '📊'
        else:
            return '▪️'
    
    @staticmethod
    async def animate_scanning(message: Message, url: str):
        """Animate scanning process"""
        animations = [
            f"🔍 Scanning {url}\n▱▱▱▱▱▱▱▱▱▱ 0%",
            f"🔍 Scanning {url}\n▰▰▱▱▱▱▱▱▱▱ 20%",
            f"🔍 Scanning {url}\n▰▰▰▰▱▱▱▱▱▱ 40%",
            f"🔍 Scanning {url}\n▰▰▰▰▰▰▱▱▱▱ 60%",
            f"🔍 Scanning {url}\n▰▰▰▰▰▰▰▰▱▱ 80%",
            f"🔍 Scanning {url}\n▰▰▰▰▰▰▰▰▰▰ 100%",
        ]
        
        for anim in animations:
            try:
                await message.edit_text(anim)
                await asyncio.sleep(0.3)
            except Exception:
                break