mimic/backend/tests/test_export_render.py

"""Unit tests for render functions in backend.app.services.export."""
from __future__ import annotations

from datetime import UTC, datetime
from types import SimpleNamespace
from typing import Any

from backend.app.services.export import (
    render_engagement_csv,
    render_engagement_markdown,
    render_engagement_pdf,
)

# ---------------------------------------------------------------------------
# Fixtures / factories
# ---------------------------------------------------------------------------


def _make_engagement(**kw) -> Any:
    from datetime import date

    defaults: dict[str, Any] = {
        "id": 1,
        "name": "Test Engagement",
        "description": "A purple team exercise",
        "start_date": date(2026, 6, 1),
        "end_date": date(2026, 6, 30),
        "status": SimpleNamespace(value="active"),
        "created_at": datetime(2026, 6, 1, 10, 0, tzinfo=UTC),
        "created_by": SimpleNamespace(username="alice"),
    }
    defaults.update(kw)
    return SimpleNamespace(**defaults)


def _make_sim(sid: int = 1, name: str = "Sim Alpha", **kw) -> Any:
    defaults: dict[str, Any] = {
        "id": sid,
        "name": name,
        "status": SimpleNamespace(value="pending"),
        "techniques": [{"id": "T1059", "name": "Command and Scripting Interpreter"}],
        "tactic_ids": ["TA0002"],
        "description": "Execute a script",
        "commands": "whoami",
        "prerequisites": "local admin",
        "executed_at": None,
        "execution_result": None,
        "log_source": None,
        "logs": None,
        "soc_comment": None,
        "incident_number": None,
        "created_at": datetime(2026, 6, 1, 10, 0, tzinfo=UTC),
        "updated_at": None,
        "created_by": SimpleNamespace(username="bob"),
    }
    defaults.update(kw)
    return SimpleNamespace(**defaults)


# ---------------------------------------------------------------------------
# Markdown tests
# ---------------------------------------------------------------------------


def test_render_engagement_markdown_includes_header_fields(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        result = render_engagement_markdown(eng, [])
        assert "Test Engagement" in result
        assert "2026-06-01" in result
        assert "2026-06-30" in result
        assert "active" in result
        assert "alice" in result


def test_render_engagement_markdown_lists_all_simulations_in_order(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sims = [_make_sim(1, "First Sim"), _make_sim(2, "Second Sim")]
        result = render_engagement_markdown(eng, sims)
        first_pos = result.index("First Sim")
        second_pos = result.index("Second Sim")
        assert first_pos < second_pos


def test_render_engagement_markdown_includes_techniques_with_id_and_name(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(
            techniques=[{"id": "T1059", "name": "Command and Scripting Interpreter"}]
        )
        result = render_engagement_markdown(eng, [sim])
        assert "T1059" in result
        assert "Command and Scripting Interpreter" in result


def test_render_engagement_markdown_includes_tactics(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(tactic_ids=["TA0002"])
        result = render_engagement_markdown(eng, [sim])
        # TA0002 = Execution — should appear as tactic name or id
        assert "TA0002" in result or "Execution" in result


def test_render_engagement_markdown_includes_soc_fields_even_when_blank(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(log_source=None, logs=None, soc_comment=None, incident_number=None)
        result = render_engagement_markdown(eng, [sim])
        assert "Log source" in result
        assert "SOC comment" in result
        assert "Incident number" in result


def test_render_engagement_markdown_escapes_backticks_in_commands(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(commands="echo `whoami`")
        result = render_engagement_markdown(eng, [sim])
        # Commands must be inside a fenced code block using ~~~ (tilde fences)
        assert "~~~" in result
        # The backtick content must still be present
        assert "`whoami`" in result


# ---------------------------------------------------------------------------
# CSV tests
# ---------------------------------------------------------------------------


def test_render_engagement_csv_has_header_row(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        result = render_engagement_csv(eng, [])
        first_line = result.splitlines()[0]
        assert "id" in first_line
        assert "name" in first_line
        assert "status" in first_line
        assert "techniques" in first_line


def test_render_engagement_csv_joins_multi_techniques_with_pipe(app) -> None:
    with app.app_context():
        import csv
        import io

        eng = _make_engagement()
        sim = _make_sim(
            techniques=[
                {"id": "T1059", "name": "Command and Scripting Interpreter"},
                {"id": "T1078", "name": "Valid Accounts"},
            ]
        )
        result = render_engagement_csv(eng, [sim])
        rows = list(csv.reader(io.StringIO(result)))
        # row[3] = techniques column
        tech_cell = rows[1][3]
        assert "|" in tech_cell
        assert "T1059" in tech_cell
        assert "T1078" in tech_cell


# ---------------------------------------------------------------------------
# PDF test
# ---------------------------------------------------------------------------


def test_render_engagement_pdf_starts_with_pdf_magic(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim()
        result = render_engagement_pdf(eng, [sim])
        assert isinstance(result, bytes)
        assert result[:4] == b"%PDF"


# ---------------------------------------------------------------------------
# MITRE bundle not loaded
# ---------------------------------------------------------------------------


def test_render_engagement_markdown_with_mitre_bundle_not_loaded_does_not_crash(
    app,
) -> None:
    with app.app_context():
        import backend.app.services.mitre as mitre_svc

        original = mitre_svc.mitre_loaded
        original_tactics = mitre_svc._tactics_by_technique.copy()
        try:
            mitre_svc.mitre_loaded = False
            mitre_svc._tactics_by_technique = {}
            eng = _make_engagement()
            sim = _make_sim(
                techniques=[{"id": "T1059", "name": "Command and Scripting Interpreter"}],
                tactic_ids=["TA0002"],
            )
            result = render_engagement_markdown(eng, [sim])
            # Must not crash and must include the technique id
            assert "T1059" in result
        finally:
            mitre_svc.mitre_loaded = original
            mitre_svc._tactics_by_technique = original_tactics


# ---------------------------------------------------------------------------
# CSV formula injection defense (security fix — 2026-06-08)
# ---------------------------------------------------------------------------

import csv as _csv  # noqa: E402  (sectioned import to keep test diff localized)
import io as _io  # noqa: E402


def _parse_csv_data_row(csv_text: str, row_index: int = 1) -> list[str]:
    """Return cells of row N (0=header) parsed by csv.reader to handle multilines."""
    return list(_csv.reader(_io.StringIO(csv_text)))[row_index]


def test_render_engagement_csv_escapes_formula_injection_in_name(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(name="=cmd|'/c calc'!A1")
        result = render_engagement_csv(eng, [sim])
        cells = _parse_csv_data_row(result)
        # name is column index 1 (after id)
        assert cells[1] == "'=cmd|'/c calc'!A1", (
            "name beginning with = must be apostrophe-prefixed to defuse Excel formula"
        )


def test_render_engagement_csv_escapes_formula_injection_in_commands(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(commands="@SUM(1+1)")
        result = render_engagement_csv(eng, [sim])
        cells = _parse_csv_data_row(result)
        # commands is column index 6 per _CSV_HEADERS order
        assert cells[6].startswith("'@"), "commands beginning with @ must be apostrophe-prefixed"


def test_render_engagement_csv_does_not_alter_safe_strings(app) -> None:
    with app.app_context():
        eng = _make_engagement()
        sim = _make_sim(name="Mimikatz LSASS Dump", commands="whoami /all")
        result = render_engagement_csv(eng, [sim])
        cells = _parse_csv_data_row(result)
        assert cells[1] == "Mimikatz LSASS Dump", "safe name must not be modified"
        assert cells[6] == "whoami /all", "safe commands must not be modified"