security_test.py

"""
Security test for the Custom Template Engine
Demonstrates protection against SSTI (Server-Side Template Injection)
"""

from template_engine import CustomTemplateEngine

def test_ssti_protection():
    """Test various SSTI attack vectors to ensure they're blocked."""
    
    # Use lenient mode to see filtering behavior
    engine = CustomTemplateEngine(strict_mode=False)
    
    print("SSTI Security Tests")
    print("=" * 50)
    
    # Test 1: Dangerous variable names should be blocked
    dangerous_context = {
        '__import__': __import__,  # This should be filtered out
        'eval': eval,  # This should be filtered out
        'safe_var': 'This is safe'
    }
    
    template1 = "$safe_var $__import__ $eval"
    result1 = engine.render(template1, dangerous_context)
    print("Test 1 - Dangerous imports blocked:")
    print(f"Template: {template1}")
    print(f"Result: {result1}")
    print()
    
    # Test 2: Object attribute access should be blocked
    class DangerousObject:
        def __init__(self):
            self.safe_attr = "safe"
            
        def dangerous_method(self):
            return "This should not execute"
    
    obj_context = {
        'obj': DangerousObject(),
        'safe_string': 'hello'
    }
    
    template2 = "$safe_string $obj.dangerous_method"
    result2 = engine.render(template2, obj_context)
    print("Test 2 - Object method access blocked:")
    print(f"Template: {template2}")
    print(f"Result: {result2}")
    print()
    
    # Test 3: XSS protection with HTML escaping
    xss_context = {
        'user_input': '<script>alert("XSS")</script>',
        'title': 'Page & Title'
    }
    
    template3 = "<h1>$title</h1><div>$user_input</div>"
    result3 = engine.render(template3, xss_context)
    print("Test 3 - XSS protection with HTML escaping:")
    print(f"Template: {template3}")
    print(f"Result: {result3}")
    print()
    
    # Test 4: Invalid variable names should be blocked
    template4 = "$invalid_path $__class__"  # Removed $1invalid as it's syntactically invalid
    safe_context = {'valid_var': 'safe'}
    result4 = engine.render(template4, safe_context)
    print("Test 4 - Invalid variable names blocked:")
    print(f"Template: {template4}")
    print(f"Result: {result4}")
    print()
    
    # Test 5: Safe usage should work normally
    safe_context = {
        'users': [
            {'name': 'John', 'role': 'admin'},
            {'name': 'Jane', 'role': 'user'}
        ],
        'app_name': 'My App'
    }
    
    template5 = """
$app_name Users:
{% for user in users %}
- $user.name ($user.role)
{% endfor %}
"""
    result5 = engine.render(template5, safe_context)
    print("Test 5 - Safe usage works normally:")
    print(f"Result: {result5}")

def test_performance_with_large_data():
    """Test performance and security with larger datasets."""
    
    engine = CustomTemplateEngine()
    
    # Create large dataset
    large_context = {
        'items': [
            {'id': i, 'name': f'Item {i}', 'value': f'Value {i}'}
            for i in range(1000)
        ]
    }
    
    template = """
Total items: $items
{% for item in items %}
$item.id: $item.name = $item.value
{% endfor %}
"""
    
    print("\nPerformance Test with 1000 items:")
    print("=" * 50)
    
    import time
    start_time = time.time()
    result = engine.render(template, large_context)
    end_time = time.time()
    
    # Just show first few lines to avoid spam
    lines = result.strip().split('\n')
    print(f"First 10 lines of output:")
    for line in lines[:10]:
        print(line)
    print(f"... (total {len(lines)} lines)")
    print(f"Rendering time: {end_time - start_time:.4f} seconds")

if __name__ == "__main__":
    test_ssti_protection()
    test_performance_with_large_data()