feat(backend): c2 poll-on-read + output mapping (sprint 8 M3)
- adapter.py: add completed_at field to C2TaskStatus dataclass - mythic.py: implement get_task() (GraphQL task query) and get_task_output() (response query + decode_response_text concat) - fake.py: deterministic state progression via per-instance call counter; get_task_output raises C2Error until completed - mapping.py: apply_task_to_simulation() idempotent output mapper (mapping_applied anchor prevents double-writes) - migration 0007: add mapping_applied BOOLEAN NOT NULL DEFAULT false to c2_task - c2_task model: mapping_applied column added - api/c2.py: GET /api/simulations/<id>/c2/tasks poll-on-read endpoint; refreshes incomplete tasks from C2, fetches output on completion, applies mapping, skips re-polling for completed tasks; best-effort (C2Error on individual task skipped, returns 200 with stale status) - 51 new tests (396 total); pytest/ruff/mypy all green Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -4,7 +4,8 @@ from __future__ import annotations
|
||||
import base64
|
||||
import binascii
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class C2Error(Exception):
|
||||
@@ -32,6 +33,7 @@ class C2TaskStatus:
|
||||
display_id: int
|
||||
status: str
|
||||
completed: bool
|
||||
completed_at: datetime | None = field(default=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -8,6 +8,7 @@ from __future__ import annotations
|
||||
from backend.app.services.c2.adapter import (
|
||||
C2Adapter,
|
||||
C2Callback,
|
||||
C2Error,
|
||||
C2Health,
|
||||
C2TaskPage,
|
||||
C2TaskStatus,
|
||||
@@ -46,11 +47,17 @@ 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)
|
||||
@@ -78,20 +85,44 @@ class FakeAdapter(C2Adapter):
|
||||
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 task is None:
|
||||
return C2TaskStatus(display_id=task_display_id, status="unknown", completed=False)
|
||||
|
||||
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=task["status"],
|
||||
completed=task["completed"],
|
||||
status=status,
|
||||
completed=completed,
|
||||
)
|
||||
|
||||
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)
|
||||
if task is None:
|
||||
return ""
|
||||
return task.get("output") or ""
|
||||
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,
|
||||
|
||||
38
backend/app/services/c2/mapping.py
Normal file
38
backend/app/services/c2/mapping.py
Normal file
@@ -0,0 +1,38 @@
|
||||
"""C2 task → Simulation output mapping.
|
||||
|
||||
apply_task_to_simulation() writes task output into the simulation's
|
||||
execution_result field and marks the task as mapping_applied=True so that
|
||||
the operation is idempotent (safe to call multiple times for the same task).
|
||||
|
||||
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:
|
||||
"""Write task output into simulation.execution_result (append, newline-separated).
|
||||
|
||||
No-op if task.mapping_applied is already True or task.output is empty.
|
||||
Marks task.mapping_applied = True on completion.
|
||||
"""
|
||||
if task.mapping_applied:
|
||||
return
|
||||
|
||||
output = (task.output or "").strip()
|
||||
if not output:
|
||||
task.mapping_applied = True
|
||||
return
|
||||
|
||||
existing = (simulation.execution_result or "").rstrip("\n")
|
||||
if existing:
|
||||
simulation.execution_result = existing + "\n" + output
|
||||
else:
|
||||
simulation.execution_result = output
|
||||
|
||||
simulation.updated_at = datetime.now(UTC)
|
||||
task.mapping_applied = True
|
||||
@@ -12,6 +12,8 @@ M4: list_callback_tasks()
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import requests
|
||||
|
||||
from backend.app.services.c2.adapter import (
|
||||
@@ -21,6 +23,7 @@ from backend.app.services.c2.adapter import (
|
||||
C2Health,
|
||||
C2TaskPage,
|
||||
C2TaskStatus,
|
||||
decode_response_text,
|
||||
)
|
||||
|
||||
_HEALTH_QUERY = "{ __typename }"
|
||||
@@ -54,6 +57,28 @@ mutation CreateTask($callback_id: Int!, $command: String!, $params: String!) {
|
||||
}
|
||||
"""
|
||||
|
||||
_GET_TASK_QUERY = """
|
||||
query GetTask($display_id: Int!) {
|
||||
task(where: {display_id: {_eq: $display_id}}) {
|
||||
display_id
|
||||
status
|
||||
completed
|
||||
timestamp
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
_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."""
|
||||
@@ -144,10 +169,52 @@ class MythicAdapter(C2Adapter):
|
||||
return int(task_data["display_id"])
|
||||
|
||||
def get_task(self, task_display_id: int) -> C2TaskStatus:
|
||||
raise NotImplementedError("M3")
|
||||
"""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(str(exc)) 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,
|
||||
)
|
||||
|
||||
def get_task_output(self, task_display_id: int) -> str:
|
||||
raise NotImplementedError("M3")
|
||||
"""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(str(exc)) 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,
|
||||
|
||||
Reference in New Issue
Block a user