mimic/backend/tests/test_c2_tasks_list.py

"""Tests for GET /api/simulations/<id>/c2/tasks — poll-on-read endpoint."""
from __future__ import annotations

import pytest
from cryptography.fernet import Fernet
from flask import Flask
from flask.testing import FlaskClient

from backend.app.extensions import db
from backend.app.models.c2_task import C2Task
from backend.app.models.simulation import Simulation
from backend.app.services.c2.adapter import C2Error, C2TaskStatus
from backend.tests.conftest import auth_headers as _h

_FERNET_KEY = Fernet.generate_key().decode()


@pytest.fixture(autouse=True)
def set_encryption_key(monkeypatch):
    monkeypatch.setenv("MIMIC_ENCRYPTION_KEY", _FERNET_KEY)


@pytest.fixture(autouse=True)
def use_fake_adapter(monkeypatch):
    monkeypatch.setenv("MIMIC_C2_ADAPTER", "fake")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_engagement(client: FlaskClient, token: str) -> dict:
    resp = client.post(
        "/api/engagements",
        headers=_h(token),
        json={"name": "Op Alpha", "start_date": "2026-06-10"},
    )
    assert resp.status_code == 201
    return resp.get_json()


def _put_config(client: FlaskClient, token: str, eid: int) -> None:
    resp = client.put(
        f"/api/engagements/{eid}/c2-config",
        headers=_h(token),
        json={"url": "https://c2.internal:7443", "api_token": "s3cr3t", "verify_tls": True},
    )
    assert resp.status_code == 200


def _make_sim(client: FlaskClient, token: str, eid: int) -> dict:
    resp = client.post(
        f"/api/engagements/{eid}/simulations",
        headers=_h(token),
        json={"name": "Sim Alpha"},
    )
    assert resp.status_code == 201
    return resp.get_json()


def _execute(client: FlaskClient, token: str, sid: int, commands: list, callback_display_id: int = 1):
    return client.post(
        f"/api/simulations/{sid}/c2/execute",
        headers=_h(token),
        json={"callback_display_id": callback_display_id, "commands": commands},
    )


def _list_tasks(client: FlaskClient, token: str, sid: int):
    return client.get(
        f"/api/simulations/{sid}/c2/tasks",
        headers=_h(token),
    )


# ---------------------------------------------------------------------------
# Happy path
# ---------------------------------------------------------------------------


class TestListTasksHappyPath:
    def test_returns_empty_list_when_no_tasks(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 200
        assert resp.get_json()["tasks"] == []

    def test_returns_task_after_execute(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 200
        tasks = resp.get_json()["tasks"]
        assert len(tasks) == 1
        assert tasks[0]["command"] == "whoami"

    def test_task_shape(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["hostname"])

        resp = _list_tasks(client, admin_token, sim["id"])
        task = resp.get_json()["tasks"][0]
        for field in ("id", "mythic_task_display_id", "callback_display_id",
                      "command", "params", "status", "completed", "output",
                      "mapping_applied", "created_at", "completed_at"):
            assert field in task, f"missing field: {field}"

    def test_first_poll_returns_submitted(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        # First GET — FakeAdapter.get_task() first call → submitted.
        resp = _list_tasks(client, admin_token, sim["id"])
        task = resp.get_json()["tasks"][0]
        assert task["status"] == "submitted"
        assert task["completed"] is False

    def test_poll_marks_completed_when_adapter_returns_completed(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """When adapter.get_task returns completed=True the task is updated in DB."""
        from datetime import UTC, datetime

        from backend.app.services.c2 import fake as fake_mod

        def _completed(self, task_display_id: int) -> C2TaskStatus:
            return C2TaskStatus(
                display_id=task_display_id,
                status="completed",
                completed=True,
                completed_at=datetime.now(UTC),
            )

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _completed)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        resp = _list_tasks(client, admin_token, sim["id"])
        task = resp.get_json()["tasks"][0]
        assert task["completed"] is True
        assert task["status"] == "completed"

    def test_output_populated_after_completion(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """Output is fetched and stored when task transitions to completed."""
        from datetime import UTC, datetime

        from backend.app.services.c2 import fake as fake_mod

        def _completed(self, task_display_id: int) -> C2TaskStatus:
            return C2TaskStatus(
                display_id=task_display_id,
                status="completed",
                completed=True,
                completed_at=datetime.now(UTC),
            )

        def _output(self, task_display_id: int) -> str:
            return f"whoami result for task {task_display_id}"

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _completed)
        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task_output", _output)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        resp = _list_tasks(client, admin_token, sim["id"])
        task = resp.get_json()["tasks"][0]
        assert task["output"] is not None
        assert "whoami" in task["output"]

    def test_mapping_applied_set_after_completion(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        from datetime import UTC, datetime

        from backend.app.services.c2 import fake as fake_mod

        def _completed(self, task_display_id: int) -> C2TaskStatus:
            return C2TaskStatus(
                display_id=task_display_id,
                status="completed",
                completed=True,
                completed_at=datetime.now(UTC),
            )

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _completed)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        _list_tasks(client, admin_token, sim["id"])

        with app.app_context():
            task = C2Task.query.filter_by(simulation_id=sim["id"]).first()
            assert task is not None
            assert task.mapping_applied is True

    def test_execution_result_updated_on_simulation(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        from datetime import UTC, datetime

        from backend.app.services.c2 import fake as fake_mod

        def _completed(self, task_display_id: int) -> C2TaskStatus:
            return C2TaskStatus(
                display_id=task_display_id,
                status="completed",
                completed=True,
                completed_at=datetime.now(UTC),
            )

        def _output(self, task_display_id: int) -> str:
            return f"WORKSTATION-01\\whoami output {task_display_id}"

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _completed)
        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task_output", _output)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        _list_tasks(client, admin_token, sim["id"])

        with app.app_context():
            updated_sim = db.session.get(Simulation, sim["id"])
            assert updated_sim is not None
            assert updated_sim.execution_result is not None
            assert "whoami" in updated_sim.execution_result

    def test_completed_task_not_re_polled(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """Once task.completed=True in DB, subsequent GETs skip polling (no re-poll)."""
        from datetime import UTC, datetime

        from backend.app.services.c2 import fake as fake_mod

        call_count = {"n": 0}

        def _completed(self, task_display_id: int) -> C2TaskStatus:
            call_count["n"] += 1
            return C2TaskStatus(
                display_id=task_display_id,
                status="completed",
                completed=True,
                completed_at=datetime.now(UTC),
            )

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _completed)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        _list_tasks(client, admin_token, sim["id"])  # 1st GET — marks task completed (1 call)
        first_count = call_count["n"]

        _list_tasks(client, admin_token, sim["id"])  # 2nd GET — task already completed, skip poll

        # get_task should NOT have been called again on the 2nd GET.
        assert call_count["n"] == first_count, "completed task should not be re-polled"

        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 200
        task = resp.get_json()["tasks"][0]
        assert task["completed"] is True

    def test_redteam_can_list_tasks(
        self, client: FlaskClient, admin_token: str, redteam_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        resp = _list_tasks(client, redteam_token, sim["id"])
        assert resp.status_code == 200


# ---------------------------------------------------------------------------
# Error cases
# ---------------------------------------------------------------------------


class TestListTasksErrors:
    def test_404_simulation_not_found(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        resp = _list_tasks(client, admin_token, 9999)
        assert resp.status_code == 404

    def test_403_soc_forbidden(
        self, client: FlaskClient, admin_token: str, soc_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        resp = _list_tasks(client, soc_token, sim["id"])
        assert resp.status_code == 403

    def test_503_no_encryption_key(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        monkeypatch.delenv("MIMIC_ENCRYPTION_KEY", raising=False)
        eng = _make_engagement(client, admin_token)
        sim = _make_sim(client, admin_token, eng["id"])

        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 503

    def test_404_no_c2_config(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        eng = _make_engagement(client, admin_token)
        sim = _make_sim(client, admin_token, eng["id"])

        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 404

    def test_adapter_error_during_poll_is_tolerated(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """If get_task raises C2Error during poll, the task is skipped (best-effort)."""
        from backend.app.services.c2 import fake as fake_mod

        def _boom(self, task_display_id: int):
            raise C2Error("upstream unavailable")

        monkeypatch.setattr(fake_mod.FakeAdapter, "get_task", _boom)

        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _execute(client, admin_token, sim["id"], ["whoami"])

        # Should still return 200 with the task (un-refreshed status).
        resp = _list_tasks(client, admin_token, sim["id"])
        assert resp.status_code == 200
        tasks = resp.get_json()["tasks"]
        assert len(tasks) == 1
        # Status is stale (not updated due to error) — still "submitted".
        assert tasks[0]["status"] == "submitted"