mimic/backend/tests/test_c2_import.py

"""Tests for POST /api/simulations/<id>/c2/import."""
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, C2TaskSource
from backend.app.models.simulation import Simulation, SimulationStatus
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 _import(client: FlaskClient, token: str, sid: int, task_display_ids: list, callback_display_id: int = 1):
    return client.post(
        f"/api/simulations/{sid}/c2/import",
        headers=_h(token),
        json={"callback_display_id": callback_display_id, "task_display_ids": task_display_ids},
    )


def _make_completed_get_task(monkeypatch, command: str = "whoami"):
    """Patch FakeAdapter.get_task to return completed=True with a command."""
    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),
            command=command,
        )

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

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

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


def _advance_to_review_required(client, token, sid):
    client.patch(f"/api/simulations/{sid}", headers=_h(token), json={"name": "Sim Alpha"})
    client.post(f"/api/simulations/{sid}/transition", headers=_h(token), json={"to": "review_required"})


def _advance_to_done(client, admin_token, soc_token, sid):
    _advance_to_review_required(client, admin_token, sid)
    client.post(f"/api/simulations/{sid}/transition", headers=_h(soc_token), json={"to": "done"})


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


class TestImportHappyPath:
    def test_imports_two_completed_tasks(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch, command="whoami")
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        resp = _import(client, admin_token, sim["id"], [100, 101])
        assert resp.status_code == 200
        body = resp.get_json()
        assert body["imported"] == 2
        assert body["skipped"] == 0

        with app.app_context():
            rows = C2Task.query.filter_by(simulation_id=sim["id"]).all()
            assert len(rows) == 2

    def test_imported_tasks_have_source_import(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        _import(client, admin_token, sim["id"], [100])

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

    def test_completed_tasks_get_mapping_applied(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        _import(client, admin_token, sim["id"], [100])

        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_idempotent_import_counts_skipped(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        # First import.
        _import(client, admin_token, sim["id"], [100, 101])

        # Second import with one overlap.
        resp = _import(client, admin_token, sim["id"], [100, 102])
        body = resp.get_json()
        assert body["imported"] == 1
        assert body["skipped"] == 1

    def test_auto_transition_pending_to_in_progress(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        assert sim["status"] == "pending"

        _import(client, admin_token, sim["id"], [100])

        with app.app_context():
            updated = db.session.get(Simulation, sim["id"])
            assert updated is not None
            assert updated.status == SimulationStatus.IN_PROGRESS

    def test_no_transition_when_already_in_progress(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        # Advance to in_progress manually.
        client.patch(
            f"/api/simulations/{sim['id']}",
            headers=_h(admin_token),
            json={"name": "Sim Alpha"},
        )

        _import(client, admin_token, sim["id"], [100])

        with app.app_context():
            updated = db.session.get(Simulation, sim["id"])
            assert updated is not None
            assert updated.status == SimulationStatus.IN_PROGRESS

    def test_no_transition_when_review_required(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])
        _advance_to_review_required(client, admin_token, sim["id"])

        _import(client, admin_token, sim["id"], [100])

        with app.app_context():
            updated = db.session.get(Simulation, sim["id"])
            assert updated is not None
            assert updated.status == SimulationStatus.REVIEW_REQUIRED

    def test_incomplete_task_stored_without_mapping(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """An incomplete task is stored as-is; mapping_applied stays False."""
        from backend.app.services.c2 import fake as fake_mod

        def _submitted(self, task_display_id: int) -> C2TaskStatus:
            return C2TaskStatus(
                display_id=task_display_id,
                status="submitted",
                completed=False,
                command="shell",
            )

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

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

        resp = _import(client, admin_token, sim["id"], [200])
        assert resp.status_code == 200
        assert resp.get_json()["imported"] == 1

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

    def test_command_stored_from_get_task(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """Command field on the stored row comes from adapter.get_task().command."""
        _make_completed_get_task(monkeypatch, command="net user /domain")
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        _import(client, admin_token, sim["id"], [100])

        with app.app_context():
            task = C2Task.query.filter_by(simulation_id=sim["id"]).first()
            assert task is not None
            assert task.command == "net user /domain"

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

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

    def test_source_field_is_import_in_tasks_listing(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """Imported tasks appear with source='import' in GET /c2/tasks response."""
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        _import(client, admin_token, sim["id"], [100])

        resp = client.get(
            f"/api/simulations/{sim['id']}/c2/tasks",
            headers=_h(admin_token),
        )
        assert resp.status_code == 200
        task = resp.get_json()["tasks"][0]
        assert task["source"] == "import"

    def test_no_transition_when_all_skipped(
        self, app: Flask, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        """If imported=0 (all skipped), do not transition pending→in_progress."""
        _make_completed_get_task(monkeypatch)
        eng = _make_engagement(client, admin_token)
        _put_config(client, admin_token, eng["id"])
        sim = _make_sim(client, admin_token, eng["id"])

        _import(client, admin_token, sim["id"], [100])  # first import
        _import(client, admin_token, sim["id"], [])      # empty — should 400 before this matters

        # Reset to pending state via a fresh sim (can't undo, just verify the 0-skipped case).
        # We test: importing same task again = skipped=1, imported=0 → no double-transition.
        resp = _import(client, admin_token, sim["id"], [100])
        body = resp.get_json()
        assert body["imported"] == 0
        assert body["skipped"] == 1


# ---------------------------------------------------------------------------
# Validation errors
# ---------------------------------------------------------------------------


class TestImportValidation:
    def test_400_empty_task_display_ids(
        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 = _import(client, admin_token, sim["id"], [])
        assert resp.status_code == 400

    def test_400_non_int_task_display_id(
        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 = client.post(
            f"/api/simulations/{sim['id']}/c2/import",
            headers=_h(admin_token),
            json={"callback_display_id": 1, "task_display_ids": ["not-an-int"]},
        )
        assert resp.status_code == 400

    def test_400_missing_callback_display_id(
        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 = client.post(
            f"/api/simulations/{sim['id']}/c2/import",
            headers=_h(admin_token),
            json={"task_display_ids": [100]},
        )
        assert resp.status_code == 400

    def test_409_done_simulation(
        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"])
        _advance_to_done(client, admin_token, soc_token, sim["id"])

        resp = _import(client, admin_token, sim["id"], [100])
        assert resp.status_code == 409

    def test_404_simulation_not_found(
        self, client: FlaskClient, admin_token: str
    ) -> None:
        resp = _import(client, admin_token, 9999, [100])
        assert resp.status_code == 404


# ---------------------------------------------------------------------------
# Authorization / error cases
# ---------------------------------------------------------------------------


class TestImportErrors:
    def test_403_soc(
        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 = _import(client, soc_token, sim["id"], [100])
        assert resp.status_code == 403

    def test_503_no_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 = _import(client, admin_token, sim["id"], [100])
        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 = _import(client, admin_token, sim["id"], [100])
        assert resp.status_code == 404

    def test_502_adapter_error_on_get_task(
        self, monkeypatch, client: FlaskClient, admin_token: str
    ) -> None:
        from backend.app.services.c2 import fake as fake_mod

        def _boom(self, task_display_id: int) -> C2TaskStatus:
            raise C2Error("Mythic unreachable")

        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"])

        resp = _import(client, admin_token, sim["id"], [100])
        assert resp.status_code == 502
        assert "Mythic unreachable" in resp.get_json().get("error", "")