diff --git a/DESIGN.md b/DESIGN.md index f0e843f..a4cd1d0 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -2,7 +2,7 @@ Mimic is a **BAS (Breach and Attack Simulation) purple-team console** — not a product catalog, not a marketing page. The aesthetic is **Bloomberg Terminal / SOC dashboard**: dense, angular, semantic-color-driven, zero ornament. Every surface decision reinforces operational trust: data is primary, chrome is invisible. -The system sits on a **pure-white canvas** (light) / **dark slab** (dark) with one chromatic action color — **Electric Blue** (`{colors.primary}` — `#024ad8`) — and semantic status signals (`success`, `warn`, `bloom-deep`). Inter handles body/headers/labels. JetBrains Mono carries data: IDs, ISO dates, commands, execution output, MITRE technique codes, metrics — anything that must be read at a glance without typographic distraction. +The system sits on a **pale-tinted canvas** (light: `#f3f5f8`) / **dark slab** (dark) with one chromatic action color — **Electric Blue** (`{colors.primary}` — `#024ad8`) — and semantic status signals (`success`, `warn`, `bloom-deep`). Inter handles body/headers/labels. JetBrains Mono carries data: IDs, ISO dates, commands, execution output, MITRE technique codes, metrics — anything that must be read at a glance without typographic distraction. **No rounded corners on containers.** No shadows on interactive surfaces. No transitions. Hover is instantaneous. Focus rings are sharp. This is a tool, not a storefront. @@ -25,7 +25,7 @@ The system sits on a **pure-white canvas** (light) / **dark slab** (dark) with o - **Soft Blue** (`{colors.primary-soft}` — `#c9e0fc`): selection highlight, chip background on light surfaces. ### Surface -- **Canvas** (`{colors.canvas}` — `#ffffff` light / `#111827` dark): universal page background. +- **Canvas** (`{colors.canvas}` — `#f3f5f8` light / `#111827` dark): universal page background. In light mode, canvas is tinted while paper stays pure white so cards lift without shadow or radius, preserving brutalism. - **Paper** (`{colors.paper}` — `#ffffff` light / `#1f2937` dark): card and panel surfaces. - **Cloud** (`{colors.cloud}` — `#f7f7f7` light / `#1f2937` dark): alternating section band, table row zebra. - **Fog** (`{colors.fog}` — `#e8e8e8` light / `#374151` dark): secondary section band, input background on dense panels. diff --git a/SPEC.md b/SPEC.md index 20244cd..7331b08 100644 --- a/SPEC.md +++ b/SPEC.md @@ -59,8 +59,70 @@ CSV : exactement 1 ligne d'en-tête + 1 ligne par simulation. Markdown : en-têt Prévoir un module d'authentification : dans un premier temps local à la bdd. Dans un premier temps, il s'agit juste de notifier manuellement de l'exécution et les résultats des tests. -Dans un second temps, après que la V1 soit terminée, nous ajouterons une couche permettant de se connecter à un C2 (mythic ou maison) afin d'exécuter des simulation au travers du C2. - +Dans un second temps, après que la V1 soit terminée, nous ajouterons une couche permettant de se connecter à un C2 (mythic ou maison) afin d'exécuter des simulation au travers du C2. + +## Intégration C2 (Sprint 8+) + +Couche d'intégration C2 permettant d'exécuter les commandes d'une simulation à travers un Command & Control distant, suivre l'avancement des tâches en quasi-temps réel, et importer l'historique d'exécutions existant. **Implémentation de référence : Mythic 3.x**, derrière une interface `C2Adapter` mince qui ne ferme pas la porte à un C2 maison ultérieur. + +**RBAC C2 = ressource Red Team uniquement** (précédent Templates + Export) : admin et redteam ont accès complet (config + exécution + import). SOC retourne 403 sur tous les endpoints C2 (pas de nav link, pas d'affichage du panneau C2). + +**Configuration par engagement** : chaque engagement possède au plus une `c2_config` (URL Mythic + API token + flag `verify_tls`). Le token est **chiffré au repos** via `cryptography.Fernet` ; la clé est dérivée de l'env var `MIMIC_ENCRYPTION_KEY` (variable obligatoire pour activer la fonctionnalité C2 — jamais hardcodée, conforme à la règle OPSEC zero-secret-in-code). Le token n'est jamais renvoyé en clair par l'API — `GET /api/engagements//c2-config` retourne `has_token: bool` uniquement. Mise à jour via `PUT` ; suppression via `DELETE`. La suppression d'un engagement supprime en cascade sa `c2_config`. + +**Sélection d'adapter** via l'env var `MIMIC_C2_ADAPTER` : +- `mythic` (défaut) : adapter Mythic réel (GraphQL via Hasura). +- `fake` : adapter en mémoire déterministe utilisé pour la validation Playwright et le dev local sans instance Mythic. + +**Modèle de données — additions** : + +`c2_config` (1 ligne par engagement au max) : +| Colonne | Type | Notes | +|---|---|---| +| `id` | int PK | | +| `engagement_id` | int FK `engagements.id` ON DELETE CASCADE, **UNIQUE** | | +| `url` | text | endpoint Mythic, ex. `https://lab.internal:7443` | +| `api_token_encrypted` | text | Fernet ciphertext, jamais en clair | +| `verify_tls` | bool, défaut `true` | `false` autorisé pour labs auto-signés | +| `created_at`, `updated_at` | datetime | | + +`c2_task` (lien simulation ↔ tâche Mythic) : +| Colonne | Type | Notes | +|---|---|---| +| `id` | int PK | | +| `simulation_id` | int FK `simulations.id` ON DELETE CASCADE | | +| `mythic_task_display_id` | int | identifiant côté Mythic | +| `callback_display_id` | int | callback Mythic sur lequel la tâche tourne | +| `command` | text | commande envoyée | +| `params` | text nullable | paramètres associés | +| `status` | text | statut brut Mythic (`submitted`, `completed`, `error`, …) | +| `completed` | bool | `true` quand la tâche est terminée | +| `output` | text nullable | sortie décodée (base64 → utf-8 ; binaire → préfixe `` + hex) | +| `source` | enum `mimic` \| `import` | tâche lancée depuis Mimic ou importée a posteriori | +| `created_at` | datetime | | +| `completed_at` | datetime nullable | timestamp de complétion | + +**Endpoints C2** (tous admin+redteam ; SOC = 403) : +- `GET /api/engagements//c2-config` — `{has_token, url, verify_tls}` (jamais le token en clair). +- `PUT /api/engagements//c2-config` — `{url, api_token?, verify_tls}`. +- `DELETE /api/engagements//c2-config`. +- `POST /api/engagements//c2-config/test` — test de connectivité via l'adapter, renvoie `{ok, error?}`. +- `GET /api/engagements//c2/callbacks` — callbacks actifs de l'instance Mythic configurée. +- `POST /api/simulations//c2/execute` `{callback_display_id, commands: [str]}` — une tâche Mythic par commande, stockées dans `c2_task` (source=`mimic`). **Auto-transition** : si la simulation est `pending`, elle passe à `in_progress` (même règle que l'édition manuelle RT — cf. § Fonctionnement). +- `GET /api/simulations//c2/tasks` — poll-on-read : à la lecture, rafraîchit le statut et l'output des `c2_task` non terminées depuis Mythic, applique le mapping de sortie (voir ci-dessous) à la simulation pour chaque tâche qui vient de se terminer (idempotent — appliqué une seule fois par tâche). +- `GET /api/engagements//c2/callbacks//history?page=` — historique paginé des tâches d'un callback, pour l'import. +- `POST /api/simulations//c2/import` `{task_display_ids: [int]}` — import sélectif de tâches (source=`import`) + mapping de sortie. + +**Mapping de sortie vers la simulation** (appliqué une fois par tâche, lors de la complétion ou de l'import) : +- `simulation.execution_result` reçoit en append le bloc `\n$ \n\n` (préserve l'existant, jamais d'écrasement). +- `simulation.executed_at` est renseigné depuis le timestamp de la première tâche complétée si le champ est vide ; sinon non modifié. +- `simulation.commands` reçoit en append la commande si elle n'y figure pas déjà (déduplication ligne par ligne). + +**Suivi temps réel** : polling court — le frontend re-fetch `GET /api/simulations//c2/tasks` toutes les **2 500 ms** via TanStack Query `refetchInterval` tant qu'une tâche attachée n'est pas terminée ; le polling s'arrête automatiquement quand toutes les tâches sont `completed`. Pas d'infrastructure ajoutée côté serveur (pas de WebSocket, pas de scheduler). + +**UI** : les contrôles C2 vivent dans la carte Red Team de l'écran simulation — bouton `[Execute via C2]` ouvrant une modale (picker de callback + textarea de commandes pré-remplie depuis `commands`), panneau des tâches attachées sous la carte, et modale d'import historique. Configuration C2 visible/éditable depuis l'écran de détail/édition d'engagement. + +**Validation** : MVP entièrement mocké — pytest utilise un adapter mocké (zéro HTTP live), Playwright utilise l'adapter `fake` (déterministe). Le branchement contre une instance Mythic réelle est repoussé au premier usage opérationnel et peut nécessiter un patch mineur du contrat GraphQL. + ## Stacks techniques * **FrontEnd** : WebUI - Stacks standard : ReactJS, Vite, TailWind etc... diff --git a/backend/app/__init__.py b/backend/app/__init__.py index 350a0e2..3333543 100644 --- a/backend/app/__init__.py +++ b/backend/app/__init__.py @@ -6,7 +6,15 @@ from pathlib import Path from flask import Flask, jsonify, send_from_directory -from backend.app.api import auth_bp, engagements_bp, simulations_bp, templates_bp, users_bp +from backend.app.api import ( + auth_bp, + c2_bp, + engagements_bp, + sims_c2_bp, + simulations_bp, + templates_bp, + users_bp, +) from backend.app.cli import register_cli from backend.app.config import Config, TestConfig from backend.app.errors import register_error_handlers @@ -38,6 +46,8 @@ def create_app(config_object: object | None = None) -> Flask: app.register_blueprint(engagements_bp) app.register_blueprint(simulations_bp) app.register_blueprint(templates_bp) + app.register_blueprint(c2_bp) + app.register_blueprint(sims_c2_bp) from backend.app.services import mitre as mitre_svc mitre_svc.load_bundle() diff --git a/backend/app/api/__init__.py b/backend/app/api/__init__.py index 780821a..7faef15 100644 --- a/backend/app/api/__init__.py +++ b/backend/app/api/__init__.py @@ -1,8 +1,17 @@ """API blueprints.""" from backend.app.api.auth import auth_bp +from backend.app.api.c2 import c2_bp, sims_c2_bp from backend.app.api.engagements import engagements_bp from backend.app.api.simulations import simulations_bp from backend.app.api.templates import templates_bp from backend.app.api.users import users_bp -__all__ = ["auth_bp", "users_bp", "engagements_bp", "simulations_bp", "templates_bp"] +__all__ = [ + "auth_bp", + "c2_bp", + "sims_c2_bp", + "users_bp", + "engagements_bp", + "simulations_bp", + "templates_bp", +] diff --git a/backend/app/api/c2.py b/backend/app/api/c2.py new file mode 100644 index 0000000..745637c --- /dev/null +++ b/backend/app/api/c2.py @@ -0,0 +1,517 @@ +"""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("//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("//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("//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("//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("//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("//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("//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("//c2/callbacks//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("//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) + 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 diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index e432347..693b091 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -1,4 +1,6 @@ """SQLAlchemy models.""" +from backend.app.models.c2_config import C2Config +from backend.app.models.c2_task import C2Task, C2TaskSource from backend.app.models.engagement import Engagement, EngagementStatus from backend.app.models.simulation import Simulation, SimulationStatus from backend.app.models.simulation_template import SimulationTemplate @@ -12,4 +14,7 @@ __all__ = [ "Simulation", "SimulationStatus", "SimulationTemplate", + "C2Config", + "C2Task", + "C2TaskSource", ] diff --git a/backend/app/models/c2_config.py b/backend/app/models/c2_config.py new file mode 100644 index 0000000..4015a13 --- /dev/null +++ b/backend/app/models/c2_config.py @@ -0,0 +1,34 @@ +"""C2Config model — per-engagement Mythic connection settings.""" +from __future__ import annotations + +from datetime import UTC, datetime + +from backend.app.extensions import db + + +class C2Config(db.Model): # type: ignore[name-defined] + __tablename__ = "c2_config" + + id = db.Column(db.Integer, primary_key=True) + engagement_id = db.Column( + db.Integer, + db.ForeignKey("engagements.id", ondelete="CASCADE"), + nullable=False, + unique=True, + index=True, + ) + url = db.Column(db.Text, nullable=False) + api_token_encrypted = db.Column(db.Text, nullable=False) + verify_tls = db.Column(db.Boolean, nullable=False, default=True) + created_at = db.Column( + db.DateTime, nullable=False, default=lambda: datetime.now(UTC) + ) + updated_at = db.Column(db.DateTime, nullable=True) + + engagement = db.relationship( + "Engagement", + backref=db.backref("c2_config", uselist=False, cascade="all, delete-orphan"), + ) + + def __repr__(self) -> str: + return f"" diff --git a/backend/app/models/c2_task.py b/backend/app/models/c2_task.py new file mode 100644 index 0000000..b44a08f --- /dev/null +++ b/backend/app/models/c2_task.py @@ -0,0 +1,48 @@ +"""C2Task model — link between a Mimic simulation and a Mythic task.""" +from __future__ import annotations + +import enum +from datetime import UTC, datetime + +from backend.app.extensions import db + + +class C2TaskSource(str, enum.Enum): + MIMIC = "mimic" + IMPORT = "import" + + +class C2Task(db.Model): # type: ignore[name-defined] + __tablename__ = "c2_task" + + id = db.Column(db.Integer, primary_key=True) + simulation_id = db.Column( + db.Integer, + db.ForeignKey("simulations.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + mythic_task_display_id = db.Column(db.Integer, nullable=False) + callback_display_id = db.Column(db.Integer, nullable=False) + command = db.Column(db.Text, nullable=False) + params = db.Column(db.Text, nullable=True) + status = db.Column(db.Text, nullable=False) + completed = db.Column(db.Boolean, nullable=False, default=False) + output = db.Column(db.Text, nullable=True) + source = db.Column( + db.Enum(C2TaskSource, name="c2task_source"), + nullable=False, + ) + created_at = db.Column( + db.DateTime, nullable=False, default=lambda: datetime.now(UTC) + ) + completed_at = db.Column(db.DateTime, nullable=True) + mapping_applied = db.Column(db.Boolean, nullable=False, default=False) + + simulation = db.relationship( + "Simulation", + backref=db.backref("c2_tasks", cascade="all, delete-orphan", lazy="dynamic"), + ) + + def __repr__(self) -> str: + return f"" diff --git a/backend/app/services/c2/__init__.py b/backend/app/services/c2/__init__.py new file mode 100644 index 0000000..768b149 --- /dev/null +++ b/backend/app/services/c2/__init__.py @@ -0,0 +1,22 @@ +"""C2 adapter package. Import the factory from here.""" +from backend.app.services.c2.adapter import ( + C2Adapter, + C2Callback, + C2Error, + C2Health, + C2TaskPage, + C2TaskStatus, + decode_response_text, +) +from backend.app.services.c2.factory import get_adapter + +__all__ = [ + "C2Adapter", + "C2Callback", + "C2Error", + "C2Health", + "C2TaskPage", + "C2TaskStatus", + "decode_response_text", + "get_adapter", +] diff --git a/backend/app/services/c2/adapter.py b/backend/app/services/c2/adapter.py new file mode 100644 index 0000000..6d45a82 --- /dev/null +++ b/backend/app/services/c2/adapter.py @@ -0,0 +1,117 @@ +"""Abstract C2 adapter interface and shared dataclasses.""" +from __future__ import annotations + +import base64 +import binascii +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from datetime import datetime + + +class C2Error(Exception): + """Raised by adapters when the C2 returns an application-level error.""" + + +@dataclass +class C2Health: + ok: bool + error: str | None = None + + +@dataclass +class C2Callback: + display_id: int + active: bool + host: str + user: str + domain: str + last_checkin: str # ISO-8601 string + + +@dataclass +class C2TaskStatus: + display_id: int + status: str + completed: bool + completed_at: datetime | None = field(default=None) + # command_name is populated by get_task() so import doesn't need a second round-trip. + command: str | None = field(default=None) + + +@dataclass +class C2HistoricalTask: + """A task entry from callback history (carries command + params, unlike C2TaskStatus).""" + + display_id: int + command: str + params: str | None + status: str + completed: bool + timestamp: str | None # ISO-8601 or None + + +@dataclass +class C2TaskPage: + items: list[C2HistoricalTask] + total: int + page: int + page_size: int + + +def decode_response_text(raw: str) -> str: + """Decode a base64-encoded Mythic response_text field. + + On binascii.Error (binary payload) returns " " + hex string + so execution_result never silently corrupts. + """ + try: + return base64.b64decode(raw).decode("utf-8") + except binascii.Error: + return " " + raw.encode().hex() + except UnicodeDecodeError: + raw_bytes = base64.b64decode(raw) + return " " + raw_bytes.hex() + + +class C2Adapter(ABC): + """Thin interface over a C2 backend (Mythic or custom).""" + + @abstractmethod + def test_connection(self) -> C2Health: + """Verify that the C2 is reachable and the token is valid.""" + ... + + @abstractmethod + def list_callbacks(self) -> list[C2Callback]: + """Return active callbacks visible to this API token.""" + ... + + @abstractmethod + def create_task( + self, + callback_display_id: int, + command: str, + params: str | None = None, + ) -> int: + """Issue a task and return its Mythic display_id.""" + ... + + @abstractmethod + def get_task(self, task_display_id: int) -> C2TaskStatus: + """Return current status of a task.""" + ... + + @abstractmethod + def get_task_output(self, task_display_id: int) -> str: + """Return decoded, concatenated output for a completed task.""" + ... + + @abstractmethod + def list_callback_tasks( + self, + callback_display_id: int, + page: int = 1, + page_size: int = 25, + ) -> C2TaskPage: + """Return a paginated history of tasks for a callback.""" + ... diff --git a/backend/app/services/c2/factory.py b/backend/app/services/c2/factory.py new file mode 100644 index 0000000..0c46370 --- /dev/null +++ b/backend/app/services/c2/factory.py @@ -0,0 +1,19 @@ +"""Factory that resolves the C2Adapter implementation from MIMIC_C2_ADAPTER env.""" +from __future__ import annotations + +import os + +from backend.app.services.c2.adapter import C2Adapter + + +def get_adapter(url: str, api_token: str, verify_tls: bool = True) -> C2Adapter: + """Return the correct C2Adapter based on MIMIC_C2_ADAPTER (default: mythic).""" + adapter_name = os.environ.get("MIMIC_C2_ADAPTER", "mythic").lower() + + if adapter_name == "fake": + from backend.app.services.c2.fake import FakeAdapter + return FakeAdapter() + + # Default: real Mythic adapter + from backend.app.services.c2.mythic import MythicAdapter + return MythicAdapter(url=url, api_token=api_token, verify_tls=verify_tls) diff --git a/backend/app/services/c2/fake.py b/backend/app/services/c2/fake.py new file mode 100644 index 0000000..7835d9c --- /dev/null +++ b/backend/app/services/c2/fake.py @@ -0,0 +1,176 @@ +"""Deterministic in-memory C2 adapter — used when MIMIC_C2_ADAPTER=fake. + +Intended for integration tests and local development without a live Mythic instance. +Task state is per-instance so parallel tests don't interfere with each other. +""" +from __future__ import annotations + +from backend.app.services.c2.adapter import ( + C2Adapter, + C2Callback, + C2Error, + C2Health, + C2HistoricalTask, + C2TaskPage, + C2TaskStatus, +) + +# Frozen base timestamp — all fake history tasks share this prefix for determinism. +_BASE_TS = "2026-06-10T00:00:00Z" + +# Deterministic history for list_callback_tasks: +# callback 1 → 12 tasks, callback 2 → 0 tasks, callback 3 → 5 tasks. +# Commands cycle through a fixed set; even-indexed tasks are completed. +_HISTORY_COMMANDS = ["whoami", "hostname", "id", "ipconfig", "net user", "pwd"] + +_FAKE_HISTORY: dict[int, list[C2HistoricalTask]] = { + 1: [ + C2HistoricalTask( + display_id=100 + i, + command=_HISTORY_COMMANDS[i % len(_HISTORY_COMMANDS)], + params=None, + status="completed" if i % 2 == 0 else "submitted", + completed=i % 2 == 0, + timestamp=_BASE_TS if i % 2 == 0 else None, + ) + for i in range(12) + ], + 2: [], + 3: [ + C2HistoricalTask( + display_id=200 + i, + command=_HISTORY_COMMANDS[i % len(_HISTORY_COMMANDS)], + params=None, + status="completed" if i % 2 == 0 else "submitted", + completed=i % 2 == 0, + timestamp=_BASE_TS if i % 2 == 0 else None, + ) + for i in range(5) + ], +} + +# Three fixed callbacks the test suite can pin against. +_FAKE_CALLBACKS = [ + C2Callback( + display_id=1, + active=True, + host="WORKSTATION-01", + user="jdoe", + domain="LAB", + last_checkin="2026-06-10T00:00:00Z", + ), + C2Callback( + display_id=2, + active=True, + host="SERVER-DC01", + user="svc_backup", + domain="LAB", + last_checkin="2026-06-10T00:01:00Z", + ), + C2Callback( + display_id=3, + active=True, + host="LAPTOP-RT", + user="admin", + domain="LAB", + last_checkin="2026-06-10T00:02:00Z", + ), +] + + +class FakeAdapter(C2Adapter): + """In-memory adapter with deterministic behaviour. + + Each instance starts with an empty task store and display_ids from 1000. + + get_task() state progression per task (keyed by display_id): + - First call after create_task → submitted, completed=False + - Second and subsequent calls → completed=True, status="completed" + """ + + def __init__(self) -> None: + self._tasks: dict[int, dict] = {} + self._next_task_id = 1000 + # Tracks how many times get_task has been called per display_id. + self._get_task_calls: dict[int, int] = {} + + def test_connection(self) -> C2Health: + return C2Health(ok=True) + + def list_callbacks(self) -> list[C2Callback]: + return list(_FAKE_CALLBACKS) + + def create_task( + self, + callback_display_id: int, + command: str, + params: str | None = None, + ) -> int: + tid = self._next_task_id + self._next_task_id += 1 + self._tasks[tid] = { + "display_id": tid, + "callback_display_id": callback_display_id, + "command": command, + "params": params, + "status": "submitted", + "completed": False, + "output": None, + } + return tid + + def get_task(self, task_display_id: int) -> C2TaskStatus: + """Deterministic state progression: first call → submitted, second+ → completed. + + Tracks call count regardless of whether the task was created by this instance, + so the endpoint poll-on-read flow works across separate adapter instantiations. + """ + call_count = self._get_task_calls.get(task_display_id, 0) + 1 + self._get_task_calls[task_display_id] = call_count + + task = self._tasks.get(task_display_id) + + if call_count >= 2: + completed = True + status = "completed" + if task is not None: + task["status"] = "completed" + task["completed"] = True + else: + completed = False + status = task["status"] if task is not None else "submitted" + + return C2TaskStatus( + display_id=task_display_id, + status=status, + completed=completed, + command=task["command"] if task is not None else None, + ) + + def get_task_output(self, task_display_id: int) -> str: + """Returns deterministic output once task is completed; raises C2Error before that.""" + # Check call count — completed if get_task was called at least twice. + if self._get_task_calls.get(task_display_id, 0) < 2: + # Also allow tasks in _tasks that were explicitly set to completed. + task = self._tasks.get(task_display_id) + if task is None or not task.get("completed", False): + raise C2Error("task not completed") + + task = self._tasks.get(task_display_id) + command = task["command"] if task is not None else "unknown" + return f"output for task {task_display_id}: {command}\n" + + def list_callback_tasks( + self, + callback_display_id: int, + page: int = 1, + page_size: int = 25, + ) -> C2TaskPage: + all_items = _FAKE_HISTORY.get(callback_display_id, []) + start = (page - 1) * page_size + return C2TaskPage( + items=all_items[start : start + page_size], + total=len(all_items), + page=page, + page_size=page_size, + ) diff --git a/backend/app/services/c2/mapping.py b/backend/app/services/c2/mapping.py new file mode 100644 index 0000000..230b6e6 --- /dev/null +++ b/backend/app/services/c2/mapping.py @@ -0,0 +1,53 @@ +"""C2 task → Simulation output mapping. + +apply_task_to_simulation() implements the full §0.11 contract: + 1. execution_result — append "$ \n\n" block. + 2. executed_at — set from task.completed_at when currently null. + 3. commands — append task.command deduplicated line-by-line. + +Caller is responsible for committing the session. +""" +from __future__ import annotations + +from datetime import UTC, datetime + +from backend.app.models.c2_task import C2Task +from backend.app.models.simulation import Simulation + + +def apply_task_to_simulation(task: C2Task, simulation: Simulation) -> None: + """Apply completed task data to simulation fields per §0.11. + + Idempotent: no-op when task.mapping_applied is already True. + Always sets mapping_applied = True on exit so the task is never re-processed. + """ + if task.mapping_applied: + return + + output = (task.output or "").strip() + + # 1) execution_result — "$ \n\n" block, only when output is non-empty. + if output: + block = f"$ {task.command}\n{output}\n" + existing = simulation.execution_result or "" + if existing: + sep = "" if existing.endswith("\n") else "\n" + simulation.execution_result = existing + sep + block + else: + simulation.execution_result = block + + # 2) executed_at — set once from the first completed task's timestamp. + if simulation.executed_at is None and task.completed_at is not None: + simulation.executed_at = task.completed_at + + # 3) commands — append deduplicated line. + if task.command: + existing_cmds = (simulation.commands or "").splitlines() + if task.command.strip() not in (line.strip() for line in existing_cmds): + if simulation.commands: + simulation.commands = simulation.commands + "\n" + task.command + else: + simulation.commands = task.command + + simulation.updated_at = datetime.now(UTC) + task.mapping_applied = True diff --git a/backend/app/services/c2/mythic.py b/backend/app/services/c2/mythic.py new file mode 100644 index 0000000..2eec290 --- /dev/null +++ b/backend/app/services/c2/mythic.py @@ -0,0 +1,293 @@ +# Contract pinned from MythicMeta/Mythic_Scripting master @ 2026-06-10 (raw.githubusercontent.com/MythicMeta/Mythic_Scripting/master/mythic/mythic.py) +"""Mythic 3.x C2 adapter. + +Transport: POST https://:7443/graphql +Header: apitoken: +Backend: Hasura-proxied Postgres behind nginx. + +M1: test_connection() +M2: list_callbacks(), create_task() +M3: get_task(), get_task_output() +M4: list_callback_tasks() +""" +from __future__ import annotations + +from datetime import datetime + +import requests + +from backend.app.services.c2.adapter import ( + C2Adapter, + C2Callback, + C2Error, + C2Health, + C2HistoricalTask, + C2TaskPage, + C2TaskStatus, + decode_response_text, +) + +_HEALTH_QUERY = "{ __typename }" + +_CALLBACKS_QUERY = """ +query { + callback(order_by: {id: asc}, where: {active: {_eq: true}}) { + id + display_id + active + host + user + domain + last_checkin + } +} +""" + +_CREATE_TASK_MUTATION = """ +mutation CreateTask($callback_id: Int!, $command: String!, $params: String!) { + createTask( + callback_id: $callback_id, + command: $command, + params: $params, + tasking_location: "command_line" + ) { + id + display_id + error + } +} +""" + +_GET_TASK_QUERY = """ +query GetTask($display_id: Int!) { + task(where: {display_id: {_eq: $display_id}}) { + display_id + command_name + status + completed + timestamp + } +} +""" + +_LIST_CALLBACK_TASKS_QUERY = """ +query ListCallbackTasks($callback_display_id: Int!, $limit: Int!, $offset: Int!) { + task( + where: {callback: {display_id: {_eq: $callback_display_id}}} + order_by: {id: desc} + limit: $limit + offset: $offset + ) { + display_id + command_name + params + status + completed + timestamp + } +} +""" + +_COUNT_CALLBACK_TASKS_QUERY = """ +query CountCallbackTasks($callback_display_id: Int!) { + task_aggregate(where: {callback: {display_id: {_eq: $callback_display_id}}}) { + aggregate { + count + } + } +} +""" + +_GET_TASK_OUTPUT_QUERY = """ +query GetTaskOutput($display_id: Int!) { + response( + where: {task: {display_id: {_eq: $display_id}}} + order_by: {id: asc} + ) { + response_text + } +} +""" + + +class MythicAdapter(C2Adapter): + """Real Mythic 3.x adapter using GraphQL over HTTP.""" + + def __init__(self, url: str, api_token: str, verify_tls: bool = True) -> None: + self._url = url.rstrip("/") + "/graphql" + self._token = api_token + self._verify = verify_tls + + def _headers(self) -> dict[str, str]: + return { + "Content-Type": "application/json", + "apitoken": self._token, + } + + def _post(self, body: dict) -> dict: + resp = requests.post( + self._url, + json=body, + headers=self._headers(), + verify=self._verify, + timeout=10, + allow_redirects=False, + ) + resp.raise_for_status() + return resp.json() + + def test_connection(self) -> C2Health: + """POST a trivial introspection query to verify reachability and token validity.""" + try: + resp = requests.post( + self._url, + json={"query": _HEALTH_QUERY}, + headers=self._headers(), + verify=self._verify, + timeout=10, + allow_redirects=False, + ) + if resp.status_code == 200: + return C2Health(ok=True) + return C2Health(ok=False, error=f"HTTP {resp.status_code}") + except requests.RequestException as exc: + return C2Health(ok=False, error=str(exc)) + + def list_callbacks(self) -> list[C2Callback]: + """Return active callbacks from Mythic (filtered server-side: active=true).""" + try: + data = self._post({"query": _CALLBACKS_QUERY}) + except requests.RequestException as exc: + raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc + + callbacks_raw = data.get("data", {}).get("callback", []) + return [ + C2Callback( + display_id=cb["display_id"], + active=cb["active"], + host=cb.get("host") or "", + user=cb.get("user") or "", + domain=cb.get("domain") or "", + last_checkin=cb.get("last_checkin") or "", + ) + for cb in callbacks_raw + ] + + def create_task( + self, + callback_display_id: int, + command: str, + params: str | None = None, + ) -> int: + """Issue a task on a callback; return Mythic task display_id.""" + try: + data = self._post({ + "query": _CREATE_TASK_MUTATION, + "variables": { + "callback_id": callback_display_id, + "command": command, + "params": params or "", + }, + }) + except requests.RequestException as exc: + raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc + + task_data = data.get("data", {}).get("createTask", {}) + error_msg = task_data.get("error") + if error_msg: + raise C2Error(error_msg) + return int(task_data["display_id"]) + + def get_task(self, task_display_id: int) -> C2TaskStatus: + """Return current task status from Mythic.""" + try: + data = self._post({ + "query": _GET_TASK_QUERY, + "variables": {"display_id": task_display_id}, + }) + except requests.RequestException as exc: + raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc + + rows = data.get("data", {}).get("task", []) + if not rows: + raise C2Error(f"task {task_display_id} not found in Mythic") + row = rows[0] + + completed_at: datetime | None = None + if row.get("completed") and row.get("timestamp"): + try: + completed_at = datetime.fromisoformat( + row["timestamp"].replace("Z", "+00:00") + ) + except ValueError: + completed_at = None + + return C2TaskStatus( + display_id=row["display_id"], + status=row["status"], + completed=bool(row.get("completed", False)), + completed_at=completed_at, + command=row.get("command_name") or None, + ) + + def get_task_output(self, task_display_id: int) -> str: + """Return decoded, concatenated output for a task.""" + try: + data = self._post({ + "query": _GET_TASK_OUTPUT_QUERY, + "variables": {"display_id": task_display_id}, + }) + except requests.RequestException as exc: + raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc + + rows = data.get("data", {}).get("response", []) + return "".join( + decode_response_text(r["response_text"]) + for r in rows + if r.get("response_text") + ) + + def list_callback_tasks( + self, + callback_display_id: int, + page: int = 1, + page_size: int = 25, + ) -> C2TaskPage: + """Return a paginated, most-recent-first history of tasks for a callback.""" + offset = (page - 1) * page_size + try: + data = self._post({ + "query": _LIST_CALLBACK_TASKS_QUERY, + "variables": { + "callback_display_id": callback_display_id, + "limit": page_size, + "offset": offset, + }, + }) + count_data = self._post({ + "query": _COUNT_CALLBACK_TASKS_QUERY, + "variables": {"callback_display_id": callback_display_id}, + }) + except requests.RequestException as exc: + raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc + + rows = data.get("data", {}).get("task", []) + total: int = ( + count_data.get("data", {}) + .get("task_aggregate", {}) + .get("aggregate", {}) + .get("count", 0) + ) + + items = [ + C2HistoricalTask( + display_id=r["display_id"], + command=r.get("command_name") or "", + params=r.get("params") or None, + status=r.get("status") or "", + completed=bool(r.get("completed", False)), + timestamp=r.get("timestamp") or None, + ) + for r in rows + ] + return C2TaskPage(items=items, total=total, page=page, page_size=page_size) diff --git a/backend/app/services/crypto.py b/backend/app/services/crypto.py new file mode 100644 index 0000000..4c13854 --- /dev/null +++ b/backend/app/services/crypto.py @@ -0,0 +1,40 @@ +"""Fernet-based encryption service for sensitive fields. + +Key is read from the MIMIC_ENCRYPTION_KEY env var (Fernet base64-urlsafe 32-byte key). +When the key is absent the service raises C2Disabled so callers can return 503. +The key is never logged or returned in any response. +""" +from __future__ import annotations + +import os + +from cryptography.fernet import Fernet, InvalidToken + + +class C2Disabled(Exception): + """Raised when MIMIC_ENCRYPTION_KEY is not set.""" + + +def _get_fernet() -> Fernet: + key = os.environ.get("MIMIC_ENCRYPTION_KEY") + if not key: + raise C2Disabled("C2 disabled: MIMIC_ENCRYPTION_KEY not set") + return Fernet(key.encode() if isinstance(key, str) else key) + + +def encrypt(plaintext: str) -> str: + """Encrypt *plaintext* and return a Fernet token (str).""" + f = _get_fernet() + return f.encrypt(plaintext.encode()).decode() + + +def decrypt(ciphertext: str) -> str: + """Decrypt a Fernet token and return the plaintext string.""" + f = _get_fernet() + try: + return f.decrypt(ciphertext.encode()).decode() + except InvalidToken as exc: + raise ValueError("Invalid ciphertext") from exc + + +__all__ = ["C2Disabled", "encrypt", "decrypt"] diff --git a/backend/app/services/simulation_workflow.py b/backend/app/services/simulation_workflow.py index cb0b1c5..daad7fe 100644 --- a/backend/app/services/simulation_workflow.py +++ b/backend/app/services/simulation_workflow.py @@ -98,6 +98,19 @@ def _maybe_activate_engagement(simulation: Simulation) -> None: db.session.add(engagement) +def promote_to_in_progress(simulation: Simulation) -> None: + """Transition simulation pending → in_progress if it is currently pending. + + Also advances the engagement planned → active via _maybe_activate_engagement. + No-op when the simulation is already in any other status. + Caller must commit. + """ + if simulation.status == SimulationStatus.PENDING: + simulation.status = SimulationStatus.IN_PROGRESS + simulation.updated_at = datetime.now(UTC) + _maybe_activate_engagement(simulation) + + def apply_patch( simulation: Simulation, payload: dict[str, Any], user: User ) -> tuple[Any, int] | None: diff --git a/backend/migrations/versions/0006_c2_layer.py b/backend/migrations/versions/0006_c2_layer.py new file mode 100644 index 0000000..619a394 --- /dev/null +++ b/backend/migrations/versions/0006_c2_layer.py @@ -0,0 +1,67 @@ +"""create c2_config and c2_task tables + +Revision ID: 0006 +Revises: 0005 +Create Date: 2026-06-10 00:00:00.000000 +""" +import sqlalchemy as sa +from alembic import op + +revision = "0006" +down_revision = "0005" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "c2_config", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "engagement_id", + sa.Integer(), + sa.ForeignKey("engagements.id", ondelete="CASCADE"), + nullable=False, + unique=True, + ), + sa.Column("url", sa.Text(), nullable=False), + sa.Column("api_token_encrypted", sa.Text(), nullable=False), + sa.Column("verify_tls", sa.Boolean(), nullable=False, server_default=sa.true()), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("updated_at", sa.DateTime(), nullable=True), + ) + op.create_index("ix_c2_config_engagement_id", "c2_config", ["engagement_id"]) + + op.create_table( + "c2_task", + sa.Column("id", sa.Integer(), primary_key=True), + sa.Column( + "simulation_id", + sa.Integer(), + sa.ForeignKey("simulations.id", ondelete="CASCADE"), + nullable=False, + index=True, + ), + sa.Column("mythic_task_display_id", sa.Integer(), nullable=False), + sa.Column("callback_display_id", sa.Integer(), nullable=False), + sa.Column("command", sa.Text(), nullable=False), + sa.Column("params", sa.Text(), nullable=True), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("completed", sa.Boolean(), nullable=False, server_default=sa.false()), + sa.Column("output", sa.Text(), nullable=True), + sa.Column( + "source", + sa.Enum("mimic", "import", name="c2task_source"), + nullable=False, + ), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.Column("completed_at", sa.DateTime(), nullable=True), + ) + + +def downgrade() -> None: + op.drop_table("c2_task") + op.drop_index("ix_c2_config_engagement_id", "c2_config") + op.drop_table("c2_config") + # Remove the enum type (no-op on SQLite, required on Postgres) + sa.Enum(name="c2task_source").drop(op.get_bind(), checkfirst=True) diff --git a/backend/migrations/versions/0007_c2_task_mapping_applied.py b/backend/migrations/versions/0007_c2_task_mapping_applied.py new file mode 100644 index 0000000..ff60ea4 --- /dev/null +++ b/backend/migrations/versions/0007_c2_task_mapping_applied.py @@ -0,0 +1,30 @@ +"""add mapping_applied column to c2_task + +Revision ID: 0007 +Revises: 0006 +Create Date: 2026-06-10 00:00:00.000000 +""" +import sqlalchemy as sa +from alembic import op + +revision = "0007" +down_revision = "0006" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("c2_task") as batch_op: + batch_op.add_column( + sa.Column( + "mapping_applied", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ) + ) + + +def downgrade() -> None: + with op.batch_alter_table("c2_task") as batch_op: + batch_op.drop_column("mapping_applied") diff --git a/backend/requirements.txt b/backend/requirements.txt index 878005e..f22875d 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -4,6 +4,10 @@ Flask-Migrate==4.0.7 PyJWT==2.9.0 argon2-cffi==23.1.0 weasyprint>=60.0 +cryptography==44.0.0 +requests==2.32.3 pytest==8.3.3 ruff==0.6.9 mypy==1.11.2 +types-requests==2.32.0.20240914 +requests-mock==1.12.1 diff --git a/backend/tests/test_c2_adapter_fake.py b/backend/tests/test_c2_adapter_fake.py new file mode 100644 index 0000000..ca464cf --- /dev/null +++ b/backend/tests/test_c2_adapter_fake.py @@ -0,0 +1,30 @@ +"""Tests for the FakeAdapter deterministic in-memory implementation.""" +from __future__ import annotations + +from backend.app.services.c2.adapter import C2Health +from backend.app.services.c2.fake import FakeAdapter + + +class TestFakeAdapterTestConnection: + def test_returns_ok_true(self): + adapter = FakeAdapter() + health = adapter.test_connection() + assert isinstance(health, C2Health) + assert health.ok is True + assert health.error is None + + def test_list_callbacks_returns_list(self): + adapter = FakeAdapter() + callbacks = adapter.list_callbacks() + assert isinstance(callbacks, list) + assert len(callbacks) >= 1 + + def test_list_callbacks_fields(self): + adapter = FakeAdapter() + cb = adapter.list_callbacks()[0] + assert hasattr(cb, "display_id") + assert hasattr(cb, "active") + assert hasattr(cb, "host") + assert hasattr(cb, "user") + assert hasattr(cb, "domain") + assert hasattr(cb, "last_checkin") diff --git a/backend/tests/test_c2_adapter_fake_m2.py b/backend/tests/test_c2_adapter_fake_m2.py new file mode 100644 index 0000000..3da42a4 --- /dev/null +++ b/backend/tests/test_c2_adapter_fake_m2.py @@ -0,0 +1,62 @@ +"""FakeAdapter M2 tests — list_callbacks shape, create_task monotonicity.""" +from __future__ import annotations + +from backend.app.services.c2.fake import FakeAdapter + + +class TestFakeAdapterListCallbacks: + def test_returns_three_callbacks(self): + adapter = FakeAdapter() + callbacks = adapter.list_callbacks() + assert len(callbacks) == 3 + + def test_all_active(self): + adapter = FakeAdapter() + for cb in adapter.list_callbacks(): + assert cb.active is True + + def test_display_ids_are_1_2_3(self): + adapter = FakeAdapter() + ids = [cb.display_id for cb in adapter.list_callbacks()] + assert ids == [1, 2, 3] + + def test_pinned_last_checkin_format(self): + adapter = FakeAdapter() + for cb in adapter.list_callbacks(): + assert cb.last_checkin.startswith("2026-06-10") + + def test_callbacks_have_host_user_domain(self): + adapter = FakeAdapter() + for cb in adapter.list_callbacks(): + assert cb.host + assert cb.user + assert cb.domain + + +class TestFakeAdapterCreateTask: + def test_returns_monotonic_ids_from_1000(self): + adapter = FakeAdapter() + id1 = adapter.create_task(1, "whoami") + id2 = adapter.create_task(1, "ipconfig") + assert id1 == 1000 + assert id2 == 1001 + + def test_separate_instances_start_at_1000_independently(self): + a1 = FakeAdapter() + a2 = FakeAdapter() + assert a1.create_task(1, "cmd") == 1000 + assert a2.create_task(1, "cmd") == 1000 + + def test_stores_command_and_callback(self): + adapter = FakeAdapter() + tid = adapter.create_task(callback_display_id=2, command="ls", params="-la") + task = adapter._tasks[tid] + assert task["command"] == "ls" + assert task["params"] == "-la" + assert task["callback_display_id"] == 2 + + def test_initial_status_submitted(self): + adapter = FakeAdapter() + tid = adapter.create_task(1, "hostname") + assert adapter._tasks[tid]["status"] == "submitted" + assert adapter._tasks[tid]["completed"] is False diff --git a/backend/tests/test_c2_adapter_fake_m3.py b/backend/tests/test_c2_adapter_fake_m3.py new file mode 100644 index 0000000..3da8e9a --- /dev/null +++ b/backend/tests/test_c2_adapter_fake_m3.py @@ -0,0 +1,110 @@ +"""FakeAdapter M3 state-progression tests — get_task and get_task_output.""" +from __future__ import annotations + +import pytest + +from backend.app.services.c2.adapter import C2Error +from backend.app.services.c2.fake import FakeAdapter + + +@pytest.fixture() +def adapter() -> FakeAdapter: + return FakeAdapter() + + +@pytest.fixture() +def adapter_with_task(adapter: FakeAdapter) -> tuple[FakeAdapter, int]: + tid = adapter.create_task(callback_display_id=1, command="whoami") + return adapter, tid + + +class TestFakeAdapterGetTaskProgression: + def test_first_call_returns_submitted(self, adapter_with_task): + a, tid = adapter_with_task + status = a.get_task(tid) + assert status.status == "submitted" + assert status.completed is False + + def test_second_call_returns_completed(self, adapter_with_task): + a, tid = adapter_with_task + a.get_task(tid) # first call + status = a.get_task(tid) # second call + assert status.status == "completed" + assert status.completed is True + + def test_subsequent_calls_stay_completed(self, adapter_with_task): + a, tid = adapter_with_task + for _ in range(5): + a.get_task(tid) + status = a.get_task(tid) + assert status.completed is True + + def test_unknown_task_id_returns_submitted_on_first_call(self, adapter): + """A task ID not created by this instance still goes through submitted→completed.""" + status = adapter.get_task(9999) + assert status.display_id == 9999 + assert status.status == "submitted" + assert status.completed is False + + def test_call_counters_are_per_task(self, adapter): + """Two tasks have independent state — completing one does not affect the other.""" + t1 = adapter.create_task(callback_display_id=1, command="whoami") + t2 = adapter.create_task(callback_display_id=1, command="ipconfig") + + # Advance t1 to completed via two calls. + adapter.get_task(t1) + adapter.get_task(t1) + + # t2 first call should still be submitted. + s2 = adapter.get_task(t2) + assert s2.status == "submitted" + assert s2.completed is False + + def test_instances_are_isolated(self): + """Per-instance counters — different FakeAdapter instances don't share state.""" + a1 = FakeAdapter() + a2 = FakeAdapter() + + t1 = a1.create_task(1, "cmd") + t2 = a2.create_task(1, "cmd") + + a1.get_task(t1) + a1.get_task(t1) # a1's task is now completed + + # a2's task with same display_id (both start at 1000) should be independent. + assert t1 == t2 == 1000 + s2 = a2.get_task(t2) + assert s2.status == "submitted" + + +class TestFakeAdapterGetTaskOutput: + def test_raises_before_completed(self, adapter_with_task): + a, tid = adapter_with_task + with pytest.raises(C2Error, match="task not completed"): + a.get_task_output(tid) + + def test_raises_after_first_get_task_call_only(self, adapter_with_task): + a, tid = adapter_with_task + a.get_task(tid) # first call — still submitted + with pytest.raises(C2Error, match="task not completed"): + a.get_task_output(tid) + + def test_returns_output_after_completed(self, adapter_with_task): + a, tid = adapter_with_task + a.get_task(tid) + a.get_task(tid) # now completed + output = a.get_task_output(tid) + assert "whoami" in output + assert str(tid) in output + + def test_output_format(self, adapter): + tid = adapter.create_task(callback_display_id=2, command="ipconfig /all") + adapter.get_task(tid) + adapter.get_task(tid) + output = adapter.get_task_output(tid) + assert output == f"output for task {tid}: ipconfig /all\n" + + def test_unknown_task_raises_c2error(self, adapter): + """Task ID never created and never polled — not completed → C2Error.""" + with pytest.raises(C2Error, match="task not completed"): + adapter.get_task_output(9999) diff --git a/backend/tests/test_c2_adapter_fake_m4.py b/backend/tests/test_c2_adapter_fake_m4.py new file mode 100644 index 0000000..e1e88ba --- /dev/null +++ b/backend/tests/test_c2_adapter_fake_m4.py @@ -0,0 +1,75 @@ +"""FakeAdapter M4 tests — list_callback_tasks pagination.""" +from __future__ import annotations + +import pytest + +from backend.app.services.c2.adapter import C2HistoricalTask +from backend.app.services.c2.fake import FakeAdapter + + +@pytest.fixture() +def adapter() -> FakeAdapter: + return FakeAdapter() + + +class TestFakeAdapterListCallbackTasks: + def test_callback_1_returns_12_total(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=1, page=1, page_size=25) + assert page.total == 12 + + def test_callback_2_returns_0_tasks(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=2, page=1, page_size=25) + assert page.total == 0 + assert page.items == [] + + def test_callback_3_returns_5_tasks(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=3, page=1, page_size=25) + assert page.total == 5 + assert len(page.items) == 5 + + def test_items_are_c2_historical_task_instances(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=1, page=1, page_size=5) + for item in page.items: + assert isinstance(item, C2HistoricalTask) + + def test_pagination_page1(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=1, page=1, page_size=5) + assert len(page.items) == 5 + assert page.page == 1 + assert page.page_size == 5 + + def test_pagination_page2(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=1, page=2, page_size=5) + assert len(page.items) == 5 + assert page.page == 2 + + def test_pagination_last_page_partial(self, adapter): + # 12 tasks, page_size=5 → page 3 has 2 items. + page = adapter.list_callback_tasks(callback_display_id=1, page=3, page_size=5) + assert len(page.items) == 2 + assert page.total == 12 + + def test_pagination_beyond_range_returns_empty(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=1, page=99, page_size=25) + assert len(page.items) == 0 + assert page.total == 12 + + def test_history_is_deterministic_across_instances(self): + a1 = FakeAdapter() + a2 = FakeAdapter() + p1 = a1.list_callback_tasks(callback_display_id=1, page=1, page_size=25) + p2 = a2.list_callback_tasks(callback_display_id=1, page=1, page_size=25) + assert [t.display_id for t in p1.items] == [t.display_id for t in p2.items] + + def test_completed_and_submitted_mix(self, adapter): + """Callback 1 has alternating completed/submitted tasks (even=completed).""" + page = adapter.list_callback_tasks(callback_display_id=1, page=1, page_size=12) + completed = [t for t in page.items if t.completed] + submitted = [t for t in page.items if not t.completed] + assert len(completed) == 6 + assert len(submitted) == 6 + + def test_unknown_callback_returns_empty(self, adapter): + page = adapter.list_callback_tasks(callback_display_id=999, page=1, page_size=25) + assert page.total == 0 + assert page.items == [] diff --git a/backend/tests/test_c2_adapter_mythic.py b/backend/tests/test_c2_adapter_mythic.py new file mode 100644 index 0000000..de32055 --- /dev/null +++ b/backend/tests/test_c2_adapter_mythic.py @@ -0,0 +1,151 @@ +"""MythicAdapter unit tests — mocked HTTP with requests-mock.""" +from __future__ import annotations + +import pytest +import requests +import requests_mock as rm_module + +from backend.app.services.c2.adapter import C2Error +from backend.app.services.c2.mythic import MythicAdapter + +_BASE_URL = "https://mythic.lab:7443" +_GQL_URL = _BASE_URL + "/graphql" +_TOKEN = "fake-api-token" + + +@pytest.fixture() +def adapter(): + return MythicAdapter(url=_BASE_URL, api_token=_TOKEN, verify_tls=False) + + +class TestMythicAdapterListCallbacks: + def test_returns_callbacks_from_graphql(self, adapter): + payload = { + "data": { + "callback": [ + { + "id": 1, + "display_id": 1, + "active": True, + "host": "HOST-01", + "user": "jdoe", + "domain": "LAB", + "last_checkin": "2026-06-10T00:00:00Z", + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + callbacks = adapter.list_callbacks() + + assert len(callbacks) == 1 + assert callbacks[0].display_id == 1 + assert callbacks[0].host == "HOST-01" + assert callbacks[0].user == "jdoe" + + def test_sends_apitoken_header(self, adapter): + payload = {"data": {"callback": []}} + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + adapter.list_callbacks() + sent_headers = m.last_request.headers + + assert sent_headers.get("apitoken") == _TOKEN + + def test_verify_tls_flag_passed(self): + """Adapter with verify_tls=True should pass verify=True to requests.""" + adapter_tls = MythicAdapter(url=_BASE_URL, api_token=_TOKEN, verify_tls=True) + payload = {"data": {"callback": []}} + # requests-mock intercepts before TLS — just confirm no error path triggered. + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + callbacks = adapter_tls.list_callbacks() + assert isinstance(callbacks, list) + + def test_network_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.ConnectionError("connection refused")) + with pytest.raises(C2Error): + adapter.list_callbacks() + + def test_http_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, status_code=500, text="Internal Server Error") + with pytest.raises(C2Error): + adapter.list_callbacks() + + +class TestMythicAdapterCreateTask: + def test_returns_display_id_on_success(self, adapter): + payload = { + "data": { + "createTask": { + "id": 42, + "display_id": 7, + "error": None, + } + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + tid = adapter.create_task(callback_display_id=1, command="whoami") + + assert tid == 7 + + def test_sends_apitoken_header(self, adapter): + payload = {"data": {"createTask": {"id": 1, "display_id": 1, "error": None}}} + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + adapter.create_task(1, "cmd") + sent_headers = m.last_request.headers + + assert sent_headers.get("apitoken") == _TOKEN + + def test_error_field_raises_c2error(self, adapter): + payload = { + "data": { + "createTask": { + "id": None, + "display_id": None, + "error": "callback not found", + } + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + with pytest.raises(C2Error, match="callback not found"): + adapter.create_task(1, "whoami") + + def test_network_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.Timeout("timeout")) + with pytest.raises(C2Error): + adapter.create_task(1, "whoami") + + +class TestMythicAdapterErrorSanitization: + def test_connection_error_message_does_not_contain_url(self, adapter): + """C2Error message must not expose the configured Mythic URL.""" + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.ConnectionError( + f"HTTPSConnectionPool(host='{_BASE_URL}', port=7443): Max retries exceeded" + )) + with pytest.raises(C2Error) as exc_info: + adapter.list_callbacks() + + assert _BASE_URL not in str(exc_info.value) + assert "ConnectionError" in str(exc_info.value) + + +class TestMythicAdapterNoRedirects: + def test_does_not_follow_redirect(self, adapter): + """Adapter must not follow HTTP redirects (allow_redirects=False).""" + with rm_module.Mocker() as m: + # Simulate a redirect response; requests-mock won't auto-follow it. + m.post(_GQL_URL, status_code=301, headers={"Location": "https://evil.example/graphql"}) + # With allow_redirects=False the 301 is treated as a non-2xx → raise_for_status raises. + with pytest.raises(C2Error): + adapter.list_callbacks() + # Exactly one request was made — no follow-up to Location. + assert len(m.request_history) == 1 diff --git a/backend/tests/test_c2_adapter_mythic_m3.py b/backend/tests/test_c2_adapter_mythic_m3.py new file mode 100644 index 0000000..e6d5963 --- /dev/null +++ b/backend/tests/test_c2_adapter_mythic_m3.py @@ -0,0 +1,188 @@ +"""MythicAdapter M3 tests — get_task and get_task_output, mocked HTTP.""" +from __future__ import annotations + +import pytest +import requests +import requests_mock as rm_module + +from backend.app.services.c2.adapter import C2Error +from backend.app.services.c2.mythic import MythicAdapter + +_BASE_URL = "https://mythic.lab:7443" +_GQL_URL = _BASE_URL + "/graphql" +_TOKEN = "fake-api-token" + + +@pytest.fixture() +def adapter(): + return MythicAdapter(url=_BASE_URL, api_token=_TOKEN, verify_tls=False) + + +class TestMythicAdapterGetTask: + def test_returns_status_for_incomplete_task(self, adapter): + payload = { + "data": { + "task": [ + { + "display_id": 7, + "status": "processing", + "completed": False, + "timestamp": None, + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + status = adapter.get_task(7) + + assert status.display_id == 7 + assert status.status == "processing" + assert status.completed is False + assert status.completed_at is None + + def test_returns_completed_at_for_completed_task(self, adapter): + payload = { + "data": { + "task": [ + { + "display_id": 7, + "status": "completed", + "completed": True, + "timestamp": "2026-06-10T12:00:00Z", + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + status = adapter.get_task(7) + + assert status.completed is True + assert status.completed_at is not None + assert status.completed_at.year == 2026 + + def test_raises_when_task_not_found(self, adapter): + payload = {"data": {"task": []}} + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + with pytest.raises(C2Error, match="not found"): + adapter.get_task(999) + + def test_sends_apitoken_header(self, adapter): + payload = { + "data": { + "task": [ + {"display_id": 1, "status": "submitted", "completed": False, "timestamp": None} + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + adapter.get_task(1) + assert m.last_request.headers.get("apitoken") == _TOKEN + + def test_network_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.ConnectionError("refused")) + with pytest.raises(C2Error): + adapter.get_task(1) + + def test_no_redirect_followed(self, adapter): + """get_task must not follow HTTP redirects.""" + with rm_module.Mocker() as m: + m.post(_GQL_URL, status_code=301, headers={"Location": "https://evil.example/"}) + with pytest.raises(C2Error): + adapter.get_task(1) + assert len(m.request_history) == 1 + + def test_invalid_timestamp_does_not_crash(self, adapter): + """A malformed timestamp field falls back to completed_at=None without raising.""" + payload = { + "data": { + "task": [ + { + "display_id": 5, + "status": "completed", + "completed": True, + "timestamp": "not-a-date", + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + status = adapter.get_task(5) + + assert status.completed is True + assert status.completed_at is None + + +class TestMythicAdapterGetTaskOutput: + def test_returns_decoded_output(self, adapter): + import base64 + encoded = base64.b64encode(b"Administrator\r\n").decode() + payload = { + "data": { + "response": [{"response_text": encoded}] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + output = adapter.get_task_output(7) + + assert "Administrator" in output + + def test_concatenates_multiple_responses(self, adapter): + import base64 + r1 = base64.b64encode(b"line one\n").decode() + r2 = base64.b64encode(b"line two\n").decode() + payload = { + "data": { + "response": [{"response_text": r1}, {"response_text": r2}] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + output = adapter.get_task_output(7) + + assert "line one" in output + assert "line two" in output + + def test_returns_empty_string_when_no_responses(self, adapter): + payload = {"data": {"response": []}} + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + output = adapter.get_task_output(7) + + assert output == "" + + def test_skips_empty_response_text(self, adapter): + import base64 + encoded = base64.b64encode(b"real output").decode() + payload = { + "data": { + "response": [ + {"response_text": ""}, + {"response_text": encoded}, + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + output = adapter.get_task_output(7) + + assert output == "real output" + + def test_network_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.Timeout("timeout")) + with pytest.raises(C2Error): + adapter.get_task_output(7) + + def test_no_redirect_followed(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, status_code=302, headers={"Location": "https://evil.example/"}) + with pytest.raises(C2Error): + adapter.get_task_output(1) + assert len(m.request_history) == 1 diff --git a/backend/tests/test_c2_adapter_mythic_m4.py b/backend/tests/test_c2_adapter_mythic_m4.py new file mode 100644 index 0000000..5103185 --- /dev/null +++ b/backend/tests/test_c2_adapter_mythic_m4.py @@ -0,0 +1,167 @@ +"""MythicAdapter M4 tests — list_callback_tasks, mocked HTTP.""" +from __future__ import annotations + +import pytest +import requests +import requests_mock as rm_module + +from backend.app.services.c2.adapter import C2Error, C2HistoricalTask +from backend.app.services.c2.mythic import MythicAdapter + +_BASE_URL = "https://mythic.lab:7443" +_GQL_URL = _BASE_URL + "/graphql" +_TOKEN = "fake-api-token" + + +@pytest.fixture() +def adapter(): + return MythicAdapter(url=_BASE_URL, api_token=_TOKEN, verify_tls=False) + + +def _task_list_payload(tasks: list[dict]) -> dict: + return {"data": {"task": tasks}} + + +def _count_payload(count: int) -> dict: + return {"data": {"task_aggregate": {"aggregate": {"count": count}}}} + + +class TestMythicAdapterListCallbackTasks: + def test_returns_tasks_from_graphql(self, adapter): + tasks_payload = _task_list_payload([ + { + "display_id": 7, + "command_name": "whoami", + "params": "", + "status": "completed", + "completed": True, + "timestamp": "2026-06-10T12:00:00Z", + } + ]) + count_payload = _count_payload(1) + + with rm_module.Mocker() as m: + m.post(_GQL_URL, [{"json": tasks_payload}, {"json": count_payload}]) + page = adapter.list_callback_tasks(callback_display_id=1, page=1, page_size=25) + + assert page.total == 1 + assert len(page.items) == 1 + item = page.items[0] + assert isinstance(item, C2HistoricalTask) + assert item.display_id == 7 + assert item.command == "whoami" + assert item.completed is True + + def test_pagination_offset_calculation(self, adapter): + """page=2, page_size=10 → offset=10 must be sent to Mythic.""" + tasks_payload = _task_list_payload([]) + count_payload = _count_payload(0) + + with rm_module.Mocker() as m: + m.post(_GQL_URL, [{"json": tasks_payload}, {"json": count_payload}]) + adapter.list_callback_tasks(callback_display_id=1, page=2, page_size=10) + + # First request is the task list; check variables. + first_body = m.request_history[0].json() + variables = first_body.get("variables", {}) + + assert variables.get("offset") == 10 + assert variables.get("limit") == 10 + + def test_sends_apitoken_header(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, [ + {"json": _task_list_payload([])}, + {"json": _count_payload(0)}, + ]) + adapter.list_callback_tasks(callback_display_id=1) + for req in m.request_history: + assert req.headers.get("apitoken") == _TOKEN + + def test_empty_task_list(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, [ + {"json": _task_list_payload([])}, + {"json": _count_payload(0)}, + ]) + page = adapter.list_callback_tasks(callback_display_id=1) + + assert page.total == 0 + assert page.items == [] + + def test_network_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, exc=requests.exceptions.ConnectionError("refused")) + with pytest.raises(C2Error): + adapter.list_callback_tasks(callback_display_id=1) + + def test_http_error_raises_c2error(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, status_code=500, text="error") + with pytest.raises(C2Error): + adapter.list_callback_tasks(callback_display_id=1) + + def test_no_redirect_followed(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, status_code=301, headers={"Location": "https://evil.example/"}) + with pytest.raises(C2Error): + adapter.list_callback_tasks(callback_display_id=1) + # Both requests (tasks + count) should each only make one attempt. + for req in m.request_history: + assert req.method == "POST" + + def test_page_and_page_size_in_response(self, adapter): + with rm_module.Mocker() as m: + m.post(_GQL_URL, [ + {"json": _task_list_payload([])}, + {"json": _count_payload(50)}, + ]) + page = adapter.list_callback_tasks(callback_display_id=1, page=3, page_size=10) + + assert page.page == 3 + assert page.page_size == 10 + assert page.total == 50 + + +class TestMythicAdapterGetTaskCommandField: + """Ensure command_name is surfaced via get_task() C2TaskStatus.command.""" + + def test_get_task_returns_command(self, adapter): + payload = { + "data": { + "task": [ + { + "display_id": 7, + "command_name": "shell", + "status": "completed", + "completed": True, + "timestamp": "2026-06-10T12:00:00Z", + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + status = adapter.get_task(7) + + assert status.command == "shell" + + def test_get_task_command_none_when_missing(self, adapter): + payload = { + "data": { + "task": [ + { + "display_id": 7, + "command_name": None, + "status": "submitted", + "completed": False, + "timestamp": None, + } + ] + } + } + with rm_module.Mocker() as m: + m.post(_GQL_URL, json=payload) + status = adapter.get_task(7) + + assert status.command is None diff --git a/backend/tests/test_c2_callbacks.py b/backend/tests/test_c2_callbacks.py new file mode 100644 index 0000000..5f35b54 --- /dev/null +++ b/backend/tests/test_c2_callbacks.py @@ -0,0 +1,142 @@ +"""Tests for GET /api/engagements//c2/callbacks.""" +from __future__ import annotations + +import pytest +from cryptography.fernet import Fernet +from flask.testing import FlaskClient + +from backend.app.services.c2.adapter import C2Error +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") + + +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 + + +class TestGetCallbacksHappyPath: + def test_returns_3_callbacks_with_fake_adapter( + self, client: FlaskClient, admin_token: str + ) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(admin_token), + ) + assert resp.status_code == 200 + body = resp.get_json() + assert "callbacks" in body + assert len(body["callbacks"]) == 3 + + def test_callback_shape(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(admin_token), + ) + cb = resp.get_json()["callbacks"][0] + assert "display_id" in cb + assert "active" in cb + assert "host" in cb + assert "user" in cb + assert "domain" in cb + assert "last_checkin" in cb + + def test_redteam_allowed( + self, client: FlaskClient, admin_token: str, redteam_token: str + ) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(redteam_token), + ) + assert resp.status_code == 200 + + +class TestGetCallbacksErrorCases: + def test_404_when_no_config(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(admin_token), + ) + assert resp.status_code == 404 + + def test_404_engagement_not_found(self, client: FlaskClient, admin_token: str) -> None: + resp = client.get( + "/api/engagements/9999/c2/callbacks", + headers=_h(admin_token), + ) + assert resp.status_code == 404 + + def test_403_soc( + self, client: FlaskClient, admin_token: str, soc_token: str + ) -> None: + eng = _make_engagement(client, admin_token) + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(soc_token), + ) + 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) + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(admin_token), + ) + assert resp.status_code == 503 + + def test_502_when_adapter_raises( + self, monkeypatch, client: FlaskClient, admin_token: str + ) -> None: + from backend.app.services.c2 import fake as fake_mod + + def _boom(self): + raise C2Error("mythic unreachable") + + monkeypatch.setattr(fake_mod.FakeAdapter, "list_callbacks", _boom) + + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + resp = client.get( + f"/api/engagements/{eng['id']}/c2/callbacks", + headers=_h(admin_token), + ) + assert resp.status_code == 502 + assert "mythic unreachable" in resp.get_json().get("error", "") diff --git a/backend/tests/test_c2_config.py b/backend/tests/test_c2_config.py new file mode 100644 index 0000000..fd2c617 --- /dev/null +++ b/backend/tests/test_c2_config.py @@ -0,0 +1,367 @@ +"""Tests for C2 config CRUD endpoints. + +Covers: +- GET 404 when no config exists +- PUT create (api_token required) +- PUT update with omitted token keeps old ciphertext +- GET 200 returns has_token=True, never cleartext +- DELETE 204 +- Cascade delete when engagement is deleted +- RBAC: admin OK / redteam OK / SOC 403 on all 4 endpoints +- 503 guard when MIMIC_ENCRYPTION_KEY is unset +- POST /test with fake adapter +""" +from __future__ import annotations + +import pytest +from cryptography.fernet import Fernet +from flask import Flask +from flask.testing import FlaskClient + +from backend.app.models.c2_config import C2Config +from backend.tests.conftest import auth_headers as _h + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +_FERNET_KEY = Fernet.generate_key().decode() + + +@pytest.fixture(autouse=True) +def set_encryption_key(monkeypatch): + """Default: key is present. Individual tests can override.""" + 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, resp.get_json() + return resp.get_json() + + +def _put_config( + client: FlaskClient, + token: str, + eid: int, + *, + url: str = "https://c2.internal:7443", + api_token: str | None = "s3cr3t", + verify_tls: bool = True, +) -> dict: + payload: dict = {"url": url, "verify_tls": verify_tls} + if api_token is not None: + payload["api_token"] = api_token + resp = client.put( + f"/api/engagements/{eid}/c2-config", + headers=_h(token), + json=payload, + ) + return resp + + +# --------------------------------------------------------------------------- +# GET — 404 when no config +# --------------------------------------------------------------------------- + + +def test_get_config_not_found(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)) + assert resp.status_code == 404 + + +def test_get_config_engagement_not_found(client: FlaskClient, admin_token: str) -> None: + resp = client.get("/api/engagements/9999/c2-config", headers=_h(admin_token)) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# PUT — create +# --------------------------------------------------------------------------- + + +def test_put_rejects_http_scheme(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = _put_config(client, admin_token, eng["id"], url="http://c2.internal:7443") + assert resp.status_code == 400 + assert "https" in resp.get_json().get("error", "").lower() + + +def test_put_rejects_missing_hostname(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + # urlparse("https://:7443") produces an empty hostname + resp = _put_config(client, admin_token, eng["id"], url="https://:7443") + assert resp.status_code == 400 + assert "hostname" in resp.get_json().get("error", "").lower() + + +def test_put_creates_config(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = _put_config(client, admin_token, eng["id"]) + assert resp.status_code == 200 + body = resp.get_json() + assert body["has_token"] is True + assert body["url"] == "https://c2.internal:7443" + assert body["verify_tls"] is True + + +def test_put_create_requires_api_token(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = _put_config(client, admin_token, eng["id"], api_token=None) + assert resp.status_code == 400 + + +def test_put_create_requires_url(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = client.put( + f"/api/engagements/{eng['id']}/c2-config", + headers=_h(admin_token), + json={"api_token": "tok", "verify_tls": True}, + ) + assert resp.status_code == 400 + + +# --------------------------------------------------------------------------- +# PUT — update, omitting api_token preserves old ciphertext +# --------------------------------------------------------------------------- + + +def test_put_update_omits_token_keeps_old( + app: Flask, client: FlaskClient, admin_token: str +) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"], api_token="original-token") + + # Read ciphertext from DB before update. + with app.app_context(): + cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first() + assert cfg is not None + old_cipher = cfg.api_token_encrypted + + # Update URL, omit api_token. + resp = _put_config( + client, admin_token, eng["id"], + url="https://new.internal:7443", api_token=None, + ) + assert resp.status_code == 200 + + with app.app_context(): + cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first() + assert cfg is not None + assert cfg.api_token_encrypted == old_cipher + assert cfg.url == "https://new.internal:7443" + + +def test_put_update_with_token_replaces_ciphertext( + app: Flask, client: FlaskClient, admin_token: str +) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"], api_token="original-token") + + with app.app_context(): + cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first() + assert cfg is not None + old_cipher = cfg.api_token_encrypted + + _put_config(client, admin_token, eng["id"], api_token="new-token") + + with app.app_context(): + cfg = C2Config.query.filter_by(engagement_id=eng["id"]).first() + assert cfg is not None + assert cfg.api_token_encrypted != old_cipher + + +# --------------------------------------------------------------------------- +# GET — 200, has_token=True, never cleartext +# --------------------------------------------------------------------------- + + +def test_get_config_returns_has_token_not_cleartext( + client: FlaskClient, admin_token: str +) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"], api_token="s3cr3t") + + resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)) + assert resp.status_code == 200 + body = resp.get_json() + assert body["has_token"] is True + assert "api_token" not in body + assert "api_token_encrypted" not in body + assert "s3cr3t" not in str(body) + + +def test_get_config_verify_tls_default_true( + client: FlaskClient, admin_token: str +) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)) + assert resp.get_json()["verify_tls"] is True + + +# --------------------------------------------------------------------------- +# DELETE — 204 +# --------------------------------------------------------------------------- + + +def test_delete_config_204(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + resp = client.delete( + f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token) + ) + assert resp.status_code == 204 + + # Subsequent GET returns 404. + resp2 = client.get(f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token)) + assert resp2.status_code == 404 + + +def test_delete_config_not_found(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = client.delete( + f"/api/engagements/{eng['id']}/c2-config", headers=_h(admin_token) + ) + assert resp.status_code == 404 + + +# --------------------------------------------------------------------------- +# CASCADE — delete engagement removes config +# --------------------------------------------------------------------------- + + +def test_cascade_delete_engagement_removes_config( + app: Flask, client: FlaskClient, admin_token: str +) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + with app.app_context(): + assert C2Config.query.filter_by(engagement_id=eng["id"]).count() == 1 + + client.delete(f"/api/engagements/{eng['id']}", headers=_h(admin_token)) + + with app.app_context(): + assert C2Config.query.filter_by(engagement_id=eng["id"]).count() == 0 + + +# --------------------------------------------------------------------------- +# RBAC matrix +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("method,path_suffix", [ + ("GET", "/c2-config"), + ("PUT", "/c2-config"), + ("DELETE", "/c2-config"), + ("POST", "/c2-config/test"), +]) +def test_soc_gets_403( + client: FlaskClient, admin_token: str, soc_token: str, + method: str, path_suffix: str, +) -> None: + eng = _make_engagement(client, admin_token) + url = f"/api/engagements/{eng['id']}{path_suffix}" + resp = getattr(client, method.lower())(url, headers=_h(soc_token), json={}) + assert resp.status_code == 403 + + +@pytest.mark.parametrize("method,path_suffix", [ + ("GET", "/c2-config"), + ("DELETE", "/c2-config"), + ("POST", "/c2-config/test"), +]) +def test_redteam_gets_allowed( + client: FlaskClient, admin_token: str, redteam_token: str, + method: str, path_suffix: str, +) -> None: + eng = _make_engagement(client, admin_token) + # Ensure config exists for GET/DELETE/test. + _put_config(client, admin_token, eng["id"]) + url = f"/api/engagements/{eng['id']}{path_suffix}" + resp = getattr(client, method.lower())(url, headers=_h(redteam_token), json={}) + # Not 403 and not 401. + assert resp.status_code not in (401, 403) + + +def test_redteam_can_put_config( + client: FlaskClient, admin_token: str, redteam_token: str, +) -> None: + eng = _make_engagement(client, admin_token) + resp = _put_config(client, redteam_token, eng["id"]) + assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# 503 guard when MIMIC_ENCRYPTION_KEY is unset +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("method,path_suffix", [ + ("GET", "/c2-config"), + ("PUT", "/c2-config"), + ("DELETE", "/c2-config"), + ("POST", "/c2-config/test"), +]) +def test_503_when_key_unset( + monkeypatch, + client: FlaskClient, + admin_token: str, + method: str, + path_suffix: str, +) -> None: + monkeypatch.delenv("MIMIC_ENCRYPTION_KEY", raising=False) + eng = _make_engagement(client, admin_token) + url = f"/api/engagements/{eng['id']}{path_suffix}" + resp = getattr(client, method.lower())(url, headers=_h(admin_token), json={ + "url": "https://c2", "api_token": "tok", "verify_tls": True, + }) + assert resp.status_code == 503 + assert "MIMIC_ENCRYPTION_KEY" in resp.get_json().get("error", "") + + +# --------------------------------------------------------------------------- +# POST /test — connectivity check via fake adapter +# --------------------------------------------------------------------------- + + +def test_post_test_returns_ok_true(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + + resp = client.post( + f"/api/engagements/{eng['id']}/c2-config/test", + headers=_h(admin_token), + json={}, + ) + assert resp.status_code == 200 + body = resp.get_json() + assert body["ok"] is True + assert body["error"] is None + + +def test_post_test_no_config_returns_404(client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + resp = client.post( + f"/api/engagements/{eng['id']}/c2-config/test", + headers=_h(admin_token), + json={}, + ) + assert resp.status_code == 404 diff --git a/backend/tests/test_c2_execute.py b/backend/tests/test_c2_execute.py new file mode 100644 index 0000000..8b632cc --- /dev/null +++ b/backend/tests/test_c2_execute.py @@ -0,0 +1,324 @@ +"""Tests for POST /api/simulations//c2/execute.""" +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, SimulationStatus +from backend.app.services.c2.adapter import C2Error +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 _advance_to_in_progress(client: FlaskClient, token: str, sid: int) -> None: + client.patch( + f"/api/simulations/{sid}", + headers=_h(token), + json={"name": "Sim Alpha"}, + ) + + +def _advance_to_review_required(client: FlaskClient, token: str, sid: int) -> None: + _advance_to_in_progress(client, token, sid) + client.post( + f"/api/simulations/{sid}/transition", + headers=_h(token), + json={"to": "review_required"}, + ) + + +def _advance_to_done(client: FlaskClient, redteam_token: str, soc_token: str, sid: int) -> None: + _advance_to_review_required(client, redteam_token, sid) + client.post( + f"/api/simulations/{sid}/transition", + headers=_h(soc_token), + json={"to": "done"}, + ) + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +class TestExecuteHappyPath: + def test_two_commands_create_two_tasks( + self, app: Flask, 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 = _execute(client, admin_token, sim["id"], ["whoami", "ipconfig"]) + assert resp.status_code == 200 + body = resp.get_json() + assert len(body["tasks"]) == 2 + assert body["tasks"][0]["command"] == "whoami" + assert body["tasks"][1]["command"] == "ipconfig" + + with app.app_context(): + rows = C2Task.query.filter_by(simulation_id=sim["id"]).all() + assert len(rows) == 2 + + def test_task_response_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"]) + + resp = _execute(client, admin_token, sim["id"], ["hostname"]) + task = resp.get_json()["tasks"][0] + assert "id" in task + assert "mythic_task_display_id" in task + assert "command" in task + assert "status" in task + assert "completed" in task + + def test_pending_sim_transitions_to_in_progress( + self, app: Flask, 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"]) + assert sim["status"] == "pending" + + _execute(client, admin_token, sim["id"], ["whoami"]) + + with app.app_context(): + updated = db.session.get(Simulation, sim["id"]) + assert updated is not None + assert updated.status == SimulationStatus.IN_PROGRESS + + def test_already_in_progress_stays_in_progress( + self, app: Flask, 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"]) + _advance_to_in_progress(client, admin_token, sim["id"]) + + resp = _execute(client, admin_token, sim["id"], ["whoami"]) + assert resp.status_code == 200 + + with app.app_context(): + updated = db.session.get(Simulation, sim["id"]) + assert updated is not None + assert updated.status == SimulationStatus.IN_PROGRESS + + def test_review_required_sim_still_allowed( + self, app: Flask, 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_review_required(client, admin_token, sim["id"]) + + resp = _execute(client, admin_token, sim["id"], ["net use"]) + assert resp.status_code == 200 + + # Status stays review_required — no regression to in_progress. + with app.app_context(): + updated = db.session.get(Simulation, sim["id"]) + assert updated is not None + assert updated.status == SimulationStatus.REVIEW_REQUIRED + + def test_redteam_can_execute( + 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"]) + + resp = _execute(client, redteam_token, sim["id"], ["whoami"]) + assert resp.status_code == 200 + + def test_mythic_task_display_id_stored( + self, app: Flask, 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"]) + + with app.app_context(): + task = C2Task.query.filter_by(simulation_id=sim["id"]).first() + assert task is not None + assert task.mythic_task_display_id == 1000 # FakeAdapter starts at 1000 + + +# --------------------------------------------------------------------------- +# Error cases +# --------------------------------------------------------------------------- + + +class TestExecuteValidation: + def test_400_empty_commands( + 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 = _execute(client, admin_token, sim["id"], []) + assert resp.status_code == 400 + + def test_400_non_string_command( + 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/execute", + headers=_h(admin_token), + json={"callback_display_id": 1, "commands": [42]}, + ) + 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/execute", + headers=_h(admin_token), + json={"commands": ["whoami"]}, + ) + assert resp.status_code == 400 + + def test_409_done_sim( + 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 = _execute(client, admin_token, sim["id"], ["whoami"]) + assert resp.status_code == 409 + assert "done" in resp.get_json().get("error", "").lower() + + def test_404_simulation_not_found( + self, client: FlaskClient, admin_token: str + ) -> None: + resp = _execute(client, admin_token, 9999, ["whoami"]) + assert resp.status_code == 404 + + 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 = _execute(client, admin_token, sim["id"], ["whoami"]) + assert resp.status_code == 404 + + 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 = _execute(client, soc_token, sim["id"], ["whoami"]) + 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 = _execute(client, admin_token, sim["id"], ["whoami"]) + assert resp.status_code == 503 + + def test_502_adapter_error( + self, monkeypatch, client: FlaskClient, admin_token: str + ) -> None: + from backend.app.services.c2 import fake as fake_mod + + def _boom(self, callback_display_id, command, params=None): + raise C2Error("task queue full") + + monkeypatch.setattr(fake_mod.FakeAdapter, "create_task", _boom) + + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + sim = _make_sim(client, admin_token, eng["id"]) + + resp = _execute(client, admin_token, sim["id"], ["whoami"]) + assert resp.status_code == 502 + assert "task queue full" in resp.get_json().get("error", "") diff --git a/backend/tests/test_c2_history.py b/backend/tests/test_c2_history.py new file mode 100644 index 0000000..196947c --- /dev/null +++ b/backend/tests/test_c2_history.py @@ -0,0 +1,215 @@ +"""Tests for GET /api/engagements//c2/callbacks//history.""" +from __future__ import annotations + +import pytest +from cryptography.fernet import Fernet +from flask.testing import FlaskClient + +from backend.app.services.c2.adapter import C2Error +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 _history(client: FlaskClient, token: str, eid: int, cid: int, **params): + return client.get( + f"/api/engagements/{eid}/c2/callbacks/{cid}/history", + headers=_h(token), + query_string=params, + ) + + +# --------------------------------------------------------------------------- +# Happy path +# --------------------------------------------------------------------------- + + +class TestHistoryHappyPath: + def test_returns_200(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.status_code == 200 + + def test_response_shape(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + body = resp.get_json() + assert "tasks" in body + assert "total" in body + assert "page" in body + assert "page_size" in body + + def test_task_shape(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + task = resp.get_json()["tasks"][0] + for field in ("display_id", "command", "params", "status", "completed", "timestamp"): + assert field in task, f"missing field: {field}" + + def test_default_page_is_1(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.get_json()["page"] == 1 + + def test_default_page_size_is_25(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.get_json()["page_size"] == 25 + + def test_callback_1_has_12_total(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.get_json()["total"] == 12 + + def test_callback_2_has_0_tasks(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 2) + body = resp.get_json() + assert body["total"] == 0 + assert body["tasks"] == [] + + def test_pagination_page_size_applied(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page=1, page_size=5) + body = resp.get_json() + assert len(body["tasks"]) == 5 + assert body["page_size"] == 5 + + def test_redteam_can_view_history( + self, client: FlaskClient, admin_token: str, redteam_token: str + ) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, redteam_token, eng["id"], 1) + assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# Validation errors +# --------------------------------------------------------------------------- + + +class TestHistoryValidation: + def test_400_page_size_too_large(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page_size=101) + assert resp.status_code == 400 + + def test_400_page_zero(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page=0) + assert resp.status_code == 400 + + def test_400_page_size_zero(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page_size=0) + assert resp.status_code == 400 + + def test_400_page_negative(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page=-1) + assert resp.status_code == 400 + + def test_400_page_size_100_is_ok(self, client: FlaskClient, admin_token: str) -> None: + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1, page_size=100) + assert resp.status_code == 200 + + +# --------------------------------------------------------------------------- +# Authorization / error cases +# --------------------------------------------------------------------------- + + +class TestHistoryErrors: + 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"]) + resp = _history(client, soc_token, eng["id"], 1) + 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) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.status_code == 503 + + def test_404_engagement_not_found( + self, client: FlaskClient, admin_token: str + ) -> None: + resp = _history(client, admin_token, 9999, 1) + assert resp.status_code == 404 + + def test_404_no_c2_config( + self, client: FlaskClient, admin_token: str + ) -> None: + eng = _make_engagement(client, admin_token) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.status_code == 404 + + def test_502_adapter_error( + self, monkeypatch, client: FlaskClient, admin_token: str + ) -> None: + from backend.app.services.c2 import fake as fake_mod + + def _boom(self, callback_display_id, page=1, page_size=25): + raise C2Error("upstream error") + + monkeypatch.setattr(fake_mod.FakeAdapter, "list_callback_tasks", _boom) + + eng = _make_engagement(client, admin_token) + _put_config(client, admin_token, eng["id"]) + resp = _history(client, admin_token, eng["id"], 1) + assert resp.status_code == 502 + assert "upstream error" in resp.get_json().get("error", "") diff --git a/backend/tests/test_c2_import.py b/backend/tests/test_c2_import.py new file mode 100644 index 0000000..e385530 --- /dev/null +++ b/backend/tests/test_c2_import.py @@ -0,0 +1,437 @@ +"""Tests for POST /api/simulations//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", "") diff --git a/backend/tests/test_c2_mapping.py b/backend/tests/test_c2_mapping.py new file mode 100644 index 0000000..5392c74 --- /dev/null +++ b/backend/tests/test_c2_mapping.py @@ -0,0 +1,208 @@ +"""Unit tests for apply_task_to_simulation() mapping helper — §0.11 contract.""" +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import MagicMock + +from backend.app.services.c2.mapping import apply_task_to_simulation + + +def _make_task( + command: str = "whoami", + output: str | None = "root", + mapping_applied: bool = False, + completed_at: datetime | None = None, +) -> MagicMock: + task = MagicMock() + task.command = command + task.output = output + task.mapping_applied = mapping_applied + task.completed_at = completed_at + return task + + +def _make_sim( + execution_result: str | None = None, + executed_at: datetime | None = None, + commands: str | None = None, +) -> MagicMock: + sim = MagicMock() + sim.execution_result = execution_result + sim.executed_at = executed_at + sim.commands = commands + sim.updated_at = None + return sim + + +class TestExecutionResult: + def test_first_task_produces_command_block(self): + task = _make_task(command="whoami", output="root") + sim = _make_sim() + + apply_task_to_simulation(task, sim) + + assert sim.execution_result == "$ whoami\nroot\n" + + def test_second_task_appended_with_block_separator(self): + """Two tasks → two '$ command\noutput\n' blocks separated by a single newline.""" + sim = _make_sim() + t1 = _make_task(command="whoami", output="root") + t2 = _make_task(command="hostname", output="lab-1") + + apply_task_to_simulation(t1, sim) + apply_task_to_simulation(t2, sim) + + assert sim.execution_result == "$ whoami\nroot\n$ hostname\nlab-1\n" + + def test_no_double_blank_line_when_existing_ends_with_newline(self): + """If existing result already ends with \n, no extra blank line is inserted.""" + sim = _make_sim(execution_result="$ id\nuid=0\n") + task = _make_task(command="hostname", output="lab-1") + + apply_task_to_simulation(task, sim) + + assert sim.execution_result == "$ id\nuid=0\n$ hostname\nlab-1\n" + + def test_empty_output_skips_block_but_marks_applied(self): + task = _make_task(output="") + sim = _make_sim(execution_result="$ id\nuid=0\n") + + apply_task_to_simulation(task, sim) + + assert sim.execution_result == "$ id\nuid=0\n" + assert task.mapping_applied is True + + def test_none_output_skips_block_but_marks_applied(self): + task = _make_task(output=None) + sim = _make_sim() + + apply_task_to_simulation(task, sim) + + assert sim.execution_result is None + assert task.mapping_applied is True + + def test_command_with_empty_string_produces_dollar_header(self): + """Empty command → block header is '$ \n\n' (consistent, not suppressed).""" + task = _make_task(command="", output="some output") + sim = _make_sim() + + apply_task_to_simulation(task, sim) + + assert sim.execution_result == "$ \nsome output\n" or sim.execution_result == "$ \nsome output\n" + + +class TestExecutedAt: + def test_sets_executed_at_from_task_when_null(self): + ts = datetime(2026, 6, 10, 12, 0, 0, tzinfo=UTC) + task = _make_task(completed_at=ts) + sim = _make_sim(executed_at=None) + + apply_task_to_simulation(task, sim) + + assert sim.executed_at == ts + + def test_does_not_overwrite_existing_executed_at(self): + original_ts = datetime(2026, 6, 1, 0, 0, 0, tzinfo=UTC) + later_ts = datetime(2026, 6, 10, 12, 0, 0, tzinfo=UTC) + task = _make_task(completed_at=later_ts) + sim = _make_sim(executed_at=original_ts) + + apply_task_to_simulation(task, sim) + + assert sim.executed_at == original_ts + + def test_executed_at_stays_null_when_task_completed_at_is_none(self): + task = _make_task(completed_at=None) + sim = _make_sim(executed_at=None) + + apply_task_to_simulation(task, sim) + + assert sim.executed_at is None + + def test_first_task_sets_executed_at_second_does_not_overwrite(self): + ts1 = datetime(2026, 6, 10, 10, 0, 0, tzinfo=UTC) + ts2 = datetime(2026, 6, 10, 11, 0, 0, tzinfo=UTC) + t1 = _make_task(command="whoami", output="root", completed_at=ts1) + t2 = _make_task(command="hostname", output="lab-1", completed_at=ts2) + sim = _make_sim(executed_at=None) + + apply_task_to_simulation(t1, sim) + apply_task_to_simulation(t2, sim) + + assert sim.executed_at == ts1 + + +class TestCommandsDedup: + def test_appends_command_to_empty_commands(self): + task = _make_task(command="whoami", output="root") + sim = _make_sim(commands=None) + + apply_task_to_simulation(task, sim) + + assert sim.commands == "whoami" + + def test_appends_second_distinct_command(self): + sim = _make_sim(commands=None) + t1 = _make_task(command="whoami", output="root") + t2 = _make_task(command="hostname", output="lab-1") + + apply_task_to_simulation(t1, sim) + apply_task_to_simulation(t2, sim) + + assert sim.commands == "whoami\nhostname" + + def test_deduplicates_repeated_command(self): + sim = _make_sim(commands=None) + t1 = _make_task(command="whoami", output="root") + t2 = _make_task(command="whoami", output="root2") + + apply_task_to_simulation(t1, sim) + apply_task_to_simulation(t2, sim) + + assert sim.commands == "whoami" + + def test_dedup_is_case_and_whitespace_stripped(self): + sim = _make_sim(commands="whoami") + task = _make_task(command=" whoami ", output="root") + + apply_task_to_simulation(task, sim) + + # " whoami ".strip() == "whoami" which is already present → no append. + assert sim.commands == "whoami" + + def test_empty_command_not_appended(self): + task = _make_task(command="", output="output") + sim = _make_sim(commands=None) + + apply_task_to_simulation(task, sim) + + # task.command is falsy → commands block skipped. + assert sim.commands is None + + +class TestIdempotency: + def test_no_op_when_mapping_already_applied(self): + task = _make_task(output="root", mapping_applied=True) + sim = _make_sim(execution_result="existing") + + apply_task_to_simulation(task, sim) + + assert sim.execution_result == "existing" + + def test_always_marks_mapping_applied(self): + task = _make_task(output="root") + sim = _make_sim() + + apply_task_to_simulation(task, sim) + + assert task.mapping_applied is True + + def test_updated_at_is_set(self): + task = _make_task(output="root") + sim = _make_sim() + before = datetime.now(UTC) + + apply_task_to_simulation(task, sim) + + assert sim.updated_at is not None + assert sim.updated_at >= before diff --git a/backend/tests/test_c2_tasks_list.py b/backend/tests/test_c2_tasks_list.py new file mode 100644 index 0000000..e231545 --- /dev/null +++ b/backend/tests/test_c2_tasks_list.py @@ -0,0 +1,375 @@ +"""Tests for GET /api/simulations//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", + "source", "mapping_applied", "created_at", "completed_at"): + assert field in task, f"missing field: {field}" + assert task["source"] == "mimic" + + 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" diff --git a/backend/tests/test_crypto.py b/backend/tests/test_crypto.py new file mode 100644 index 0000000..71da8d3 --- /dev/null +++ b/backend/tests/test_crypto.py @@ -0,0 +1,52 @@ +"""Tests for the Fernet crypto service.""" +from __future__ import annotations + +import pytest +from cryptography.fernet import Fernet + +from backend.app.services.crypto import C2Disabled, decrypt, encrypt + + +@pytest.fixture() +def fernet_key(monkeypatch) -> str: + key = Fernet.generate_key().decode() + monkeypatch.setenv("MIMIC_ENCRYPTION_KEY", key) + return key + + +@pytest.fixture() +def no_key(monkeypatch): + monkeypatch.delenv("MIMIC_ENCRYPTION_KEY", raising=False) + + +class TestEncryptDecrypt: + def test_round_trip(self, fernet_key): + plaintext = "s3cr3t-api-token" + ciphertext = encrypt(plaintext) + assert ciphertext != plaintext + assert decrypt(ciphertext) == plaintext + + def test_different_tokens_for_same_input(self, fernet_key): + # Fernet tokens are non-deterministic (random IV). + t1 = encrypt("same") + t2 = encrypt("same") + assert t1 != t2 + assert decrypt(t1) == decrypt(t2) == "same" + + def test_decrypt_invalid_ciphertext(self, fernet_key): + with pytest.raises(ValueError): + decrypt("not-valid-fernet-token") + + +class TestKeyAbsent: + def test_encrypt_raises_c2disabled(self, no_key): + with pytest.raises(C2Disabled): + encrypt("anything") + + def test_decrypt_raises_c2disabled(self, no_key): + with pytest.raises(C2Disabled): + decrypt("anything") + + def test_c2disabled_message(self, no_key): + with pytest.raises(C2Disabled, match="MIMIC_ENCRYPTION_KEY"): + encrypt("x") diff --git a/backend/tests/test_migration_0006_c2.py b/backend/tests/test_migration_0006_c2.py new file mode 100644 index 0000000..453d43e --- /dev/null +++ b/backend/tests/test_migration_0006_c2.py @@ -0,0 +1,199 @@ +"""Migration round-trip test for 0006_c2_layer. + +Verifies that upgrade() creates c2_config and c2_task with the expected schema, +and that downgrade() removes both tables cleanly. + +Uses the resolved-path pattern (derives path from __file__) to avoid the +hardcoded-path regression documented in lessons.md Sprint 4. +""" +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from alembic.operations import Operations +from alembic.runtime.migration import MigrationContext +from sqlalchemy import create_engine, inspect, text + + +def _load_migration(): + versions_dir = Path(__file__).resolve().parent.parent / "migrations" / "versions" + path = versions_dir / "0006_c2_layer.py" + spec = importlib.util.spec_from_file_location("migration_0006", path) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) # type: ignore[union-attr] + return mod + + +def _fresh_engine(): + """In-memory SQLite with the tables that 0006 depends on already present.""" + engine = create_engine("sqlite:///:memory:") + with engine.begin() as conn: + conn.execute( + text( + """ + CREATE TABLE users ( + id INTEGER PRIMARY KEY, + username TEXT NOT NULL UNIQUE, + password_hash TEXT NOT NULL, + role TEXT NOT NULL, + created_at DATETIME NOT NULL + ) + """ + ) + ) + conn.execute( + text( + """ + CREATE TABLE engagements ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + description TEXT, + start_date DATE NOT NULL, + end_date DATE, + status TEXT NOT NULL DEFAULT 'planned', + created_at DATETIME NOT NULL, + created_by_id INTEGER NOT NULL REFERENCES users(id) + ) + """ + ) + ) + conn.execute( + text( + """ + CREATE TABLE simulations ( + id INTEGER PRIMARY KEY, + engagement_id INTEGER NOT NULL REFERENCES engagements(id), + name TEXT NOT NULL, + techniques JSON NOT NULL DEFAULT '[]', + tactic_ids JSON NOT NULL DEFAULT '[]', + description TEXT, + commands TEXT, + prerequisites TEXT, + executed_at DATETIME, + execution_result TEXT, + log_source TEXT, + logs TEXT, + soc_comment TEXT, + incident_number TEXT, + status TEXT NOT NULL DEFAULT 'pending', + created_at DATETIME NOT NULL, + updated_at DATETIME, + created_by_id INTEGER NOT NULL REFERENCES users(id) + ) + """ + ) + ) + return engine + + +def _run_upgrade(engine, migration_mod): + with engine.begin() as conn: + ctx = MigrationContext.configure(conn) + ops = Operations(ctx) + ops._install_proxy() # type: ignore[attr-defined] + try: + migration_mod.upgrade() + finally: + ops._remove_proxy() # type: ignore[attr-defined] + + +def _run_downgrade(engine, migration_mod): + with engine.begin() as conn: + ctx = MigrationContext.configure(conn) + ops = Operations(ctx) + ops._install_proxy() # type: ignore[attr-defined] + try: + migration_mod.downgrade() + finally: + ops._remove_proxy() # type: ignore[attr-defined] + + +class TestMigration0006Upgrade: + def test_c2_config_table_created(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + + insp = inspect(engine) + assert "c2_config" in insp.get_table_names() + + def test_c2_task_table_created(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + + insp = inspect(engine) + assert "c2_task" in insp.get_table_names() + + def test_c2_config_columns(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + + insp = inspect(engine) + cols = {c["name"] for c in insp.get_columns("c2_config")} + assert {"id", "engagement_id", "url", "api_token_encrypted", + "verify_tls", "created_at", "updated_at"} <= cols + + def test_c2_config_unique_constraint_on_engagement_id(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + + # Insert a user and engagement first. + with engine.begin() as conn: + conn.execute(text( + "INSERT INTO users (id, username, password_hash, role, created_at) " + "VALUES (1, 'u', 'h', 'admin', '2026-01-01')" + )) + conn.execute(text( + "INSERT INTO engagements (id, name, start_date, status, created_at, created_by_id) " + "VALUES (1, 'Op', '2026-01-01', 'planned', '2026-01-01', 1)" + )) + conn.execute(text( + "INSERT INTO c2_config (engagement_id, url, api_token_encrypted, verify_tls, created_at) " + "VALUES (1, 'https://c2', 'tok', 1, '2026-01-01')" + )) + # Second insert on same engagement_id must fail. + try: + conn.execute(text( + "INSERT INTO c2_config (engagement_id, url, api_token_encrypted, verify_tls, created_at) " + "VALUES (1, 'https://c2b', 'tok2', 1, '2026-01-01')" + )) + raised = False + except Exception: + raised = True + assert raised, "UNIQUE constraint on c2_config.engagement_id must be enforced" + + def test_c2_task_columns(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + + insp = inspect(engine) + cols = {c["name"] for c in insp.get_columns("c2_task")} + assert {"id", "simulation_id", "mythic_task_display_id", "callback_display_id", + "command", "params", "status", "completed", "output", "source", + "created_at", "completed_at"} <= cols + + +class TestMigration0006Downgrade: + def test_downgrade_removes_c2_config(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + _run_downgrade(engine, mod) + + insp = inspect(engine) + assert "c2_config" not in insp.get_table_names() + + def test_downgrade_removes_c2_task(self): + engine = _fresh_engine() + mod = _load_migration() + _run_upgrade(engine, mod) + _run_downgrade(engine, mod) + + insp = inspect(engine) + assert "c2_task" not in insp.get_table_names() diff --git a/backend/tests/test_migration_0007_c2.py b/backend/tests/test_migration_0007_c2.py new file mode 100644 index 0000000..11b122d --- /dev/null +++ b/backend/tests/test_migration_0007_c2.py @@ -0,0 +1,124 @@ +"""Migration round-trip test for 0007_c2_task_mapping_applied. + +Verifies that upgrade() adds the mapping_applied column and downgrade() removes it. +Uses the resolved-path pattern per lessons.md Sprint 4. +""" +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from alembic.operations import Operations +from alembic.runtime.migration import MigrationContext +from sqlalchemy import create_engine, inspect, text + + +def _load_migration(name: str): + versions_dir = Path(__file__).resolve().parent.parent / "migrations" / "versions" + path = versions_dir / name + spec = importlib.util.spec_from_file_location(name.removesuffix(".py"), path) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) # type: ignore[union-attr] + return mod + + +def _fresh_engine_with_c2_task(): + """In-memory SQLite with c2_task already created (as left by 0006 upgrade).""" + engine = create_engine("sqlite:///:memory:") + with engine.begin() as conn: + conn.execute(text(""" + CREATE TABLE c2_task ( + id INTEGER PRIMARY KEY, + simulation_id INTEGER NOT NULL, + mythic_task_display_id INTEGER NOT NULL, + callback_display_id INTEGER NOT NULL, + command TEXT NOT NULL, + params TEXT, + status TEXT NOT NULL, + completed BOOLEAN NOT NULL DEFAULT 0, + output TEXT, + source TEXT NOT NULL, + created_at DATETIME NOT NULL, + completed_at DATETIME + ) + """)) + return engine + + +def _run_upgrade(engine, migration_mod): + with engine.begin() as conn: + ctx = MigrationContext.configure(conn) + ops = Operations(ctx) + ops._install_proxy() # type: ignore[attr-defined] + try: + migration_mod.upgrade() + finally: + ops._remove_proxy() # type: ignore[attr-defined] + + +def _run_downgrade(engine, migration_mod): + with engine.begin() as conn: + ctx = MigrationContext.configure(conn) + ops = Operations(ctx) + ops._install_proxy() # type: ignore[attr-defined] + try: + migration_mod.downgrade() + finally: + ops._remove_proxy() # type: ignore[attr-defined] + + +class TestMigration0007Upgrade: + def test_mapping_applied_column_added(self): + engine = _fresh_engine_with_c2_task() + mod = _load_migration("0007_c2_task_mapping_applied.py") + _run_upgrade(engine, mod) + + insp = inspect(engine) + cols = {c["name"] for c in insp.get_columns("c2_task")} + assert "mapping_applied" in cols + + def test_mapping_applied_defaults_to_false(self): + engine = _fresh_engine_with_c2_task() + mod = _load_migration("0007_c2_task_mapping_applied.py") + + # Insert a row before upgrading (no mapping_applied column yet). + with engine.begin() as conn: + conn.execute(text( + "INSERT INTO c2_task " + "(simulation_id, mythic_task_display_id, callback_display_id, " + "command, status, completed, source, created_at) " + "VALUES (1, 1000, 1, 'whoami', 'submitted', 0, 'mimic', '2026-01-01')" + )) + + _run_upgrade(engine, mod) + + with engine.begin() as conn: + row = conn.execute( + text("SELECT mapping_applied FROM c2_task WHERE id = 1") + ).fetchone() + assert row is not None + # SQLite stores booleans as 0/1. + assert row[0] == 0 or row[0] is False + + +class TestMigration0007Downgrade: + def test_downgrade_removes_mapping_applied(self): + engine = _fresh_engine_with_c2_task() + mod = _load_migration("0007_c2_task_mapping_applied.py") + _run_upgrade(engine, mod) + _run_downgrade(engine, mod) + + insp = inspect(engine) + cols = {c["name"] for c in insp.get_columns("c2_task")} + assert "mapping_applied" not in cols + + def test_downgrade_does_not_drop_other_columns(self): + engine = _fresh_engine_with_c2_task() + mod = _load_migration("0007_c2_task_mapping_applied.py") + _run_upgrade(engine, mod) + _run_downgrade(engine, mod) + + insp = inspect(engine) + cols = {c["name"] for c in insp.get_columns("c2_task")} + assert {"id", "simulation_id", "command", "status", "completed"} <= cols diff --git a/frontend/src/api/c2.ts b/frontend/src/api/c2.ts new file mode 100644 index 0000000..b0aa68d --- /dev/null +++ b/frontend/src/api/c2.ts @@ -0,0 +1,94 @@ +import { apiClient } from './client'; +import type { + C2CallbackHistoryResponse, + C2Config, + C2ConfigInput, + C2TestResult, + C2CallbacksResponse, + C2ExecuteInput, + C2ExecuteResponse, + C2ImportInput, + C2ImportResponse, + C2TasksResponse, +} from './types'; + +export async function getC2Config(engagementId: number): Promise { + try { + const { data } = await apiClient.get(`/engagements/${engagementId}/c2-config`); + return data; + } catch (err: unknown) { + const e = err as { response?: { status?: number } }; + if (e?.response?.status === 404) return null; + throw err; + } +} + +export async function putC2Config( + engagementId: number, + input: C2ConfigInput, +): Promise { + const { data } = await apiClient.put( + `/engagements/${engagementId}/c2-config`, + input, + ); + return data; +} + +export async function deleteC2Config(engagementId: number): Promise { + await apiClient.delete(`/engagements/${engagementId}/c2-config`); +} + +export async function testC2Config(engagementId: number): Promise { + const { data } = await apiClient.post( + `/engagements/${engagementId}/c2-config/test`, + ); + return data; +} + +export async function listCallbacks(engagementId: number): Promise { + const { data } = await apiClient.get( + `/engagements/${engagementId}/c2/callbacks`, + ); + return data; +} + +export async function executeC2( + simulationId: number, + input: C2ExecuteInput, +): Promise { + const { data } = await apiClient.post( + `/simulations/${simulationId}/c2/execute`, + input, + ); + return data; +} + +export async function getC2Tasks(simulationId: number): Promise { + const { data } = await apiClient.get( + `/simulations/${simulationId}/c2/tasks`, + ); + return data; +} + +export async function listCallbackHistory( + engagementId: number, + callbackDisplayId: number, + params: { page: number; pageSize: number }, +): Promise { + const { data } = await apiClient.get( + `/engagements/${engagementId}/c2/callbacks/${callbackDisplayId}/history`, + { params: { page: params.page, page_size: params.pageSize } }, + ); + return data; +} + +export async function importC2( + simulationId: number, + input: C2ImportInput, +): Promise { + const { data } = await apiClient.post( + `/simulations/${simulationId}/c2/import`, + input, + ); + return data; +} diff --git a/frontend/src/api/types.ts b/frontend/src/api/types.ts index 3c977eb..11ff621 100644 --- a/frontend/src/api/types.ts +++ b/frontend/src/api/types.ts @@ -154,3 +154,101 @@ export interface SimulationPatchInput { soc_comment?: string | null; incident_number?: string | null; } + +// C2 types + +export interface C2Config { + has_token: boolean; + url: string; + verify_tls: boolean; +} + +export interface C2ConfigInput { + url: string; + api_token?: string; + verify_tls: boolean; +} + +export interface C2TestResult { + ok: boolean; + error: string | null; +} + +export interface C2Callback { + display_id: number; + active: boolean; + host: string; + user: string; + domain: string; + last_checkin: string; +} + +export interface C2CallbacksResponse { + callbacks: C2Callback[]; +} + +// Thin shape returned by the execute endpoint +export interface C2ExecuteTask { + id: number; + mythic_task_display_id: number; + command: string; + status: string; + completed: boolean; +} + +export interface C2ExecuteInput { + callback_display_id: number; + commands: string[]; +} + +export interface C2ExecuteResponse { + tasks: C2ExecuteTask[]; +} + +// Full shape returned by the tasks list endpoint (M3) +export interface C2TaskListItem { + id: number; + mythic_task_display_id: number; + callback_display_id: number; + command: string; + params: string | null; + status: string; + completed: boolean; + output: string | null; + mapping_applied: boolean; + source: 'mimic' | 'import'; + created_at: string; + completed_at: string | null; +} + +export interface C2TasksResponse { + tasks: C2TaskListItem[]; +} + +// Callback history (M4) +export interface C2HistoryTask { + display_id: number; + command: string; + status: string; + completed: boolean; + completed_at: string | null; + created_at: string; +} + +export interface C2CallbackHistoryResponse { + tasks: C2HistoryTask[]; + total: number; + page: number; + page_size: number; +} + +// Import (M4) +export interface C2ImportInput { + callback_display_id: number; + task_display_ids: number[]; +} + +export interface C2ImportResponse { + imported: number; + skipped: number; +} diff --git a/frontend/src/components/C2CallbackPicker.tsx b/frontend/src/components/C2CallbackPicker.tsx new file mode 100644 index 0000000..7c760c0 --- /dev/null +++ b/frontend/src/components/C2CallbackPicker.tsx @@ -0,0 +1,95 @@ +import { extractApiError } from '@/api/client'; +import type { C2Callback } from '@/api/types'; + +interface C2CallbackPickerProps { + callbacks: C2Callback[]; + isLoading: boolean; + isError: boolean; + error: unknown; + selectedId: number | null; + onSelect: (id: number) => void; + rowTestId?: string; +} + +export function C2CallbackPicker({ + callbacks, + isLoading, + isError, + error, + selectedId, + onSelect, + rowTestId = 'c2-callback-row', +}: C2CallbackPickerProps): JSX.Element { + if (isLoading) { + return

Loading callbacks…

; + } + + if (isError) { + return ( +

+ Could not load callbacks: {extractApiError(error, 'Unknown error')} +

+ ); + } + + if (callbacks.length === 0) { + return

No callbacks available.

; + } + + return ( +
+ + + + + + + + + + + + + {callbacks.map((cb) => { + const isSelected = selectedId === cb.display_id; + return ( + onSelect(cb.display_id)} + onKeyDown={(e) => { + if (e.key === 'Enter' || e.key === ' ') { + e.preventDefault(); + onSelect(cb.display_id); + } + }} + tabIndex={0} + role="button" + className={`cursor-pointer border-b border-hairline focus:outline-none focus-visible:ring-2 focus-visible:ring-primary ${ + isSelected ? 'bg-primary-soft' : 'hover:bg-cloud' + }`} + > + + + + + + + + ); + })} + +
Display IDActiveHostUserDomainLast check-in
{cb.display_id} + + {cb.active ? 'Active' : 'Inactive'} + + {cb.host}{cb.user}{cb.domain}{cb.last_checkin}
+
+ ); +} diff --git a/frontend/src/components/C2ConfigCard.tsx b/frontend/src/components/C2ConfigCard.tsx new file mode 100644 index 0000000..1363310 --- /dev/null +++ b/frontend/src/components/C2ConfigCard.tsx @@ -0,0 +1,240 @@ +import { useEffect, useState, type FormEvent } from 'react'; +import { extractApiError } from '@/api/client'; +import { useC2Config, useDeleteC2Config, useTestC2Config, useUpdateC2Config } from '@/hooks/useC2'; +import { ConfirmDialog } from './ConfirmDialog'; +import { FormField, TextInput } from './FormField'; +import { useToast } from '@/hooks/useToast'; + +interface C2ConfigCardProps { + engagementId: number; +} + +export function C2ConfigCard({ engagementId }: C2ConfigCardProps): JSX.Element { + const { push } = useToast(); + + const configQuery = useC2Config(engagementId); + const updateMutation = useUpdateC2Config(engagementId); + const deleteMutation = useDeleteC2Config(engagementId); + const testMutation = useTestC2Config(engagementId); + + const config = configQuery.data; + const is503 = configQuery.error + ? (configQuery.error as { response?: { status?: number } })?.response?.status === 503 + : false; + + const [url, setUrl] = useState(''); + const [token, setToken] = useState(''); + const [verifyTls, setVerifyTls] = useState(true); + const [replaceToken, setReplaceToken] = useState(false); + const [showDeleteConfirm, setShowDeleteConfirm] = useState(false); + const [testResult, setTestResult] = useState<{ ok: boolean; message: string } | null>(null); + + // Sync URL and verifyTls from loaded config (but not token — write-only at API level) + useEffect(() => { + if (config) { + setUrl(config.url); + setVerifyTls(config.verify_tls); + } + }, [config]); + + const disabled = is503 || configQuery.isLoading; + + const onSave = async (e: FormEvent) => { + e.preventDefault(); + if (is503) return; + setTestResult(null); + + const input: { url: string; verify_tls: boolean; api_token?: string } = { + url: url.trim(), + verify_tls: verifyTls, + }; + // Send token only if: no existing config, OR user explicitly chose to replace + if (!config?.has_token || replaceToken) { + if (token.trim()) input.api_token = token.trim(); + } + + try { + await updateMutation.mutateAsync(input); + push('C2 configuration saved', 'success'); + setToken(''); + setReplaceToken(false); + } catch (err) { + push(extractApiError(err, 'Could not save C2 configuration'), 'error'); + } + }; + + const onDelete = async () => { + setShowDeleteConfirm(false); + setTestResult(null); + try { + await deleteMutation.mutateAsync(); + push('C2 configuration removed', 'success'); + setUrl(''); + setToken(''); + setVerifyTls(true); + setReplaceToken(false); + } catch (err) { + push(extractApiError(err, 'Could not remove C2 configuration'), 'error'); + } + }; + + const onTest = async () => { + setTestResult(null); + try { + const result = await testMutation.mutateAsync(); + setTestResult({ + ok: result.ok, + message: result.ok ? 'Connected' : (result.error ?? 'Connection failed'), + }); + } catch (err) { + setTestResult({ ok: false, message: extractApiError(err, 'Test failed') }); + } + }; + + const submitting = updateMutation.isPending || deleteMutation.isPending; + + return ( +
+

C2 configuration

+ + {is503 && ( +
+ C2 features are disabled (server has no encryption key configured). +
+ )} + + {configQuery.isLoading ? ( +

Loading…

+ ) : ( +
+ + setUrl(e.target.value)} + disabled={disabled} + /> + + + + {config?.has_token && !replaceToken ? ( +
+ + +
+ ) : ( + setToken(e.target.value)} + disabled={disabled} + /> + )} + + +
+ setVerifyTls(e.target.checked)} + disabled={disabled} + className="h-4 w-4 accent-primary" + /> + +
+ +
+ + + + + {testResult !== null && ( + + {testResult.message} + + )} + + {config?.has_token && ( + + )} +
+
+ )} + + {showDeleteConfirm && ( + setShowDeleteConfirm(false)} + /> + )} +
+ ); +} diff --git a/frontend/src/components/C2TaskStatusBadge.tsx b/frontend/src/components/C2TaskStatusBadge.tsx new file mode 100644 index 0000000..dcd04c1 --- /dev/null +++ b/frontend/src/components/C2TaskStatusBadge.tsx @@ -0,0 +1,27 @@ +// Dedicated badge for Mythic task statuses — separate from simulation status badges. +// submitted / processed → primary-soft (in-flight) +// completed → success-soft +// error* / fail* → warn-soft (task-level issue, not system error) +// anything else → cloud / graphite (unknown/neutral) + +interface C2TaskStatusBadgeProps { + status: string; +} + +function badgeClass(status: string): string { + const s = status.toLowerCase(); + if (s === 'completed') return 'bg-success-soft text-success'; + if (s.startsWith('error') || s.startsWith('fail')) return 'bg-warn-soft text-warn'; + if (s === 'submitted' || s === 'processed') return 'bg-primary-soft text-primary-deep'; + return 'bg-cloud text-graphite border border-hairline'; +} + +export function C2TaskStatusBadge({ status }: C2TaskStatusBadgeProps): JSX.Element { + return ( + + {status} + + ); +} diff --git a/frontend/src/components/C2TasksPanel.tsx b/frontend/src/components/C2TasksPanel.tsx new file mode 100644 index 0000000..138729a --- /dev/null +++ b/frontend/src/components/C2TasksPanel.tsx @@ -0,0 +1,135 @@ +import { Fragment, useState } from 'react'; +import { ChevronRight, ChevronDown } from 'lucide-react'; +import { useC2Tasks } from '@/hooks/useC2'; +import { C2TaskStatusBadge } from './C2TaskStatusBadge'; + +interface C2TasksPanelProps { + simulationId: number; +} + +export function C2TasksPanel({ simulationId }: C2TasksPanelProps): JSX.Element { + const query = useC2Tasks(simulationId, { enabled: true }); + const [expandedIds, setExpandedIds] = useState>(new Set()); + + const tasks = query.data?.tasks ?? []; + const isRefreshing = query.isFetching && !query.isLoading; + + function toggleExpand(id: number, completed: boolean) { + if (!completed) return; + setExpandedIds((prev) => { + const next = new Set(prev); + if (next.has(id)) { + next.delete(id); + } else { + next.add(id); + } + return next; + }); + } + + return ( +
+
+

C2 Tasks

+ {isRefreshing && ( + + Refreshing… + + )} +
+ + {tasks.length === 0 ? ( +
+

+ No C2 tasks yet. Use Execute via C2 to launch commands. +

+
+ ) : ( +
+ + + + + + + + + + + + {tasks.map((task) => { + const isExpanded = expandedIds.has(task.id); + const canExpand = task.completed && Boolean(task.output); + return ( + + toggleExpand(task.id, task.completed)} + onKeyDown={(e) => { + if (canExpand && (e.key === 'Enter' || e.key === ' ')) { + e.preventDefault(); + toggleExpand(task.id, task.completed); + } + }} + tabIndex={canExpand ? 0 : undefined} + role={canExpand ? 'button' : undefined} + aria-expanded={canExpand ? isExpanded : undefined} + className={`border-b border-hairline ${canExpand ? 'cursor-pointer hover:bg-cloud focus:outline-none focus-visible:ring-2 focus-visible:ring-primary' : ''}`} + > + + + + + + + + {isExpanded && task.output && ( + + + + )} + + ); + })} + +
+ TaskCommandSourceStatusCompleted at
+ {canExpand ? ( + isExpanded ? ( + + ) : ( + + ) + ) : null} + #{task.mythic_task_display_id} + {task.command} + + + {task.source === 'mimic' ? 'MIMIC' : 'IMPORT'} + + + + + {task.completed_at ?? '—'} +
+ 
+                            {task.output}
+
+
+
+ )} +
+ ); +} diff --git a/frontend/src/components/ExecuteViaC2Modal.tsx b/frontend/src/components/ExecuteViaC2Modal.tsx new file mode 100644 index 0000000..5eee92b --- /dev/null +++ b/frontend/src/components/ExecuteViaC2Modal.tsx @@ -0,0 +1,129 @@ +import { useState } from 'react'; +import { extractApiError } from '@/api/client'; +import { useC2Callbacks, useExecuteC2 } from '@/hooks/useC2'; +import { C2CallbackPicker } from './C2CallbackPicker'; +import { useToast } from '@/hooks/useToast'; + +interface ExecuteViaC2ModalProps { + simulationId: number; + engagementId: number; + initialCommands: string; + onClose: () => void; +} + +export function ExecuteViaC2Modal({ + simulationId, + engagementId, + initialCommands, + onClose, +}: ExecuteViaC2ModalProps): JSX.Element { + const { push } = useToast(); + + const callbacksQuery = useC2Callbacks(engagementId, { enabled: true }); + const executeMutation = useExecuteC2(simulationId, engagementId); + + const [selectedId, setSelectedId] = useState(null); + const [commands, setCommands] = useState(initialCommands); + const [submitError, setSubmitError] = useState(null); + + const callbacks = callbacksQuery.data?.callbacks ?? []; + + const commandLines = commands + .split('\n') + .map((l) => l.trim()) + .filter(Boolean); + + const canLaunch = selectedId !== null && commandLines.length > 0; + + const onLaunch = async () => { + if (!canLaunch) return; + setSubmitError(null); + try { + const result = await executeMutation.mutateAsync({ + callback_display_id: selectedId, + commands: commandLines, + }); + push(`${result.tasks.length} task(s) submitted`, 'success'); + onClose(); + } catch (err) { + setSubmitError(extractApiError(err, 'Could not execute via C2')); + } + }; + + return ( +
+