mimic/backend/app/services/c2/fake.py

"""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,
        )