"""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/ERB delimiters it must be neutralised, otherwise attacker-influenced
config content becomes live template code that executes when Salt/Ansible/Puppet
later 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 re
import subprocess
import sys
from pathlib import Path

import jinja2
import pytest
import yaml as pyyaml

from jinjaturtle.escape import (
    escape_jinja_literal,
    escape_erb_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 }}",
    "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


def test_escape_jinja_literal_wraps_erb_markup_for_translated_erb_safety():
    payload = "<%= boom.run('erb-through-jinja') %>"
    escaped = escape_jinja_literal(payload)
    assert escaped != payload
    env = jinja2.Environment(undefined=jinja2.ChainableUndefined)
    rendered = env.from_string(escaped).render(boom=_Boom())
    assert rendered == payload
    assert TRIP not in rendered


@pytest.mark.parametrize(
    "payload",
    [
        "<%= system('id') %>",
        "<% require 'open3' %>",
        "text <%= 1+1 %> more",
        "-%> orphan <%",
    ],
)
def test_escape_erb_literal_removes_executable_tags(payload):
    escaped = escape_erb_literal(payload)
    # No raw executable ERB tag should survive that contains the original code.
    live = re.findall(r"<%[-=#]?(.*?)-?%>", escaped, re.S)
    for chunk in live:
        assert "system" not in chunk
        assert "require" not in chunk
        # Numeric/expression payloads must be reduced to literal-string prints.


def test_escape_literal_dispatches_by_engine():
    """The public ``escape_literal`` wrapper routes to the right engine."""
    payload = "{{ 7*7 }}"
    # Default engine is Jinja2.
    assert escape_literal(payload) == escape_jinja_literal(payload)
    assert escape_literal(payload, engine="jinja2") == escape_jinja_literal(payload)
    # An unknown engine falls back to the safer Jinja2 escaping.
    assert escape_literal(payload, engine="nonsense") == escape_jinja_literal(payload)
    # ERB routing.
    erb_payload = "<%= system('id') %>"
    assert escape_literal(erb_payload, engine="erb") == escape_erb_literal(erb_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


def test_json_source_keys_are_escaped_and_do_not_execute(tmp_path):
    body = '{"evil_{{ boom.run(\'json-key\') }}": "ok"}\n'
    template_text, defaults = _run_jinjaturtle(tmp_path, "evil.json", body, "json")
    rendered = _render_jinja(template_text, {**defaults, "boom": _Boom()})
    assert TRIP not in rendered
    assert "evil_{{ boom.run('json-key') }}" in rendered


def test_xml_source_comments_cannot_spoof_internal_markers(tmp_path):
    body = (
        "<root>\n"
        "  <!--IF:boom.run('xml-marker')-->\n"
        "  <item>ok</item>\n"
        "  <!--ENDIF:any-->\n"
        "</root>\n"
    )
    template_text, defaults = _run_jinjaturtle(tmp_path, "evil.xml", body, "xml")
    assert "{% if boom.run" not in template_text
    rendered = _render_jinja(template_text, {**defaults, "boom": _Boom()})
    assert TRIP not in rendered
    assert "IF:boom.run('xml-marker')" in rendered


def test_folder_yaml_keys_are_escaped_and_do_not_execute(tmp_path):
    from jinjaturtle.multi import process_directory

    (tmp_path / "one.yaml").write_text(
        'evil_{{ boom.run("folder-yaml") }}: ok\n', encoding="utf-8"
    )
    defaults_yaml, outputs = process_directory(
        tmp_path, recursive=False, role_prefix="role"
    )
    defaults = pyyaml.safe_load(defaults_yaml) or {}
    assert len(outputs) == 1
    template_text = outputs[0].template
    rendered = _render_jinja(
        template_text, {**defaults["role_items"][0], "boom": _Boom()}
    )
    assert TRIP not in rendered
    assert 'evil_{{ boom.run("folder-yaml") }}' in rendered


def test_erb_source_comment_payload_is_literal_not_executed(tmp_path):
    import shutil

    marker = tmp_path / "erb-owned"
    src = tmp_path / "evil.ini"
    src.write_text(
        "[s]\n" "good = ok\n" f"# <%= File.write({str(marker)!r}, 'owned') %>\n",
        encoding="utf-8",
    )
    tpl = tmp_path / "out.erb"
    dfl = tmp_path / "hiera.yml"
    res = subprocess.run(
        [
            sys.executable,
            "-m",
            "jinjaturtle.cli",
            str(src),
            "-f",
            "ini",
            "--template-engine",
            "erb",
            "--role-name",
            "role",
            "-t",
            str(tpl),
            "-d",
            str(dfl),
        ],
        capture_output=True,
        text=True,
    )
    assert res.returncode == 0, res.stderr
    erb_text = tpl.read_text(encoding="utf-8")
    live_chunks = re.findall(r"<%[-=#]?(.*?)-?%>", erb_text, re.S)
    assert not any("File.write" in chunk for chunk in live_chunks)

    ruby = shutil.which("ruby")
    if ruby is None:
        pytest.skip("ruby is not installed")
    run = subprocess.run(
        [
            ruby,
            "-rerb",
            "-e",
            "puts ERB.new(File.read(ARGV[0])).result(binding)",
            str(tpl),
        ],
        capture_output=True,
        text=True,
    )
    assert run.returncode == 0, run.stderr
    assert not marker.exists()


def test_ansible_defaults_tag_dangerous_strings_as_unsafe_and_still_load():
    from jinjaturtle.core import generate_ansible_yaml

    out = generate_ansible_yaml(
        "role",
        [
            (("jinja",), "{{ boom.run('unsafe') }}"),
            (("erb",), "<%= File.write('/tmp/nope', 'x') %>"),
            (("plain",), "ordinary"),
        ],
    )
    assert "role_jinja: !unsafe" in out
    assert "role_erb: !unsafe" in out
    assert "role_plain: ordinary" in out
    loaded = pyyaml.safe_load(out)
    assert loaded["role_jinja"] == "{{ boom.run('unsafe') }}"


def test_cli_rejects_template_unsafe_role_name(tmp_path):
    from jinjaturtle import cli

    src = tmp_path / "a.ini"
    src.write_text("[s]\nkey=value\n", encoding="utf-8")
    with pytest.raises(SystemExit) as exc:
        cli._main([str(src), "--role-name", "bad-{{ role }}"])
    assert exc.value.code == 2


def test_yaml_loop_source_key_cannot_become_jinja_expression(tmp_path):
    from jinjaturtle.core import (
        analyze_loops,
        flatten_config,
        generate_ansible_yaml,
        generate_jinja2_template,
        parse_config,
    )

    malicious_key = (
        "__class__.__mro__[1].__subclasses__()[166].__init__"
        ".__globals__['__builtins__']['__import__']('os')"
        ".popen('echo JT_RCE').read()"
    )
    src = tmp_path / "evil.yaml"
    src.write_text(
        "servers:\n" f"  - {malicious_key!r}: a\n" f"  - {malicious_key!r}: b\n",
        encoding="utf-8",
    )
    fmt, parsed = parse_config(src, "yaml")
    loops = analyze_loops(fmt, parsed)
    flat = flatten_config(fmt, parsed, loops)
    template = generate_jinja2_template(fmt, parsed, "role", src.read_text(), loops)
    defaults = pyyaml.safe_load(generate_ansible_yaml("role", flat, loops)) or {}

    assert f"server.{malicious_key}" not in template
    assert "server.jt_" in template
    rendered = _render_jinja(template, defaults)
    assert TRIP not in rendered
    assert "JT_RCE" in rendered  # literal source key only, not command output


def test_toml_array_table_loop_source_key_cannot_become_jinja_expression(tmp_path):
    from jinjaturtle.core import (
        analyze_loops,
        flatten_config,
        generate_ansible_yaml,
        generate_jinja2_template,
        parse_config,
    )

    malicious_key = (
        "__class__.__mro__[1].__subclasses__()[166].__init__"
        ".__globals__['__builtins__']['__import__']('os')"
        ".popen('echo JT_TOML_RCE').read()"
    )
    src = tmp_path / "evil.toml"
    src.write_text(
        f'[[servers]]\n"{malicious_key}" = "a"\n\n'
        f'[[servers]]\n"{malicious_key}" = "b"\n',
        encoding="utf-8",
    )
    fmt, parsed = parse_config(src, "toml")
    loops = analyze_loops(fmt, parsed)
    flat = flatten_config(fmt, parsed, loops)
    template = generate_jinja2_template(fmt, parsed, "role", src.read_text(), loops)
    defaults = pyyaml.safe_load(generate_ansible_yaml("role", flat, loops)) or {}

    assert f"server.{malicious_key}" not in template
    assert "server.jt_" in template
    rendered = _render_jinja(template, defaults)
    assert TRIP not in rendered
    assert "JT_TOML_RCE" in rendered


def test_erb_loop_source_key_cannot_become_ruby_method_call(tmp_path):
    import shutil

    from jinjaturtle.core import (
        analyze_loops,
        flatten_config,
        generate_erb_template,
        generate_puppet_hiera_yaml,
        parse_config,
    )

    marker = tmp_path / "erb-loop-owned"
    key = f"instance_eval(%q{{File.write({str(marker)!r}, 'owned')}})"
    src = tmp_path / "evil.yaml"
    src.write_text(
        "servers:\n" f"  - {key!r}: a\n" f"  - {key!r}: b\n",
        encoding="utf-8",
    )
    fmt, parsed = parse_config(src, "yaml")
    loops = analyze_loops(fmt, parsed)
    flat = flatten_config(fmt, parsed, loops)
    erb = generate_erb_template(
        fmt,
        parsed,
        "role",
        original_text=src.read_text(),
        loop_candidates=loops,
        flat_items=flat,
    )
    assert "server.instance_eval" not in erb
    assert "server['jt_" in erb

    ruby = shutil.which("ruby")
    if ruby is None:
        pytest.skip("ruby is not installed")

    hiera_path = tmp_path / "hiera.yml"
    hiera_path.write_text(
        generate_puppet_hiera_yaml("role", flat, loops), encoding="utf-8"
    )
    tpl = tmp_path / "safe.erb"
    tpl.write_text(erb, encoding="utf-8")
    run = subprocess.run(
        [
            ruby,
            "-rerb",
            "-ryaml",
            "-e",
            "@servers = YAML.load_file(ARGV[1])['role::servers']; "
            "puts ERB.new(File.read(ARGV[0]), trim_mode: '-').result(binding)",
            str(tpl),
            str(hiera_path),
        ],
        capture_output=True,
        text=True,
    )
    assert run.returncode == 0, run.stderr
    assert not marker.exists()


def test_canonical_ini_and_postfix_escape_source_keys():
    from configparser import ConfigParser

    from jinjaturtle.core import generate_jinja2_template

    parser = ConfigParser()
    parser.add_section('s_{{ boom.run("ini") }}')
    parser.set('s_{{ boom.run("ini") }}', "key", "ok")
    ini_template = generate_jinja2_template("ini", parser, "role", original_text=None)
    rendered = _render_jinja(
        ini_template, {"role_s_boom_run_ini_key": "ok", "boom": _Boom()}
    )
    assert TRIP not in rendered
    assert 's_{{ boom.run("ini") }}' in rendered

    postfix_template = generate_jinja2_template(
        "postfix",
        {'evil_{{ boom.run("postfix") }}': "ok"},
        "role",
        original_text=None,
    )
    rendered = _render_jinja(
        postfix_template,
        {"role_evil_boom_run_postfix": "ok", "boom": _Boom()},
    )
    assert TRIP not in rendered
    assert 'evil_{{ boom.run("postfix") }}' in rendered


def test_template_safety_validator_rejects_future_expression_injection():
    from jinjaturtle.template_safety import (
        UnsafeTemplateError,
        validate_generated_template,
    )

    with pytest.raises(UnsafeTemplateError):
        validate_generated_template("{{ item.__class__.__mro__[1] }}")
    with pytest.raises(UnsafeTemplateError):
        validate_generated_template("{{ boom.run('x') }}")
    # Generated field names and loop.last remain allowed.
    validate_generated_template(
        "{% for item in role_items %}{{ item.jt_name_82a3537f }}"
        "{% if not loop.last %},{% endif %}{% endfor %}"
    )


def test_defaults_generation_rejects_colliding_variable_names():
    from jinjaturtle.core import generate_ansible_yaml

    with pytest.raises(ValueError, match="same variable name"):
        generate_ansible_yaml("role", [(("a-b",), 1), (("a_b",), 2)])