205 lines
7.0 KiB
Python
205 lines
7.0 KiB
Python
|
|
"""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
|