mimic/backend/app/models/simulation.py

"""Simulation model."""
from __future__ import annotations

import enum
from datetime import UTC, datetime

from backend.app.extensions import db


class SimulationStatus(str, enum.Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    REVIEW_REQUIRED = "review_required"
    DONE = "done"


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

    id = db.Column(db.Integer, primary_key=True)
    engagement_id = db.Column(
        db.Integer,
        db.ForeignKey("engagements.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    name = db.Column(db.String(255), nullable=False)
    techniques = db.Column(db.JSON, nullable=False, default=list)
    tactic_ids = db.Column(db.JSON, nullable=False, default=list)
    description = db.Column(db.Text, nullable=True)
    commands = db.Column(db.Text, nullable=True)
    prerequisites = db.Column(db.Text, nullable=True)
    executed_at = db.Column(db.DateTime, nullable=True)
    execution_result = db.Column(db.Text, nullable=True)
    log_source = db.Column(db.Text, nullable=True)
    logs = db.Column(db.Text, nullable=True)
    soc_comment = db.Column(db.Text, nullable=True)
    incident_number = db.Column(db.String(128), nullable=True)
    status = db.Column(
        db.Enum(SimulationStatus, name="simulation_status"),
        nullable=False,
        default=SimulationStatus.PENDING,
    )
    created_at = db.Column(
        db.DateTime, nullable=False, default=lambda: datetime.now(UTC)
    )
    updated_at = db.Column(db.DateTime, nullable=True)
    created_by_id = db.Column(
        db.Integer,
        db.ForeignKey("users.id", ondelete="RESTRICT"),
        nullable=False,
        index=True,
    )

    engagement = db.relationship(
        "Engagement",
        backref=db.backref("simulations", cascade="all, delete-orphan", lazy="dynamic"),
    )
    created_by = db.relationship("User", backref="simulations", lazy="joined")

    def __repr__(self) -> str:
        return f"<Simulation {self.id} {self.name!r}>"