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

"""Abstract C2 adapter interface and shared dataclasses."""
from __future__ import annotations

import base64
import binascii
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime


class C2Error(Exception):
    """Raised by adapters when the C2 returns an application-level error."""


@dataclass
class C2Health:
    ok: bool
    error: str | None = None


@dataclass
class C2Callback:
    display_id: int
    active: bool
    host: str
    user: str
    domain: str
    last_checkin: str  # ISO-8601 string


@dataclass
class C2TaskStatus:
    display_id: int
    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[C2HistoricalTask]
    total: int
    page: int
    page_size: int


def decode_response_text(raw: str) -> str:
    """Decode a base64-encoded Mythic response_text field.

    On binascii.Error (binary payload) returns "<binary> " + hex string
    so execution_result never silently corrupts.
    """
    try:
        return base64.b64decode(raw).decode("utf-8")
    except binascii.Error:
        return "<binary> " + raw.encode().hex()
    except UnicodeDecodeError:
        raw_bytes = base64.b64decode(raw)
        return "<binary> " + raw_bytes.hex()


class C2Adapter(ABC):
    """Thin interface over a C2 backend (Mythic or custom)."""

    @abstractmethod
    def test_connection(self) -> C2Health:
        """Verify that the C2 is reachable and the token is valid."""
        ...

    @abstractmethod
    def list_callbacks(self) -> list[C2Callback]:
        """Return active callbacks visible to this API token."""
        ...

    @abstractmethod
    def create_task(
        self,
        callback_display_id: int,
        command: str,
        params: str | None = None,
    ) -> int:
        """Issue a task and return its Mythic display_id."""
        ...

    @abstractmethod
    def get_task(self, task_display_id: int) -> C2TaskStatus:
        """Return current status of a task."""
        ...

    @abstractmethod
    def get_task_output(self, task_display_id: int) -> str:
        """Return decoded, concatenated output for a completed task."""
        ...

    @abstractmethod
    def list_callback_tasks(
        self,
        callback_display_id: int,
        page: int = 1,
        page_size: int = 25,
    ) -> C2TaskPage:
        """Return a paginated history of tasks for a callback."""
        ...