# Copyright 2025 Cisco Systems, Inc. and its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

"""Example 1: Complex Jinja2 SSTI with Dynamic Template Construction"""

from jinja2 import Template
from mcp import FastMCP

app = FastMCP("template-injection-service-1")


class ReportGenerator:
    """Generates reports using Jinja2 templates"""

    REPORT_TEMPLATES = {
        "user_activity": "User: {{ username }} - Activity: {{ activity }}",
        "system_status": "System: {{ system_name }} - Status: {{ status }}",
        "custom": "{{ content }}",
    }

    @staticmethod
    def build_template(template_type: str, custom_sections: list = None) -> str:
        """Build template string with custom sections"""
        base_template = ReportGenerator.REPORT_TEMPLATES.get(
            template_type, ReportGenerator.REPORT_TEMPLATES["custom"]
        )

        if custom_sections:
            # VULNERABLE: Concatenating user input into template
            for section in custom_sections:
                base_template += f"\n{section}"

        return base_template

    @staticmethod
    def render_template(template_str: str, context: dict) -> str:
        """Render Jinja2 template with context"""
        # VULNERABLE: User-controlled template string
        template = Template(template_str)
        return template.render(**context)


@app.tool()
def generate_report(
    username: str,
    activity_data: str,
    template_type: str = "user_activity",
    custom_format: str = "",
    include_timestamp: bool = True,
) -> str:
    """
    Generate formatted user activity report with custom templates.
    """
    generator = ReportGenerator()

    # Build custom sections
    custom_sections = []

    if custom_format:
        # VULNERABLE: User-provided format string added to template
        custom_sections.append(custom_format)

    if include_timestamp:
        custom_sections.append("Timestamp: {{ timestamp }}")

    # VULNERABLE: Template construction with user input
    # User can inject: "{{ ''.__class__.__mro__[1].__subclasses__() }}" to access Python internals
    # Or: "{% for item in [].__class__.__base__.__subclasses__() %}{{ item }}{% endfor %}"
    template_str = generator.build_template(template_type, custom_sections)

    # Build context
    context = {
        "username": username,
        "activity": activity_data,
        "system_name": "MCP Server",
        "status": "Active",
        "content": activity_data,
        "timestamp": "2024-01-01 12:00:00",
    }

    # VULNERABLE: Render user-controlled template
    result = generator.render_template(template_str, context)

    return f"Report generated:\n{result}"
