"""Security regression tests for template injection (SSTI).

JinjaTurtle copies parts of the source config (comments, unrecognised lines,
structural keys) verbatim into the generated template.  If that text contains
Jinja2 delimiters it must be neutralised, otherwise attacker-influenced
config content becomes live template code when a downstream tool renders the
template.

These tests render the *generated* template the way a downstream tool would and
assert that an injected payload never executes.  A tripwire object is exposed
under every name a payload might reference; if the rendered output ever contains
the tripwire sentinel, an injected expression executed and the test fails.
"""

from __future__ import annotations

import subprocess
import sys
from pathlib import Path

import jinja2
import pytest
import yaml as pyyaml

from jinjaturtle.escape import (
    escape_jinja_literal,
    escape_literal,
)

TRIP = "__TRIPWIRE_FIRED__"


class _Boom:
    """Returns the tripwire sentinel for any access/call an SSTI payload makes."""

    def run(self, *a, **k):
        return TRIP

    def __call__(self, *a, **k):
        return TRIP

    def __getitem__(self, k):
        return self

    def __getattr__(self, n):
        return _Boom()

    def __str__(self):
        return TRIP


def _render_jinja(template_text: str, defaults: dict) -> str:
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    env.filters.setdefault("to_json", lambda v, **k: __import__("json").dumps(v))
    env.filters.setdefault("lower", lambda v: str(v).lower())
    ctx = dict(defaults or {})
    for name in ("salt", "cmd", "os", "subprocess", "cycler", "lipsum", "namespace"):
        ctx.setdefault(name, _Boom())
    return env.from_string(template_text).render(**ctx)


def _run_jinjaturtle(tmp_path: Path, source_name: str, body: str, fmt: str):
    src = tmp_path / source_name
    src.write_text(body, encoding="utf-8")
    tpl = tmp_path / "out.tpl"
    dfl = tmp_path / "defaults.yml"
    res = subprocess.run(
        [
            sys.executable,
            "-m",
            "jinjaturtle.cli",
            str(src),
            "-f",
            fmt,
            "--role-name",
            "role",
            "-t",
            str(tpl),
            "-d",
            str(dfl),
        ],
        capture_output=True,
        text=True,
    )
    assert res.returncode == 0, f"generation failed: {res.stderr}"
    defaults = pyyaml.safe_load(dfl.read_text()) or {}
    return tpl.read_text(), defaults


# A representative payload for each format, placed where the format allows
# attacker-controlled verbatim text (comments / unrecognised lines).
FORMAT_CASES = [
    (
        "ini",
        "evil.ini",
        "[s]\n"
        "good = ok ; {{ salt['cmd.run']('id') }}\n"
        "# {{ cmd.run('whoami') }}\n",
    ),
    (
        "yaml",
        "evil.yaml",
        "server:\n" "  motd: ok\n" "  # {{ salt['cmd.run']('id') }}\n",
    ),
    (
        "toml",
        "evil.toml",
        "[s]\n" 'good = "ok"\n' "# {{ cmd.run('id') }}\n",
    ),
    (
        "xml",
        "evil.xml",
        "<config>\n"
        "  <!-- {{ salt['cmd.run']('id') }} -->\n"
        '  <server name="ok"><motd>ok</motd></server>\n'
        "</config>\n",
    ),
    (
        "postfix",
        "main.cf",
        "myhostname = mail.example.com\n" "# {{ salt['cmd.run']('id') }}\n",
    ),
    (
        "systemd",
        "evil.service",
        "[Unit]\n"
        "Description=ok\n"
        "# {{ cmd.run('id') }}\n"
        "RawLineNoEquals {% for x in ().__class__.__bases__ %}\n"
        "[Service]\n"
        "ExecStart=/bin/true\n",
    ),
    (
        "ssh",
        "sshd_config",
        "# {{ salt['cmd.run']('id') }}\n" "Port 22\n" "PermitRootLogin no\n",
    ),
]


@pytest.mark.parametrize("fmt,name,body", FORMAT_CASES)
def test_comment_payload_does_not_execute(tmp_path, fmt, name, body):
    template_text, defaults = _run_jinjaturtle(tmp_path, name, body, fmt)
    rendered = _render_jinja(template_text, defaults)
    assert TRIP not in rendered, f"injected payload executed for {fmt}:\n{rendered}"


@pytest.mark.parametrize("fmt,name,body", FORMAT_CASES)
def test_generated_template_is_renderable(tmp_path, fmt, name, body):
    # A correct escape must still produce a syntactically valid template.
    template_text, defaults = _run_jinjaturtle(tmp_path, name, body, fmt)
    # Should not raise a TemplateSyntaxError.
    _render_jinja(template_text, defaults)


# --- Unit-level guarantees for the escaper itself ---------------------------

SSTI_PAYLOADS = [
    "{{ 7*7 }}",
    "{{ salt['cmd.run']('id') }}",
    "{% set x = cycler.__init__.__globals__ %}{{ x }}",
    "{# comment payload #}",
    "text {% endraw %} breakout {{ evil }}",
    "nested {%  endraw  %} spacing {{ evil }}",
    "{%- endraw -%}{{ evil }}",
    # Whitespace-control markers: Jinja2 accepts "-", "+" or none adjacent to a
    # tag's delimiters, and every variant closes a {% raw %} block.  The "+"
    # forms in particular were a raw-wrapper breakout vector (the defang regex
    # historically only matched "-"), so all combinations must be neutralised.
    "{%+ endraw %}{{ evil }}",
    "{% endraw +%}{{ evil }}",
    "{%+ endraw +%}{{ evil }}",
    "{%- endraw +%}{{ evil }}",
    "{%+ endraw -%}{{ evil }}",
    # Full breakout attempt: close raw early, inject live code, re-open raw to
    # swallow our trailing {% endraw %} so the template would otherwise compile.
    "{%+ endraw %}{{ evil }}{%+ raw %}",
    "mixed {{ a }} and {% b %} and {# c #}",
    "}}{{ orphan delimiters %}{%",
]


@pytest.mark.parametrize("payload", SSTI_PAYLOADS)
def test_escape_jinja_literal_renders_back_to_original(payload):
    """Escaped text must render to the exact original characters, inertly."""
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    escaped = escape_jinja_literal(payload)
    rendered = env.from_string(escaped).render(evil="EVIL", x="X", a="A")
    assert rendered == payload
    assert TRIP not in rendered


def test_escape_jinja_literal_noop_on_plain_text():
    for plain in ["", "hello world", "# a normal comment", "port = 8080", "key: value"]:
        assert escape_jinja_literal(plain) == plain


def test_escape_jinja_literal_actually_blocks_execution():
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    payload = "{{ boom.run('x') }}"
    escaped = escape_jinja_literal(payload)
    rendered = env.from_string(escaped).render(boom=_Boom())
    assert TRIP not in rendered
    # Sanity: the *unescaped* payload would have fired the tripwire.
    fired = env.from_string(payload).render(boom=_Boom())
    assert TRIP in fired


@pytest.mark.parametrize(
    "endraw",
    [
        "{% endraw %}",
        "{% endraw -%}",
        "{% endraw +%}",
        "{%- endraw %}",
        "{%- endraw -%}",
        "{%- endraw +%}",
        "{%+ endraw %}",
        "{%+ endraw -%}",
        "{%+ endraw +%}",
    ],
)
def test_endraw_whitespace_control_cannot_break_out(endraw):
    """Every whitespace-control form of endraw closes a {% raw %} block in
    Jinja2, so each must be defanged.  A payload that closes raw early, injects
    a live tripwire call, then re-opens raw to balance the wrapper must still
    render inertly back to its original characters."""
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    payload = f"{endraw}{{{{ boom.run('x') }}}}{{%+ raw %}}"
    escaped = escape_jinja_literal(payload)
    rendered = env.from_string(escaped).render(boom=_Boom())
    assert TRIP not in rendered
    assert rendered == payload


def test_escape_literal_is_jinja_alias():
    """The public ``escape_literal`` wrapper is Jinja2-only."""
    payload = "{{ 7*7 }}"
    assert escape_literal(payload) == escape_jinja_literal(payload)


def test_escape_literal_jinja_output_is_inert():
    """End-to-end: text routed through escape_literal does not execute."""
    escaped = escape_literal("{{ boom.run('x') }}")
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    rendered = env.from_string(escaped).render(boom=_Boom())
    assert TRIP not in rendered