fix(backend): complete c2 task→simulation mapping per spec + sanitize adapter errors (sprint 8 code-review)

mapping.py — full §0.11 contract:
1. execution_result: append '$ <command>\n<output>\n' block (previously
   wrote raw output without command header, making multi-task blobs
   unreadable in exports)
2. executed_at: set from task.completed_at when currently null (was
   completely missing — simulation.executed_at stayed null forever)
3. commands: append task.command deduplicated line-by-line (was
   completely missing — simulation.commands stayed empty)

mythic.py — sanitize transport errors:
Replace 'raise C2Error(str(exc))' (which leaks the Mythic URL via
requests exception repr) with 'raise C2Error(f"C2 transport error:
{type(exc).__name__}")'. Original exc stays chained for backend logs.

api/c2.py — remove redundant 'task.mapping_applied = True' in import
endpoint (apply_task_to_simulation() already sets it).

test_c2_mapping.py — full rewrite: 19 tests covering command blocks,
executed_at set/preserve, commands dedup, idempotency.

test_c2_adapter_mythic.py — add URL-leak sanitization assertion.

468 passed; ruff + mypy clean.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Knacky
2026-06-10 20:28:49 +02:00
parent 7d3d39639e
commit 38e282a126
5 changed files with 205 additions and 66 deletions

View File

@@ -500,7 +500,6 @@ def import_tasks(sid: int):
db.session.add(task)
db.session.flush()
apply_task_to_simulation(task, sim)
task.mapping_applied = True
else:
db.session.add(task)

View File

@@ -1,8 +1,9 @@
"""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).
apply_task_to_simulation() implements the full §0.11 contract:
1. execution_result — append "$ <command>\n<output>\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.
"""
@@ -15,24 +16,38 @@ 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).
"""Apply completed task data to simulation fields per §0.11.
No-op if task.mapping_applied is already True or task.output is empty.
Marks task.mapping_applied = True on completion.
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()
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
# 1) execution_result — "$ <command>\n<output>\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

View File

@@ -158,7 +158,7 @@ class MythicAdapter(C2Adapter):
try:
data = self._post({"query": _CALLBACKS_QUERY})
except requests.RequestException as exc:
raise C2Error(str(exc)) from exc
raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc
callbacks_raw = data.get("data", {}).get("callback", [])
return [
@@ -190,7 +190,7 @@ class MythicAdapter(C2Adapter):
},
})
except requests.RequestException as exc:
raise C2Error(str(exc)) from 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")
@@ -206,7 +206,7 @@ class MythicAdapter(C2Adapter):
"variables": {"display_id": task_display_id},
})
except requests.RequestException as exc:
raise C2Error(str(exc)) from exc
raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc
rows = data.get("data", {}).get("task", [])
if not rows:
@@ -238,7 +238,7 @@ class MythicAdapter(C2Adapter):
"variables": {"display_id": task_display_id},
})
except requests.RequestException as exc:
raise C2Error(str(exc)) from exc
raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc
rows = data.get("data", {}).get("response", [])
return "".join(
@@ -269,7 +269,7 @@ class MythicAdapter(C2Adapter):
"variables": {"callback_display_id": callback_display_id},
})
except requests.RequestException as exc:
raise C2Error(str(exc)) from exc
raise C2Error(f"C2 transport error: {type(exc).__name__}") from exc
rows = data.get("data", {}).get("task", [])
total: int = (