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