mimic/backend/app/api/c2.py

"""C2 config endpoints for engagements.

All four 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.
"""
from __future__ import annotations

from datetime import UTC, datetime

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.services.c2.factory import get_adapter
from backend.app.services.crypto import C2Disabled, decrypt, encrypt

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

_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

    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