Metamorph/backend/app/models/template.py

"""Reusable templates: test_templates and scenario_templates.

A `mission_scenarios` row is a snapshot copy of a `scenario_templates` row at
mission-creation time. Templates can therefore be edited freely without
disturbing already-running missions.
"""

from __future__ import annotations

import uuid
from typing import Any

from sqlalchemy import (
    ARRAY,
    CheckConstraint,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
    Uuid,
)
from sqlalchemy.orm import Mapped, mapped_column, relationship

from app.db.base import Base
from app.db.mixins import SoftDeleteMixin, TimestampMixin, UuidPkMixin
from app.db.types import MITRE_KINDS, OPSEC_LEVELS


class TestTemplate(Base, UuidPkMixin, TimestampMixin, SoftDeleteMixin):
    __tablename__ = "test_templates"

    name: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str | None] = mapped_column(Text, nullable=True)
    objective: Mapped[str | None] = mapped_column(Text, nullable=True)
    procedure_md: Mapped[str | None] = mapped_column(Text, nullable=True)
    prerequisites_md: Mapped[str | None] = mapped_column(Text, nullable=True)
    expected_result_red_md: Mapped[str | None] = mapped_column(Text, nullable=True)
    expected_detection_blue_md: Mapped[str | None] = mapped_column(Text, nullable=True)
    opsec_level: Mapped[str] = mapped_column(String(8), default="medium", nullable=False)
    tags: Mapped[list[str]] = mapped_column(
        ARRAY(String(64)), nullable=False, server_default="{}"
    )
    expected_iocs: Mapped[list[str]] = mapped_column(
        ARRAY(String(255)), nullable=False, server_default="{}"
    )

    mitre_tags: Mapped[list["TestTemplateMitreTag"]] = relationship(
        back_populates="test_template",
        cascade="all, delete-orphan",
        lazy="selectin",
    )

    __table_args__ = (
        CheckConstraint(
            f"opsec_level IN ({', '.join(repr(v) for v in OPSEC_LEVELS)})",
            name="opsec_level_valid",
        ),
        Index("ix_test_templates_active", "deleted_at", postgresql_where="deleted_at IS NULL"),
        Index("ix_test_templates_name", "name"),
    )


class TestTemplateMitreTag(Base):
    """Polymorphic MITRE tag on a test template.

    Exactly one of `tactic_id`, `technique_id`, `subtechnique_id` is set —
    enforced by the CHECK constraint. This keeps FK integrity per MITRE level
    while letting a single conceptual table answer "what's tagged on this test".
    """

    __tablename__ = "test_template_mitre_tags"

    id: Mapped[uuid.UUID] = mapped_column(
        Uuid(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False
    )
    test_template_id: Mapped[uuid.UUID] = mapped_column(
        Uuid(as_uuid=True),
        ForeignKey("test_templates.id", ondelete="CASCADE"),
        nullable=False,
    )
    mitre_kind: Mapped[str] = mapped_column(String(16), nullable=False)
    tactic_id: Mapped[uuid.UUID | None] = mapped_column(
        Uuid(as_uuid=True), ForeignKey("mitre_tactics.id", ondelete="CASCADE"), nullable=True
    )
    technique_id: Mapped[uuid.UUID | None] = mapped_column(
        Uuid(as_uuid=True), ForeignKey("mitre_techniques.id", ondelete="CASCADE"), nullable=True
    )
    subtechnique_id: Mapped[uuid.UUID | None] = mapped_column(
        Uuid(as_uuid=True),
        ForeignKey("mitre_subtechniques.id", ondelete="CASCADE"),
        nullable=True,
    )

    test_template: Mapped[TestTemplate] = relationship(back_populates="mitre_tags")

    __table_args__: Any = (
        CheckConstraint(
            f"mitre_kind IN ({', '.join(repr(v) for v in MITRE_KINDS)})",
            name="mitre_kind_valid",
        ),
        CheckConstraint(
            "(CASE WHEN tactic_id IS NOT NULL THEN 1 ELSE 0 END) "
            "+ (CASE WHEN technique_id IS NOT NULL THEN 1 ELSE 0 END) "
            "+ (CASE WHEN subtechnique_id IS NOT NULL THEN 1 ELSE 0 END) = 1",
            name="exactly_one_mitre_fk",
        ),
        UniqueConstraint(
            "test_template_id",
            "tactic_id",
            "technique_id",
            "subtechnique_id",
            name="uq_test_template_mitre_tag",
        ),
        Index("ix_test_template_mitre_tags_template", "test_template_id"),
    )


class ScenarioTemplate(Base, UuidPkMixin, TimestampMixin, SoftDeleteMixin):
    __tablename__ = "scenario_templates"

    name: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str | None] = mapped_column(Text, nullable=True)

    tests: Mapped[list["ScenarioTemplateTest"]] = relationship(
        back_populates="scenario_template",
        cascade="all, delete-orphan",
        order_by="ScenarioTemplateTest.position",
    )

    __table_args__ = (
        Index(
            "ix_scenario_templates_active",
            "deleted_at",
            postgresql_where="deleted_at IS NULL",
        ),
        Index("ix_scenario_templates_name", "name"),
    )


class ScenarioTemplateTest(Base, UuidPkMixin):
    """Ordered membership of a test template inside a scenario template.

    UUID PK + UNIQUE(scenario_template_id, position) lets the same test appear
    multiple times at different positions (chained operations are common in
    purple-team scenarios).
    """

    __tablename__ = "scenario_template_tests"

    scenario_template_id: Mapped[uuid.UUID] = mapped_column(
        Uuid(as_uuid=True),
        ForeignKey("scenario_templates.id", ondelete="CASCADE"),
        nullable=False,
    )
    test_template_id: Mapped[uuid.UUID] = mapped_column(
        Uuid(as_uuid=True),
        ForeignKey("test_templates.id", ondelete="RESTRICT"),
        nullable=False,
    )
    position: Mapped[int] = mapped_column(Integer, nullable=False)

    scenario_template: Mapped[ScenarioTemplate] = relationship(back_populates="tests")

    __table_args__ = (
        UniqueConstraint(
            "scenario_template_id",
            "position",
            name="uq_scenario_template_tests_position",
        ),
        Index("ix_scenario_template_tests_scenario", "scenario_template_id"),
        Index("ix_scenario_template_tests_test", "test_template_id"),
    )