template_engine.py

import string
import html
import re
from typing import Dict, List, Any, Union, Set

class TemplateRenderError(Exception):
    """Base exception for template rendering errors."""
    pass

class MissingVariableError(TemplateRenderError):
    """Raised when a required template variable is missing."""
    pass

class InvalidVariableError(TemplateRenderError):
    """Raised when a variable name contains unsafe characters."""
    pass

class LoopRenderError(TemplateRenderError):
    """Raised when there's an error rendering a loop."""
    pass

class ExtendedTemplate(string.Template):
    """Extended Template that supports dot notation in variable names."""
    # Override the delimiter pattern to include dots in identifiers
    idpattern = r'[a-zA-Z_][a-zA-Z0-9_.]*'
    
    def substitute(self, mapping=None, **kws):
        """Override to provide better error messages for missing variables."""
        if mapping is None:
            mapping = kws
        elif kws:
            mapping = dict(mapping)
            mapping.update(kws)
        
        def convert(mo):
            named = mo.group('named') or mo.group('braced')
            if named is not None:
                if named not in mapping:
                    raise MissingVariableError(f"Missing required variable: '{named}'")
                return str(mapping[named])
            if mo.group('escaped') is not None:
                return self.delimiter
            if mo.group('invalid') is not None:
                self._invalid(mo)
            raise ValueError('Unrecognized named group in pattern', self.pattern)
        
        return self.pattern.sub(convert, self.template)
    
    def safe_substitute(self, mapping=None, **kws):
        """Override to track which variables were not found."""
        if mapping is None:
            mapping = kws
        elif kws:
            mapping = dict(mapping)
            mapping.update(kws)
        
        missing_vars = set()
        
        def convert(mo):
            named = mo.group('named') or mo.group('braced')
            if named is not None:
                if named not in mapping:
                    missing_vars.add(named)
                    return f"${{{named}}}"  # Keep placeholder for missing vars
                return str(mapping[named])
            if mo.group('escaped') is not None:
                return self.delimiter
            if mo.group('invalid') is not None:
                self._invalid(mo)
            raise ValueError('Unrecognized named group in pattern', self.pattern)
        
        result = self.pattern.sub(convert, self.template)
        
        # Store missing variables for later inspection
        if hasattr(self, '_missing_vars'):
            self._missing_vars.update(missing_vars)
        else:
            self._missing_vars = missing_vars
            
        return result

class CustomTemplateEngine:
    """
    A production-grade custom template engine that extends Python's string.Template 
    to support looping through lists of dictionaries with comprehensive error handling
    and full Unicode support for both text and HTML templates.
    
    Features:
    - Variables: $variable_name
    - Loops: {% for item in list_name %}...{% endfor %}
    - Dot notation: $item.key
    - Full Unicode support (UTF-8, emojis, international characters)
    - Text and HTML template modes
    - Missing variable detection
    - XSS protection for HTML templates
    - SSTI prevention
    - Comprehensive error handling
    
    Template Types:
    - Text templates: Plain text with Unicode support, no HTML escaping
    - HTML templates: XSS-safe HTML rendering with Unicode support
    
    Unicode Support:
    - Handles all Unicode characters including emojis
    - Supports international languages (Chinese, Arabic, Russian, etc.)
    
    Raises:
        MissingVariableError: When required variables are missing in strict mode
        InvalidVariableError: When variable names contain unsafe characters
        LoopRenderError: When loop processing fails
    """
    
    def __init__(self, auto_escape=True, strict_mode=True):
        """
        Initialize the template engine.
        
        Args:
            auto_escape: Whether to automatically escape HTML in values
            strict_mode: Whether to raise errors for missing variables
        """
        # Pattern for loop blocks only - let string.Template handle variables
        self.loop_pattern = re.compile(
            r'{%\s*for\s+(\w+)\s+in\s+(\w+)\s*%}(.*?){%\s*endfor\s*%}',
            re.DOTALL
        )
        self.auto_escape = auto_escape
        self.strict_mode = strict_mode
        
        # Whitelist of allowed variable names and paths
        self._safe_chars = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')
        self._safe_key = re.compile(r'^[a-zA-Z0-9_]+$')
        
        # Track missing variables across template rendering
        self._missing_variables = set()
    
    def get_missing_variables(self) -> Set[str]:
        """Get list of variables that were missing during last render."""
        return self._missing_variables.copy()
    
    def _validate_template_syntax(self, template: str) -> None:
        """
        Validate template syntax for common issues.
        
        Args:
            template: The template string to validate
            
        Raises:
            TemplateRenderError: If template syntax is invalid
        """
        # Check for unmatched loop tags
        open_tags = len(re.findall(r'{%\s*for\s+', template))
        close_tags = len(re.findall(r'{%\s*endfor\s*%}', template))
        
        if open_tags != close_tags:
            raise TemplateRenderError(
                f"Unmatched loop tags: {open_tags} opening tags, {close_tags} closing tags"
            )
        
        # Check for nested loops (not currently supported)
        loop_matches = list(self.loop_pattern.finditer(template))
        for match in loop_matches:
            loop_content = match.group(3)
            if '{%' in loop_content and 'for' in loop_content:
                raise TemplateRenderError("Nested loops are not supported")
    
    def _extract_template_variables(self, template: str) -> Set[str]:
        """
        Extract all safe variable placeholders from template.
        
        Args:
            template: The template string
            
        Returns:
            Set of safe variable names found in template
        """
        # Use string.Template's pattern to find variables
        temp_template = ExtendedTemplate(template)
        variables = set()
        
        # Find all variable patterns
        for match in temp_template.pattern.finditer(template):
            named = match.group('named') or match.group('braced')
            if named:
                # Only include safe variable names
                try:
                    self._is_safe_variable_name(named.split('.')[0])  # Check base variable name
                    variables.add(named)
                except InvalidVariableError:
                    # Skip dangerous variables silently
                    continue
        
        return variables
    
    def _sanitize_value(self, value: Any) -> str:
        """
        Sanitize values to prevent injection attacks while preserving Unicode.
        
        Args:
            value: The value to sanitize
            
        Returns:
            Sanitized string value with proper Unicode handling
        """
        if value is None:
            return ''
        
        # Convert to string with Unicode support
        if isinstance(value, bytes):
            # Handle bytes by trying to decode as UTF-8 first
            try:
                str_value = value.decode('utf-8')
            except UnicodeDecodeError:
                # Fallback to latin-1 if UTF-8 fails
                str_value = value.decode('latin-1')
        else:
            str_value = str(value)
        
        # Auto-escape HTML if enabled, but preserve Unicode characters
        if self.auto_escape:
            # Only escape HTML-dangerous characters, not Unicode
            str_value = html.escape(str_value, quote=False)
        
        return str_value
    
    def _is_safe_variable_name(self, name: str) -> bool:
        """
        Check if variable name is safe (alphanumeric + underscore only).
        
        Args:
            name: Variable name to check
            
        Returns:
            True if safe, False otherwise
            
        Raises:
            InvalidVariableError: If name contains dangerous patterns
        """
        if not name or not isinstance(name, str):
            raise InvalidVariableError(f"Invalid variable name type: {type(name)}")
        
        # Check for dangerous patterns
        dangerous_patterns = ['__', 'import', 'eval', 'exec', 'open', 'file']
        for pattern in dangerous_patterns:
            if pattern in name.lower():
                raise InvalidVariableError(f"Variable name contains dangerous pattern: {name}")
        
        if not self._safe_chars.match(name):
            raise InvalidVariableError(f"Variable name contains invalid characters: {name}")
        
        return True
    
    def _is_safe_key(self, key: str) -> bool:
        """
        Check if dictionary key is safe.
        
        Args:
            key: Dictionary key to check
            
        Returns:
            True if safe, False otherwise
        """
        return bool(self._safe_key.match(key))
        """
        Render a template with the given context.
        
        Args:
            template: Template string with placeholders
            context: Dictionary containing variables and data
            
        Returns:
            Rendered string
        """
        # First process loops
        result = self._process_loops(template, context)
        
        # Then process remaining variables
        result = self._process_variables(result, context)
        
    
    def render(self, template: str, context: Dict[str, Any], auto_escape: bool = None) -> str:
        """
        Render a template with the given context using Python's string.Template
        with added loop support and comprehensive error handling.
        
        Args:
            template: Template string with placeholders and loops
            context: Dictionary containing variables and data
            auto_escape: Override the engine's auto_escape setting (optional)
            
        Returns:
            Rendered string
            
        Raises:
            MissingVariableError: When required variables are missing (strict mode)
            InvalidVariableError: When variable names are unsafe
            LoopRenderError: When loop processing fails
            TemplateRenderError: For other template-related errors
        """
        # Store original auto_escape setting
        original_auto_escape = self.auto_escape
        
        # Override auto_escape if provided
        if auto_escape is not None:
            self.auto_escape = auto_escape
        
        try:
            # Reset missing variables tracker
            self._missing_variables = set()
            
            # Validate template syntax
            self._validate_template_syntax(template)
            
            # Extract and validate all variables in template
            template_vars = self._extract_template_variables(template)
            
            # Sanitize context to prevent injection
            safe_context = self._sanitize_context(context)
            
            # Check for missing variables before processing
            if self.strict_mode:
                self._validate_required_variables(template_vars, safe_context, template)
            
            # First process loops
            processed_template = self._process_loops(template, safe_context)
            
            # Then use Python's built-in Template for variable substitution
            # Flatten the context for string.Template (it doesn't handle nested objects)
            flat_context = self._flatten_context(safe_context)
            
            # Use our extended Template engine that supports dot notation
            template_obj = ExtendedTemplate(processed_template)
            
            if self.strict_mode:
                # Use strict substitution that raises errors for missing vars
                result = template_obj.substitute(flat_context)
            else:
                # Use safe substitution that preserves missing vars as placeholders
                result = template_obj.safe_substitute(flat_context)
                # Track missing variables
                if hasattr(template_obj, '_missing_vars'):
                    self._missing_variables.update(template_obj._missing_vars)
            
            return result
            
        except (MissingVariableError, InvalidVariableError, LoopRenderError) as e:
            # Re-raise our custom exceptions
            raise
        except Exception as e:
            # Wrap unexpected exceptions
            error_msg = f"Unexpected template rendering error: {str(e)}"
            raise TemplateRenderError(error_msg) from e
        finally:
            # Restore original auto_escape setting
            self.auto_escape = original_auto_escape
    
    def _validate_required_variables(self, template_vars: Set[str], context: Dict[str, Any], 
                                   original_template: str) -> None:
        """
        Validate that all required variables are present in context.
        
        Args:
            template_vars: Set of variables found in template
            context: The context dictionary
            original_template: Original template for loop analysis
            
        Raises:
            MissingVariableError: If required variables are missing
        """
        # Get flattened context to check all possible variable paths
        flat_context = self._flatten_context(context)
        
        # Also need to consider variables that might be created by loops
        loop_vars = self._extract_loop_variables(original_template, context)
        all_available_vars = set(flat_context.keys()) | loop_vars
        
        # Filter out variables that are inside loop blocks - they are handled separately
        variables_outside_loops = self._filter_loop_variables(original_template, template_vars)
        
        missing_vars = variables_outside_loops - all_available_vars
        
        if missing_vars:
            raise MissingVariableError(
                f"Missing required variables: {', '.join(sorted(missing_vars))}"
            )
    
    def _filter_loop_variables(self, template: str, all_vars: Set[str]) -> Set[str]:
        """
        Filter out variables that are inside loop blocks as they are validated separately.
        
        Args:
            template: The template string
            all_vars: All variables found in template
            
        Returns:
            Variables that are outside of loop blocks
        """
        variables_outside_loops = all_vars.copy()
        
        # Find all loop blocks and remove variables that are inside them
        for match in self.loop_pattern.finditer(template):
            loop_body = match.group(3)
            item_var = match.group(1)
            
            # Extract variables from loop body
            loop_template = ExtendedTemplate(loop_body)
            loop_vars = set()
            
            for var_match in loop_template.pattern.finditer(loop_body):
                named = var_match.group('named') or var_match.group('braced')
                if named and named.startswith(f"{item_var}."):
                    loop_vars.add(named)
            
            # Remove loop-specific variables from the main check
            variables_outside_loops -= loop_vars
        
        return variables_outside_loops
    
    def _extract_loop_variables(self, template: str, context: Dict[str, Any]) -> Set[str]:
        """
        Extract variables that will be available within loop contexts.
        
        Args:
            template: The template string
            context: The context dictionary
            
        Returns:
            Set of variables that will be available in loops
        """
        loop_vars = set()
        
        for match in self.loop_pattern.finditer(template):
            item_var = match.group(1)  # Variable name for each item
            list_var = match.group(2)  # List variable name
            
            if list_var in context and isinstance(context[list_var], list):
                # Add the loop item variable itself
                loop_vars.add(item_var)
                
                # If list contains dicts, add all their keys with dot notation
                for item in context[list_var]:
                    if isinstance(item, dict):
                        for key in item.keys():
                            if self._is_safe_key(str(key)):
                                loop_vars.add(f"{item_var}.{key}")
        
        return loop_vars
    
    def _flatten_context(self, context: Dict[str, Any]) -> Dict[str, str]:
        """
        Flatten nested context for string.Template compatibility.
        Converts nested objects to dot notation keys.
        """
        flat = {}
        
        for key, value in context.items():
            if isinstance(value, dict):
                # Add the dict itself for potential direct access
                flat[key] = str(value)
                # Add nested dictionary keys with dot notation
                for sub_key, sub_value in value.items():
                    if self._is_safe_key(str(sub_key)):
                        flat_key = f"{key}.{sub_key}"
                        flat[flat_key] = self._sanitize_value(sub_value)
            elif isinstance(value, list):
                # Convert list to string representation
                flat[key] = self._sanitize_value(value)
            else:
                flat[key] = self._sanitize_value(value)
        
        return flat
    
    def _sanitize_context(self, context: Dict[str, Any]) -> Dict[str, Any]:
        """
        Sanitize the context to only allow safe data types and structures.
        
        Args:
            context: The original context dictionary
            
        Returns:
            Sanitized context dictionary with only safe variables
        """
        safe_context = {}
        
        for key, value in context.items():
            # Only allow safe variable names - skip dangerous ones silently
            try:
                if not self._is_safe_variable_name(key):
                    continue
            except InvalidVariableError:
                # Skip dangerous variables silently in context sanitization
                continue
                
            # Only allow safe data types
            if isinstance(value, (str, int, float, bool, type(None))):
                safe_context[key] = value
            elif isinstance(value, list):
                # Recursively sanitize list items
                safe_list = []
                for item in value:
                    if isinstance(item, dict):
                        safe_item = {}
                        for k, v in item.items():
                            if (self._is_safe_key(str(k)) and 
                                isinstance(v, (str, int, float, bool, type(None)))):
                                safe_item[str(k)] = v
                        safe_list.append(safe_item)
                    elif isinstance(item, (str, int, float, bool, type(None))):
                        safe_list.append(item)
                safe_context[key] = safe_list
            elif isinstance(value, dict):
                # Sanitize dictionary
                safe_dict = {}
                for k, v in value.items():
                    if (self._is_safe_key(str(k)) and 
                        isinstance(v, (str, int, float, bool, type(None)))):
                        safe_dict[str(k)] = v
                safe_context[key] = safe_dict
                
        return safe_context
    
    def _process_loops(self, template: str, context: Dict[str, Any]) -> str:
        """
        Process all loop blocks in the template with comprehensive error handling.
        
        Args:
            template: Template string containing loop blocks
            context: Context dictionary with variables
            
        Returns:
            Template string with loops processed
            
        Raises:
            LoopRenderError: When loop processing fails
        """
        def replace_loop(match):
            try:
                item_var = match.group(1)  # Variable name for each item
                list_var = match.group(2)  # List variable name
                loop_body = match.group(3)  # Content inside the loop
                
                # Validate variable names
                try:
                    self._is_safe_variable_name(item_var)
                    self._is_safe_variable_name(list_var)
                except InvalidVariableError as e:
                    if self.strict_mode:
                        raise LoopRenderError(f"Invalid loop variable: {str(e)}")
                    return f"<!-- Error: {str(e)} -->"
                
                # Get the list from context
                if list_var not in context:
                    error_msg = f"Loop variable '{list_var}' not found in context"
                    if self.strict_mode:
                        raise LoopRenderError(error_msg)
                    return f"<!-- Error: {error_msg} -->"
                
                items = context[list_var]
                if not isinstance(items, list):
                    error_msg = f"Variable '{list_var}' is not a list (got {type(items).__name__})"
                    if self.strict_mode:
                        raise LoopRenderError(error_msg)
                    return f"<!-- Error: {error_msg} -->"
                
                # Validate that required loop item properties exist
                if self.strict_mode and items:
                    # Check if any loop variables reference missing properties
                    loop_template_obj = ExtendedTemplate(loop_body)
                    loop_vars = set()
                    
                    for var_match in loop_template_obj.pattern.finditer(loop_body):
                        named = var_match.group('named') or var_match.group('braced')
                        if named and named.startswith(f"{item_var}."):
                            loop_vars.add(named)
                    
                    # Check first item for required properties
                    if loop_vars and isinstance(items[0], dict):
                        sample_item = items[0]
                        missing_props = []
                        
                        for var in loop_vars:
                            prop_name = var.split('.', 1)[1]  # Get property name after item_var
                            if prop_name not in sample_item:
                                missing_props.append(var)
                        
                        if missing_props:
                            error_msg = f"Missing properties in loop items: {', '.join(missing_props)}"
                            raise LoopRenderError(error_msg)
                
                # Render the loop body for each item
                rendered_items = []
                for idx, item in enumerate(items):
                    try:
                        # Create a new context with the current item
                        loop_context = context.copy()
                        loop_context[item_var] = item
                        
                        # Flatten the loop context and use string.Template
                        flat_loop_context = self._flatten_context(loop_context)
                        
                        loop_template = ExtendedTemplate(loop_body)
                        
                        if self.strict_mode:
                            rendered_item = loop_template.substitute(flat_loop_context)
                        else:
                            rendered_item = loop_template.safe_substitute(flat_loop_context)
                            
                        rendered_items.append(rendered_item)
                        
                    except Exception as e:
                        error_msg = f"Error rendering loop item {idx}: {str(e)}"
                        if self.strict_mode:
                            raise LoopRenderError(error_msg) from e
                        rendered_items.append(f"<!-- {error_msg} -->")
                
                return ''.join(rendered_items)
                
            except LoopRenderError:
                # Re-raise our custom exceptions
                raise
            except Exception as e:
                error_msg = f"Unexpected error in loop processing: {str(e)}"
                if self.strict_mode:
                    raise LoopRenderError(error_msg) from e
                return f"<!-- {error_msg} -->"
        
        return self.loop_pattern.sub(replace_loop, template)
    
    def render_text(self, template: str, context: Dict[str, Any]) -> str:
        """
        Render a plain text template with Unicode support.
        
        Args:
            template: Plain text template string
            context: Dictionary containing variables and data
            
        Returns:
            Rendered plain text string with Unicode support
        """
        # Temporarily disable HTML escaping for plain text
        original_escape = self.auto_escape
        self.auto_escape = False
        
        try:
            return self.render(template, context)
        finally:
            self.auto_escape = original_escape
    
    def render_html(self, template: str, context: Dict[str, Any], 
                   escape_all=True) -> str:
        """
        Render an HTML template with XSS protection and Unicode support.
        
        Args:
            template: HTML template string
            context: Dictionary containing variables and data
            escape_all: Whether to escape all HTML characters (default: True)
            
        Returns:
            Rendered HTML string with XSS protection and Unicode support
        """
        # Ensure HTML escaping is enabled for HTML templates
        original_escape = self.auto_escape
        self.auto_escape = escape_all
        
        try:
            return self.render(template, context)
        finally:
            self.auto_escape = original_escape

# Convenience alias for the main class
TemplateEngine = CustomTemplateEngine

# Example usage and testing
if __name__ == "__main__":
    
    # Test in strict mode (raises errors for missing variables)
    print("=== Testing Strict Mode ===")
    strict_engine = TemplateEngine(strict_mode=True)
    
    # Test data
    context = {
        'title': 'User List <script>alert("XSS")</script>',  # Test XSS protection
        'users': [
            {'name': 'Alice', 'age': 30, 'email': 'alice@example.com'},
            {'name': 'Bob & Co', 'age': 25, 'email': 'bob@example.com'},  # Test HTML escaping
            {'name': 'Charlie', 'age': 35, 'email': 'charlie@example.com'}
        ],
        'company': 'Tech Corp'
    }
    
    # Test template with missing variable (should raise error in strict mode)
    template_with_missing = """
<h1>$title</h1>
<p>Company: $company</p>
<p>Missing: $missing_var</p>

<ul>
{% for user in users %}
    <li>
        <strong>$user.name</strong> ($user.age years old)
        <br>Email: $user.email
    </li>
{% endfor %}
</ul>
"""
    
    try:
        result = strict_engine.render(template_with_missing, context)
        print("ERROR: Should have failed with missing variable!")
    except MissingVariableError as e:
        print(f"✓ Correctly caught missing variable: {e}")
    
    # Test valid template
    valid_template = """
<h1>$title</h1>
<p>Company: $company</p>

<ul>
{% for user in users %}
    <li>
        <strong>$user.name</strong> ($user.age years old)
        <br>Email: $user.email
    </li>
{% endfor %}
</ul>
"""
    
    try:
        result = strict_engine.render(valid_template, context)
        print("✓ Valid template rendered successfully in strict mode")
    except Exception as e:
        print(f"✗ Unexpected error: {e}")
    
    print("\n=== Testing Non-Strict Mode ===")
    # Test in non-strict mode (preserves missing variables as placeholders)
    lenient_engine = TemplateEngine(strict_mode=False)
    
    result = lenient_engine.render(template_with_missing, context)
    print("✓ Template with missing variable rendered (placeholders preserved)")
    
    missing_vars = lenient_engine.get_missing_variables()
    if missing_vars:
        print(f"✓ Missing variables detected: {missing_vars}")
    
    print("\n=== Security Test ===")
    # Test security features
    dangerous_context = {
        'safe_var': 'This is safe',
        '__import__': 'dangerous',
        'title': '<script>alert("XSS")</script>',
        'users': []
    }
    
    dangerous_template = "$safe_var $__import__ $title"
    
    try:
        result = lenient_engine.render(dangerous_template, dangerous_context)
        print(f"✓ Security test result: {result}")
        print("✓ XSS protection and dangerous variable filtering working")
    except Exception as e:
        print(f"Security test: {e}")
    
    print("\n=== Loop Error Test ===")
    # Test loop with missing list
    loop_template = """
{% for item in missing_list %}
    $item.name
{% endfor %}
"""
    
    try:
        result = strict_engine.render(loop_template, {'other_var': 'test'})
        print("ERROR: Should have failed with missing list!")
    except LoopRenderError as e:
        print(f"✓ Correctly caught loop error: {e}")
    
    print("\n=== Performance Test ===")
    # Test performance with large dataset
    large_context = {
        'items': [{'id': i, 'name': f'Item {i}', 'value': f'Value {i}'} for i in range(1000)]
    }
    
    perf_template = """
Total items: $items
{% for item in items %}
$item.id: $item.name = $item.value
{% endfor %}
"""
    
    import time
    start_time = time.time()
    result = lenient_engine.render(perf_template, large_context)
    end_time = time.time()
    
    lines = result.strip().split('\n')
    print(f"✓ Rendered {len(lines)} lines in {end_time - start_time:.4f} seconds")
    print("✓ Performance test completed")