mimic/backend/app/api/c2.py

"""C2 endpoints — config CRUD and execution.

All endpoints:
- Require admin or redteam role (SOC → 403).
- Return 503 when MIMIC_ENCRYPTION_KEY is not set.
- Never include the cleartext API token in any response.
- Adapter errors → 502 with sanitized message (no URL or token in body).
"""
from __future__ import annotations

from datetime import UTC, datetime
from urllib.parse import urlparse

from flask import Blueprint, jsonify, request

from backend.app.auth import role_required
from backend.app.extensions import db
from backend.app.models import Engagement
from backend.app.models.c2_config import C2Config
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
from backend.app.services.c2.factory import get_adapter
from backend.app.services.c2.mapping import apply_task_to_simulation
from backend.app.services.crypto import C2Disabled, decrypt, encrypt
from backend.app.services.simulation_workflow import promote_to_in_progress

c2_bp = Blueprint("c2", __name__, url_prefix="/api/engagements")
sims_c2_bp = Blueprint("sims_c2", __name__, url_prefix="/api/simulations")

_503_BODY = {"error": "C2 disabled: MIMIC_ENCRYPTION_KEY not set"}


def _crypto_guard():
    """Return a 503 Response when crypto key is absent, else None."""
    try:
        # Attempt a dummy operation to test key availability.
        encrypt("probe")
        return None
    except C2Disabled:
        return jsonify(_503_BODY), 503


@c2_bp.get("/<int:eid>/c2-config")
@role_required("admin", "redteam")
def get_c2_config(eid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    cfg: C2Config | None = engagement.c2_config
    if cfg is None:
        return jsonify({"error": "C2 config not found"}), 404

    return jsonify({
        "has_token": bool(cfg.api_token_encrypted),
        "url": cfg.url,
        "verify_tls": cfg.verify_tls,
    }), 200


@c2_bp.put("/<int:eid>/c2-config")
@role_required("admin", "redteam")
def upsert_c2_config(eid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    data = request.get_json(silent=True) or {}
    url = (data.get("url") or "").strip()
    if not url:
        return jsonify({"error": "url is required"}), 400
    parsed = urlparse(url)
    if parsed.scheme != "https":
        return jsonify({"error": "url must use https"}), 400
    if not parsed.hostname:
        return jsonify({"error": "url must contain a hostname"}), 400

    verify_tls = data.get("verify_tls", True)
    if not isinstance(verify_tls, bool):
        return jsonify({"error": "verify_tls must be a boolean"}), 400

    cfg: C2Config | None = engagement.c2_config

    if cfg is None:
        # New row — api_token is required on creation.
        raw_token = data.get("api_token") or ""
        if not raw_token:
            return jsonify({"error": "api_token is required when creating a config"}), 400
        encrypted = encrypt(raw_token)
        cfg = C2Config(
            engagement_id=eid,
            url=url,
            api_token_encrypted=encrypted,
            verify_tls=verify_tls,
        )
        db.session.add(cfg)
    else:
        # Update — omitting api_token keeps the existing ciphertext.
        cfg.url = url
        cfg.verify_tls = verify_tls
        cfg.updated_at = datetime.now(UTC)
        raw_token = data.get("api_token") or ""
        if raw_token:
            cfg.api_token_encrypted = encrypt(raw_token)

    db.session.commit()
    return jsonify({
        "has_token": True,
        "url": cfg.url,
        "verify_tls": cfg.verify_tls,
    }), 200


@c2_bp.delete("/<int:eid>/c2-config")
@role_required("admin", "redteam")
def delete_c2_config(eid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    cfg: C2Config | None = engagement.c2_config
    if cfg is None:
        return jsonify({"error": "C2 config not found"}), 404

    db.session.delete(cfg)
    db.session.commit()
    return "", 204


@c2_bp.post("/<int:eid>/c2-config/test")
@role_required("admin", "redteam")
def test_c2_config(eid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    cfg: C2Config | None = engagement.c2_config
    if cfg is None:
        return jsonify({"error": "C2 config not found"}), 404

    try:
        api_token = decrypt(cfg.api_token_encrypted)
    except ValueError:
        return jsonify({"ok": False, "error": "Stored token is corrupt"}), 200

    adapter = get_adapter(
        url=cfg.url,
        api_token=api_token,
        verify_tls=cfg.verify_tls,
    )
    health = adapter.test_connection()
    return jsonify({"ok": health.ok, "error": health.error}), 200


# ---------------------------------------------------------------------------
# M2 — callbacks listing + execute
# ---------------------------------------------------------------------------


def _load_adapter_for_engagement(engagement: Engagement):
    """Decrypt token and return adapter, or return a (response, status) error tuple."""
    cfg: C2Config | None = engagement.c2_config
    if cfg is None:
        return None, (jsonify({"error": "C2 config not found"}), 404)
    try:
        api_token = decrypt(cfg.api_token_encrypted)
    except ValueError:
        return None, (jsonify({"error": "Stored token is corrupt"}), 500)
    adapter = get_adapter(url=cfg.url, api_token=api_token, verify_tls=cfg.verify_tls)
    return adapter, None


@c2_bp.get("/<int:eid>/c2/callbacks")
@role_required("admin", "redteam")
def list_callbacks(eid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    adapter, err = _load_adapter_for_engagement(engagement)
    if err is not None:
        return err

    try:
        callbacks = adapter.list_callbacks()
    except C2Error as exc:
        return jsonify({"error": str(exc)}), 502

    return jsonify({
        "callbacks": [
            {
                "display_id": cb.display_id,
                "active": cb.active,
                "host": cb.host,
                "user": cb.user,
                "domain": cb.domain,
                "last_checkin": cb.last_checkin,
            }
            for cb in callbacks
        ]
    }), 200


@sims_c2_bp.post("/<int:sid>/c2/execute")
@role_required("admin", "redteam")
def execute_simulation(sid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    sim = db.session.get(Simulation, sid)
    if sim is None:
        return jsonify({"error": "Simulation not found"}), 404

    # Done is terminal — block execution.
    if sim.status == SimulationStatus.DONE:
        return jsonify({"error": "simulation is done — reopen first"}), 409

    data = request.get_json(silent=True) or {}
    callback_display_id = data.get("callback_display_id")
    commands = data.get("commands")

    if not isinstance(callback_display_id, int):
        return jsonify({"error": "callback_display_id must be an integer"}), 400
    if not isinstance(commands, list) or len(commands) == 0:
        return jsonify({"error": "commands must be a non-empty list"}), 400
    for cmd in commands:
        if not isinstance(cmd, str):
            return jsonify({"error": "each command must be a string"}), 400

    engagement = db.session.get(Engagement, sim.engagement_id)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    adapter, err = _load_adapter_for_engagement(engagement)
    if err is not None:
        return err

    created_tasks = []
    try:
        for command in commands:
            mythic_id = adapter.create_task(
                callback_display_id=callback_display_id,
                command=command,
            )
            task = C2Task(
                simulation_id=sid,
                mythic_task_display_id=mythic_id,
                callback_display_id=callback_display_id,
                command=command,
                params=None,
                status="submitted",
                completed=False,
                source=C2TaskSource.MIMIC,
                created_at=datetime.now(UTC),
            )
            db.session.add(task)
            created_tasks.append(task)
    except C2Error as exc:
        db.session.rollback()
        return jsonify({"error": str(exc)}), 502

    # Auto-transition pending → in_progress (no-op for other statuses).
    promote_to_in_progress(sim)

    db.session.commit()

    return jsonify({
        "tasks": [
            {
                "id": t.id,
                "mythic_task_display_id": t.mythic_task_display_id,
                "command": t.command,
                "status": t.status,
                "completed": t.completed,
            }
            for t in created_tasks
        ]
    }), 200


# ---------------------------------------------------------------------------
# M3 — poll-on-read task listing
# ---------------------------------------------------------------------------


@sims_c2_bp.get("/<int:sid>/c2/tasks")
@role_required("admin", "redteam")
def list_simulation_tasks(sid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    sim = db.session.get(Simulation, sid)
    if sim is None:
        return jsonify({"error": "Simulation not found"}), 404

    engagement = db.session.get(Engagement, sim.engagement_id)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    adapter, err = _load_adapter_for_engagement(engagement)
    if err is not None:
        return err

    tasks: list[C2Task] = C2Task.query.filter_by(simulation_id=sid).all()

    for task in tasks:
        if task.completed:
            continue

        try:
            status = adapter.get_task(task.mythic_task_display_id)
        except C2Error:
            # Best-effort refresh — skip this task if the adapter fails.
            continue

        task.status = status.status
        task.completed = status.completed

        if status.completed:
            task.completed_at = status.completed_at or datetime.now(UTC)
            try:
                task.output = adapter.get_task_output(task.mythic_task_display_id)
            except C2Error:
                task.output = ""
            apply_task_to_simulation(task, sim)

    db.session.commit()

    return jsonify({
        "tasks": [
            {
                "id": t.id,
                "mythic_task_display_id": t.mythic_task_display_id,
                "callback_display_id": t.callback_display_id,
                "command": t.command,
                "params": t.params,
                "status": t.status,
                "completed": t.completed,
                "output": t.output,
                "source": t.source.value,
                "mapping_applied": t.mapping_applied,
                "created_at": t.created_at.isoformat() if t.created_at else None,
                "completed_at": t.completed_at.isoformat() if t.completed_at else None,
            }
            for t in tasks
        ]
    }), 200


# ---------------------------------------------------------------------------
# M4 — callback history + task import
# ---------------------------------------------------------------------------


@c2_bp.get("/<int:eid>/c2/callbacks/<int:cid>/history")
@role_required("admin", "redteam")
def list_callback_history(eid: int, cid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    engagement = db.session.get(Engagement, eid)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    # Validate pagination params.
    try:
        page = int(request.args.get("page", 1))
        page_size = int(request.args.get("page_size", 25))
    except (ValueError, TypeError):
        return jsonify({"error": "page and page_size must be integers"}), 400

    if page < 1 or page_size < 1:
        return jsonify({"error": "page and page_size must be >= 1"}), 400
    if page_size > 100:
        return jsonify({"error": "page_size must be <= 100"}), 400

    adapter, err = _load_adapter_for_engagement(engagement)
    if err is not None:
        return err

    try:
        page_result = adapter.list_callback_tasks(
            callback_display_id=cid,
            page=page,
            page_size=page_size,
        )
    except C2Error as exc:
        return jsonify({"error": str(exc)}), 502

    return jsonify({
        "tasks": [
            {
                "display_id": t.display_id,
                "command": t.command,
                "params": t.params,
                "status": t.status,
                "completed": t.completed,
                "timestamp": t.timestamp,
            }
            for t in page_result.items
        ],
        "total": page_result.total,
        "page": page_result.page,
        "page_size": page_result.page_size,
    }), 200


@sims_c2_bp.post("/<int:sid>/c2/import")
@role_required("admin", "redteam")
def import_tasks(sid: int):
    guard = _crypto_guard()
    if guard is not None:
        return guard

    sim = db.session.get(Simulation, sid)
    if sim is None:
        return jsonify({"error": "Simulation not found"}), 404

    if sim.status == SimulationStatus.DONE:
        return jsonify({"error": "simulation is done — reopen first"}), 409

    data = request.get_json(silent=True) or {}
    callback_display_id = data.get("callback_display_id")
    task_display_ids = data.get("task_display_ids")

    if not isinstance(callback_display_id, int):
        return jsonify({"error": "callback_display_id must be an integer"}), 400
    if not isinstance(task_display_ids, list) or len(task_display_ids) == 0:
        return jsonify({"error": "task_display_ids must be a non-empty list"}), 400
    for tid in task_display_ids:
        if not isinstance(tid, int):
            return jsonify({"error": "each task_display_id must be an integer"}), 400

    engagement = db.session.get(Engagement, sim.engagement_id)
    if engagement is None:
        return jsonify({"error": "Engagement not found"}), 404

    adapter, err = _load_adapter_for_engagement(engagement)
    if err is not None:
        return err

    imported_count = 0
    skipped_count = 0

    try:
        for task_display_id in task_display_ids:
            # Idempotency: skip if already imported for this simulation.
            existing = C2Task.query.filter_by(
                simulation_id=sid,
                mythic_task_display_id=task_display_id,
            ).first()
            if existing is not None:
                skipped_count += 1
                continue

            status = adapter.get_task(task_display_id)
            task = C2Task(
                simulation_id=sid,
                mythic_task_display_id=task_display_id,
                callback_display_id=callback_display_id,
                command=status.command or "",
                params=None,
                status=status.status,
                completed=status.completed,
                source=C2TaskSource.IMPORT,
                created_at=datetime.now(UTC),
                mapping_applied=False,
            )

            if status.completed:
                task.completed_at = status.completed_at or datetime.now(UTC)
                try:
                    task.output = adapter.get_task_output(task_display_id)
                except C2Error:
                    task.output = ""
                db.session.add(task)
                db.session.flush()
                apply_task_to_simulation(task, sim)
                task.mapping_applied = True
            else:
                db.session.add(task)

            imported_count += 1

    except C2Error as exc:
        db.session.rollback()
        return jsonify({"error": str(exc)}), 502

    # Auto-transition pending → in_progress when at least one task was imported.
    if imported_count > 0:
        promote_to_in_progress(sim)

    db.session.commit()
    return jsonify({"imported": imported_count, "skipped": skipped_count}), 200