src/aibom-generator/aibom_security.py

"""
Security module for AIBOM generator implementation.

This module provides security functions that can be integrated
into the AIBOM generator to improve input validation, error handling,
and protection against common web vulnerabilities.
"""

import re
import os
import json
import logging
from typing import Dict, Any, Optional, Union

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

def validate_model_id(model_id: str) -> str:
    """
    Validate model ID to prevent injection attacks.
    
    Args:
        model_id: The model ID to validate
        
    Returns:
        The validated model ID
        
    Raises:
        ValueError: If the model ID contains invalid characters
    """
    # Only allow alphanumeric characters, hyphens, underscores, and forward slashes
    if not model_id or not isinstance(model_id, str):
        raise ValueError("Model ID must be a non-empty string")
        
    if not re.match(r'^[a-zA-Z0-9_\-/]+$', model_id):
        raise ValueError(f"Invalid model ID format: {model_id}")
    
    # Prevent path traversal attempts
    if '..' in model_id:
        raise ValueError(f"Invalid model ID - contains path traversal sequence: {model_id}")
        
    return model_id

def safe_path_join(directory: str, filename: str) -> str:
    """
    Safely join directory and filename to prevent path traversal attacks.
    
    Args:
        directory: Base directory
        filename: Filename to append
        
    Returns:
        Safe file path
    """
    # Ensure filename doesn't contain path traversal attempts
    filename = os.path.basename(filename)
    return os.path.join(directory, filename)

def safe_json_parse(json_string: str, default: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """
    Safely parse JSON with error handling.
    
    Args:
        json_string: JSON string to parse
        default: Default value to return if parsing fails
        
    Returns:
        Parsed JSON object or default value
    """
    if default is None:
        default = {}
        
    try:
        return json.loads(json_string)
    except (json.JSONDecodeError, TypeError) as e:
        logger.error(f"Invalid JSON: {e}")
        return default

def sanitize_html_output(text: str) -> str:
    """
    Sanitize text for safe HTML output to prevent XSS attacks.
    
    Args:
        text: Text to sanitize
        
    Returns:
        Sanitized text
    """
    if not text or not isinstance(text, str):
        return ""
        
    # Replace HTML special characters with their entities
    replacements = {
        '&': '&',
        '<': '&lt;',
        '>': '&gt;',
        '"': '&quot;',
        "'": '&#x27;',
        '/': '&#x2F;',
    }
    
    for char, entity in replacements.items():
        text = text.replace(char, entity)
        
    return text

def secure_file_operations(file_path: str, operation: str, content: Optional[str] = None) -> Union[str, bool]:
    """
    Perform secure file operations with proper error handling.
    
    Args:
        file_path: Path to the file
        operation: Operation to perform ('read', 'write', 'append')
        content: Content to write (for 'write' and 'append' operations)
        
    Returns:
        File content for 'read' operation, True for successful 'write'/'append', False otherwise
    """
    try:
        if operation == 'read':
            with open(file_path, 'r', encoding='utf-8') as f:
                return f.read()
        elif operation == 'write' and content is not None:
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(content)
            return True
        elif operation == 'append' and content is not None:
            with open(file_path, 'a', encoding='utf-8') as f:
                f.write(content)
            return True
        else:
            logger.error(f"Invalid file operation: {operation}")
            return False
    except Exception as e:
        logger.error(f"File operation failed: {e}")
        return "" if operation == 'read' else False

def validate_url(url: str) -> bool:
    """
    Validate URL format to prevent malicious URL injection.
    
    Args:
        url: URL to validate
        
    Returns:
        True if URL is valid, False otherwise
    """
    # Basic URL validation
    url_pattern = re.compile(
        r'^(https?):\/\/'  # http:// or https://
        r'(([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])\.)*'  # domain segments
        r'([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-]*[a-zA-Z0-9])'  # last domain segment
        r'(:\d+)?'  # optional port
        r'(\/[-a-zA-Z0-9%_.~#+]*)*'  # path
        r'(\?[;&a-zA-Z0-9%_.~+=-]*)?'  # query string
        r'(\#[-a-zA-Z0-9%_.~+=/]*)?$'  # fragment
    )
    
    return bool(url_pattern.match(url))

def secure_template_rendering(template_content: str, context: Dict[str, Any]) -> str:
    """
    Render templates securely with auto-escaping enabled.
    
    This is a placeholder function. In a real implementation, you would use
    a template engine like Jinja2 with auto-escaping enabled.
    
    Args:
        template_content: Template content
        context: Context variables for rendering
        
    Returns:
        Rendered template
    """
    try:
        from jinja2 import Template
        template = Template(template_content, autoescape=True)
        return template.render(**context)
    except ImportError:
        logger.error("Jinja2 not available, falling back to basic rendering")
        # Very basic fallback (not recommended for production)
        result = template_content
        for key, value in context.items():
            if isinstance(value, str):
                placeholder = "{{" + key + "}}"
                result = result.replace(placeholder, sanitize_html_output(value))
        return result
    except Exception as e:
        logger.error(f"Template rendering failed: {e}")
        return ""

def implement_rate_limiting(user_id: str, action: str, limit: int, period: int) -> bool:
    """
    Implement basic rate limiting to prevent abuse.
    
    This is a placeholder function. In a real implementation, you would use
    a database or cache to track request counts.
    
    Args:
        user_id: Identifier for the user
        action: Action being performed
        limit: Maximum number of actions allowed
        period: Time period in seconds
        
    Returns:
        True if action is allowed, False if rate limit exceeded
    """
    # In a real implementation, you would:
    # 1. Check if user has exceeded limit in the given period
    # 2. If not, increment counter and allow action
    # 3. If yes, deny action
    
    # Placeholder implementation always allows action
    logger.info(f"Rate limiting check for user {user_id}, action {action}")
    return True

# Integration example for the AIBOM generator
def secure_aibom_generation(model_id: str, output_file: Optional[str] = None) -> Dict[str, Any]:
    """
    Example of how to integrate security improvements into AIBOM generation.
    
    Args:
        model_id: Model ID to generate AIBOM for
        output_file: Optional output file path
        
    Returns:
        Generated AIBOM data
    """
    try:
        # Validate input
        validated_model_id = validate_model_id(model_id)
        
        # Process model ID securely
        # (This would call your actual AIBOM generation logic)
        aibom_data = {"message": f"AIBOM for {validated_model_id}"}
        
        # Handle output file securely if provided
        if output_file:
            safe_output_path = safe_path_join(os.path.dirname(output_file), os.path.basename(output_file))
            secure_file_operations(safe_output_path, 'write', json.dumps(aibom_data, indent=2))
            
        return aibom_data
        
    except ValueError as e:
        # Handle validation errors
        logger.error(f"Validation error: {e}")
        return {"error": "Invalid input parameters"}
        
    except Exception as e:
        # Handle unexpected errors
        logger.error(f"AIBOM generation failed: {e}")
        return {"error": "An internal error occurred"}