8f23f59601
Command source decision: extended C2TaskStatus with command: str | None (default None). Added command_name to _GET_TASK_QUERY so get_task() returns command in a single round-trip — no separate history fetch needed on import. 4-line change, zero cascading test impact. adapter.py: - C2TaskStatus: add command: str | None = None field - C2HistoricalTask: new dataclass (display_id, command, params, status, completed, timestamp) for history rows - C2TaskPage.items: typed as list[C2HistoricalTask] (was list[dict]) mythic.py: - _GET_TASK_QUERY: add command_name field - _LIST_CALLBACK_TASKS_QUERY: new query (order_by id desc, limit/offset) - _COUNT_CALLBACK_TASKS_QUERY: new aggregate query for total - get_task(): surfaces command_name as status.command - list_callback_tasks(): two _post() calls (tasks + count), allow_redirects=False fake.py: - _FAKE_HISTORY: frozen deterministic history (cb1=12, cb2=0, cb3=5 tasks) - list_callback_tasks(): serves from _FAKE_HISTORY, pagination applied - get_task(): returns command from _tasks dict api/c2.py: - GET /api/engagements/<eid>/c2/callbacks/<cid>/history: page+page_size defaults 1/25, cap 100, reject <1, 502 on adapter error - POST /api/simulations/<sid>/c2/import: idempotent per (sim,mythic_id) pair, source=import, completed tasks get output+mapping_applied, incomplete tasks stored for poll-on-read pickup, auto-transition pending→in_progress 60 new tests (456 total); pytest/ruff/mypy all green Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
518 lines
16 KiB
Python
518 lines
16 KiB
Python
"""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,
|
|
"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
|