38e282a126
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>
54 lines
2.0 KiB
Python
54 lines
2.0 KiB
Python
"""C2 task → Simulation output mapping.
|
|
|
|
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.
|
|
"""
|
|
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 — "$ <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
|