security_utils.py

"""Security utilities - Input validation, sanitization, and protection"""
import re
from typing import Optional
from urllib.parse import urlparse

from loguru import logger


class InputValidator:
    """Validate and sanitize user input"""
    
    # Dangerous patterns that could indicate injection attempts
    DANGEROUS_PATTERNS = [
        r'javascript:',
        r'<script',
        r'onerror=',
        r'onload=',
        r'eval\(',
        r'expression\(',
        r'vbscript:',
        r'data:text/html',
    ]
    
    # Valid URL schemes (expanded for multiple protocols)
    VALID_SCHEMES = [
        'http', 'https',      # Web
        'ftp', 'ftps',        # File transfer
        'ws', 'wss',          # WebSocket
        'rtsp', 'rtmp',       # Streaming
        'ssh', 'sftp',        # Secure protocols
    ]
    
    # Maximum URL length
    MAX_URL_LENGTH = 2048
    
    @classmethod
    def validate_url(cls, url: str) -> tuple[bool, Optional[str]]:
        """
        Validate URL for security and format
        
        Returns:
            (is_valid, error_message)
        """
        if not url:
            return False, "URL cannot be empty"
        
        # Strip whitespace
        url = url.strip()
        
        # Check length
        if len(url) > cls.MAX_URL_LENGTH:
            return False, f"URL too long (max {cls.MAX_URL_LENGTH} characters)"
        
        # Check for dangerous patterns
        for pattern in cls.DANGEROUS_PATTERNS:
            if re.search(pattern, url, re.IGNORECASE):
                logger.warning(f"Dangerous pattern detected in URL: {pattern}")
                return False, "Invalid URL format (security check failed)"
        
        # Add scheme if missing
        supported_protocols = (
            'http://', 'https://',
            'ftp://', 'ftps://',
            'ws://', 'wss://',
            'rtsp://', 'rtmp://',
            'ssh://', 'sftp://',
        )
        if not any(url.startswith(proto) for proto in supported_protocols):
            url = 'https://' + url
        
        # Parse URL
        try:
            parsed = urlparse(url)
            
            # Check scheme
            if parsed.scheme not in cls.VALID_SCHEMES:
                return False, f"Invalid URL scheme (must be {' or '.join(cls.VALID_SCHEMES)})"
            
            # Check netloc (domain)
            if not parsed.netloc:
                return False, "Invalid URL format (no domain found)"
            
            # Check for IP address (optional - you might want to allow this)
            # Uncomment if you want to block direct IP access
            # if re.match(r'^\d+\.\d+\.\d+\.\d+$', parsed.netloc):
            #     return False, "Direct IP addresses are not allowed"
            
            return True, None
            
        except Exception as e:
            logger.error(f"URL parsing error: {e}")
            return False, "Invalid URL format"
    
    @classmethod
    def sanitize_url(cls, url: str) -> str:
        """Sanitize URL for safe processing"""
        # Strip whitespace
        url = url.strip()
        
        # Supported protocols
        supported_protocols = (
            'http://', 'https://',
            'ftp://', 'ftps://',
            'ws://', 'wss://',
            'rtsp://', 'rtmp://',
            'ssh://', 'sftp://',
        )
        
        # Add scheme if missing
        if not any(url.startswith(proto) for proto in supported_protocols):
            url = 'https://' + url
        
        # Parse and reconstruct to normalize
        try:
            parsed = urlparse(url)
            
            # Reconstruct with only safe components
            sanitized = f"{parsed.scheme}://{parsed.netloc}"
            
            if parsed.path:
                sanitized += parsed.path
            
            if parsed.query:
                sanitized += f"?{parsed.query}"
            
            return sanitized
            
        except Exception:
            return url
    
    @classmethod
    def validate_bookmark_name(cls, name: str) -> tuple[bool, Optional[str]]:
        """Validate bookmark name"""
        if not name:
            return False, "Name cannot be empty"
        
        if len(name) > 100:
            return False, "Name too long (max 100 characters)"
        
        # Check for dangerous patterns
        for pattern in cls.DANGEROUS_PATTERNS:
            if re.search(pattern, name, re.IGNORECASE):
                return False, "Invalid name format"
        
        return True, None
    
    @classmethod
    def sanitize_text(cls, text: str) -> str:
        """Sanitize text for safe display"""
        if not text:
            return ""
        
        # Remove potential XSS patterns
        text = re.sub(r'<script[^>]*>.*?</script>', '', text, flags=re.IGNORECASE | re.DOTALL)
        text = re.sub(r'javascript:', '', text, flags=re.IGNORECASE)
        text = re.sub(r'on\w+\s*=', '', text, flags=re.IGNORECASE)
        
        return text.strip()


class RateLimitTracker:
    """Track and enforce rate limits per user"""
    
    def __init__(self):
        self.user_requests = {}
    
    def check_rate_limit(
        self, 
        user_id: int, 
        max_requests: int = 10, 
        time_window: int = 60
    ) -> tuple[bool, Optional[str]]:
        """
        Check if user has exceeded rate limit
        
        Args:
            user_id: Telegram user ID
            max_requests: Maximum requests allowed
            time_window: Time window in seconds
        
        Returns:
            (is_allowed, error_message)
        """
        import time
        
        now = time.time()
        
        # Initialize user if not exists
        if user_id not in self.user_requests:
            self.user_requests[user_id] = []
        
        # Remove old requests outside time window
        self.user_requests[user_id] = [
            req_time for req_time in self.user_requests[user_id]
            if now - req_time < time_window
        ]
        
        # Check limit
        if len(self.user_requests[user_id]) >= max_requests:
            remaining_time = time_window - (now - self.user_requests[user_id][0])
            return False, f"Rate limit exceeded. Try again in {int(remaining_time)}s"
        
        # Record this request
        self.user_requests[user_id].append(now)
        
        return True, None


class SecurityHeaders:
    """Security headers for web responses"""
    
    @staticmethod
    def get_security_headers() -> dict:
        """Get recommended security headers"""
        return {
            'X-Content-Type-Options': 'nosniff',
            'X-Frame-Options': 'DENY',
            'X-XSS-Protection': '1; mode=block',
            'Strict-Transport-Security': 'max-age=31536000; includeSubDomains',
            'Content-Security-Policy': "default-src 'self'",
            'Referrer-Policy': 'strict-origin-when-cross-origin'
        }


class DataEncryption:
    """Simple encryption utilities for sensitive data"""
    
    @staticmethod
    def hash_sensitive_data(data: str) -> str:
        """Hash sensitive data (e.g., for logging)"""
        import hashlib
        return hashlib.sha256(data.encode()).hexdigest()[:16]
    
    @staticmethod
    def redact_sensitive_info(text: str) -> str:
        """Redact sensitive information from text"""
        # Redact API keys
        text = re.sub(
            r'(?i)(api[_-]?key|token|secret)["\']?\s*[:=]\s*["\']?([a-zA-Z0-9_-]+)',
            r'\1: ***REDACTED***',
            text
        )
        
        # Redact passwords
        text = re.sub(
            r'(?i)(password|passwd|pwd)["\']?\s*[:=]\s*["\']?([^\s"\']+)',
            r'\1: ***REDACTED***',
            text
        )
        
        # Redact email addresses
        text = re.sub(
            r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b',
            '***EMAIL***',
            text
        )
        
        return text


# Global instances
input_validator = InputValidator()
rate_limit_tracker = RateLimitTracker()


# Convenience functions
def validate_url(url: str) -> tuple[bool, Optional[str]]:
    """Validate URL"""
    return input_validator.validate_url(url)


def sanitize_url(url: str) -> str:
    """Sanitize URL"""
    return input_validator.sanitize_url(url)


def check_rate_limit(user_id: int) -> tuple[bool, Optional[str]]:
    """Check rate limit for user"""
    return rate_limit_tracker.check_rate_limit(user_id)


def sanitize_text(text: str) -> str:
    """Sanitize text"""
    return input_validator.sanitize_text(text)