mimic/backend/app/models/c2_task.py

"""C2Task model — link between a Mimic simulation and a Mythic task."""
from __future__ import annotations

import enum
from datetime import UTC, datetime

from backend.app.extensions import db


class C2TaskSource(str, enum.Enum):
    MIMIC = "mimic"
    IMPORT = "import"


class C2Task(db.Model):  # type: ignore[name-defined]
    __tablename__ = "c2_task"

    id = db.Column(db.Integer, primary_key=True)
    simulation_id = db.Column(
        db.Integer,
        db.ForeignKey("simulations.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    mythic_task_display_id = db.Column(db.Integer, nullable=False)
    callback_display_id = db.Column(db.Integer, nullable=False)
    command = db.Column(db.Text, nullable=False)
    params = db.Column(db.Text, nullable=True)
    status = db.Column(db.Text, nullable=False)
    completed = db.Column(db.Boolean, nullable=False, default=False)
    output = db.Column(db.Text, nullable=True)
    source = db.Column(
        db.Enum(C2TaskSource, name="c2task_source"),
        nullable=False,
    )
    created_at = db.Column(
        db.DateTime, nullable=False, default=lambda: datetime.now(UTC)
    )
    completed_at = db.Column(db.DateTime, nullable=True)
    mapping_applied = db.Column(db.Boolean, nullable=False, default=False)

    simulation = db.relationship(
        "Simulation",
        backref=db.backref("c2_tasks", cascade="all, delete-orphan", lazy="dynamic"),
    )

    def __repr__(self) -> str:
        return f"<C2Task simulation_id={self.simulation_id} mythic_id={self.mythic_task_display_id}>"