feat(backend): c2 callback history + task import (sprint 8 M4)
Command source decision: extended C2TaskStatus with command: str | None (default None). Added command_name to _GET_TASK_QUERY so get_task() returns command in a single round-trip — no separate history fetch needed on import. 4-line change, zero cascading test impact. adapter.py: - C2TaskStatus: add command: str | None = None field - C2HistoricalTask: new dataclass (display_id, command, params, status, completed, timestamp) for history rows - C2TaskPage.items: typed as list[C2HistoricalTask] (was list[dict]) mythic.py: - _GET_TASK_QUERY: add command_name field - _LIST_CALLBACK_TASKS_QUERY: new query (order_by id desc, limit/offset) - _COUNT_CALLBACK_TASKS_QUERY: new aggregate query for total - get_task(): surfaces command_name as status.command - list_callback_tasks(): two _post() calls (tasks + count), allow_redirects=False fake.py: - _FAKE_HISTORY: frozen deterministic history (cb1=12, cb2=0, cb3=5 tasks) - list_callback_tasks(): serves from _FAKE_HISTORY, pagination applied - get_task(): returns command from _tasks dict api/c2.py: - GET /api/engagements/<eid>/c2/callbacks/<cid>/history: page+page_size defaults 1/25, cap 100, reject <1, 502 on adapter error - POST /api/simulations/<sid>/c2/import: idempotent per (sim,mythic_id) pair, source=import, completed tasks get output+mapping_applied, incomplete tasks stored for poll-on-read pickup, auto-transition pending→in_progress 60 new tests (456 total); pytest/ruff/mypy all green Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -367,3 +367,151 @@ def list_simulation_tasks(sid: int):
|
||||
for t in tasks
|
||||
]
|
||||
}), 200
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# M4 — callback history + task import
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@c2_bp.get("/<int:eid>/c2/callbacks/<int:cid>/history")
|
||||
@role_required("admin", "redteam")
|
||||
def list_callback_history(eid: int, cid: int):
|
||||
guard = _crypto_guard()
|
||||
if guard is not None:
|
||||
return guard
|
||||
|
||||
engagement = db.session.get(Engagement, eid)
|
||||
if engagement is None:
|
||||
return jsonify({"error": "Engagement not found"}), 404
|
||||
|
||||
# Validate pagination params.
|
||||
try:
|
||||
page = int(request.args.get("page", 1))
|
||||
page_size = int(request.args.get("page_size", 25))
|
||||
except (ValueError, TypeError):
|
||||
return jsonify({"error": "page and page_size must be integers"}), 400
|
||||
|
||||
if page < 1 or page_size < 1:
|
||||
return jsonify({"error": "page and page_size must be >= 1"}), 400
|
||||
if page_size > 100:
|
||||
return jsonify({"error": "page_size must be <= 100"}), 400
|
||||
|
||||
adapter, err = _load_adapter_for_engagement(engagement)
|
||||
if err is not None:
|
||||
return err
|
||||
|
||||
try:
|
||||
page_result = adapter.list_callback_tasks(
|
||||
callback_display_id=cid,
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
)
|
||||
except C2Error as exc:
|
||||
return jsonify({"error": str(exc)}), 502
|
||||
|
||||
return jsonify({
|
||||
"tasks": [
|
||||
{
|
||||
"display_id": t.display_id,
|
||||
"command": t.command,
|
||||
"params": t.params,
|
||||
"status": t.status,
|
||||
"completed": t.completed,
|
||||
"timestamp": t.timestamp,
|
||||
}
|
||||
for t in page_result.items
|
||||
],
|
||||
"total": page_result.total,
|
||||
"page": page_result.page,
|
||||
"page_size": page_result.page_size,
|
||||
}), 200
|
||||
|
||||
|
||||
@sims_c2_bp.post("/<int:sid>/c2/import")
|
||||
@role_required("admin", "redteam")
|
||||
def import_tasks(sid: int):
|
||||
guard = _crypto_guard()
|
||||
if guard is not None:
|
||||
return guard
|
||||
|
||||
sim = db.session.get(Simulation, sid)
|
||||
if sim is None:
|
||||
return jsonify({"error": "Simulation not found"}), 404
|
||||
|
||||
if sim.status == SimulationStatus.DONE:
|
||||
return jsonify({"error": "simulation is done — reopen first"}), 409
|
||||
|
||||
data = request.get_json(silent=True) or {}
|
||||
callback_display_id = data.get("callback_display_id")
|
||||
task_display_ids = data.get("task_display_ids")
|
||||
|
||||
if not isinstance(callback_display_id, int):
|
||||
return jsonify({"error": "callback_display_id must be an integer"}), 400
|
||||
if not isinstance(task_display_ids, list) or len(task_display_ids) == 0:
|
||||
return jsonify({"error": "task_display_ids must be a non-empty list"}), 400
|
||||
for tid in task_display_ids:
|
||||
if not isinstance(tid, int):
|
||||
return jsonify({"error": "each task_display_id must be an integer"}), 400
|
||||
|
||||
engagement = db.session.get(Engagement, sim.engagement_id)
|
||||
if engagement is None:
|
||||
return jsonify({"error": "Engagement not found"}), 404
|
||||
|
||||
adapter, err = _load_adapter_for_engagement(engagement)
|
||||
if err is not None:
|
||||
return err
|
||||
|
||||
imported_count = 0
|
||||
skipped_count = 0
|
||||
|
||||
try:
|
||||
for task_display_id in task_display_ids:
|
||||
# Idempotency: skip if already imported for this simulation.
|
||||
existing = C2Task.query.filter_by(
|
||||
simulation_id=sid,
|
||||
mythic_task_display_id=task_display_id,
|
||||
).first()
|
||||
if existing is not None:
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
status = adapter.get_task(task_display_id)
|
||||
task = C2Task(
|
||||
simulation_id=sid,
|
||||
mythic_task_display_id=task_display_id,
|
||||
callback_display_id=callback_display_id,
|
||||
command=status.command or "",
|
||||
params=None,
|
||||
status=status.status,
|
||||
completed=status.completed,
|
||||
source=C2TaskSource.IMPORT,
|
||||
created_at=datetime.now(UTC),
|
||||
mapping_applied=False,
|
||||
)
|
||||
|
||||
if status.completed:
|
||||
task.completed_at = status.completed_at or datetime.now(UTC)
|
||||
try:
|
||||
task.output = adapter.get_task_output(task_display_id)
|
||||
except C2Error:
|
||||
task.output = ""
|
||||
db.session.add(task)
|
||||
db.session.flush()
|
||||
apply_task_to_simulation(task, sim)
|
||||
task.mapping_applied = True
|
||||
else:
|
||||
db.session.add(task)
|
||||
|
||||
imported_count += 1
|
||||
|
||||
except C2Error as exc:
|
||||
db.session.rollback()
|
||||
return jsonify({"error": str(exc)}), 502
|
||||
|
||||
# Auto-transition pending → in_progress when at least one task was imported.
|
||||
if imported_count > 0:
|
||||
promote_to_in_progress(sim)
|
||||
|
||||
db.session.commit()
|
||||
return jsonify({"imported": imported_count, "skipped": skipped_count}), 200
|
||||
|
||||
@@ -34,11 +34,25 @@ class C2TaskStatus:
|
||||
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[dict] # raw task dicts from Mythic
|
||||
items: list[C2HistoricalTask]
|
||||
total: int
|
||||
page: int
|
||||
page_size: int
|
||||
|
||||
@@ -10,10 +10,45 @@ from backend.app.services.c2.adapter import (
|
||||
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(
|
||||
@@ -109,6 +144,7 @@ class FakeAdapter(C2Adapter):
|
||||
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:
|
||||
@@ -130,14 +166,11 @@ class FakeAdapter(C2Adapter):
|
||||
page: int = 1,
|
||||
page_size: int = 25,
|
||||
) -> C2TaskPage:
|
||||
items = [
|
||||
t for t in self._tasks.values()
|
||||
if t["callback_display_id"] == callback_display_id
|
||||
]
|
||||
all_items = _FAKE_HISTORY.get(callback_display_id, [])
|
||||
start = (page - 1) * page_size
|
||||
return C2TaskPage(
|
||||
items=items[start : start + page_size],
|
||||
total=len(items),
|
||||
items=all_items[start : start + page_size],
|
||||
total=len(all_items),
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
@@ -21,6 +21,7 @@ from backend.app.services.c2.adapter import (
|
||||
C2Callback,
|
||||
C2Error,
|
||||
C2Health,
|
||||
C2HistoricalTask,
|
||||
C2TaskPage,
|
||||
C2TaskStatus,
|
||||
decode_response_text,
|
||||
@@ -61,6 +62,7 @@ _GET_TASK_QUERY = """
|
||||
query GetTask($display_id: Int!) {
|
||||
task(where: {display_id: {_eq: $display_id}}) {
|
||||
display_id
|
||||
command_name
|
||||
status
|
||||
completed
|
||||
timestamp
|
||||
@@ -68,6 +70,34 @@ query GetTask($display_id: Int!) {
|
||||
}
|
||||
"""
|
||||
|
||||
_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(
|
||||
@@ -197,6 +227,7 @@ class MythicAdapter(C2Adapter):
|
||||
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:
|
||||
@@ -222,4 +253,41 @@ class MythicAdapter(C2Adapter):
|
||||
page: int = 1,
|
||||
page_size: int = 25,
|
||||
) -> C2TaskPage:
|
||||
raise NotImplementedError("M4")
|
||||
"""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(str(exc)) 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)
|
||||
|
||||
Reference in New Issue
Block a user