Files
Aegis/backend/app/models/test.py

70 lines
4.1 KiB
Python

import uuid
from datetime import datetime
from sqlalchemy import Column, String, Text, Boolean, Integer, DateTime, ForeignKey, Enum
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import relationship
from app.database import Base
from app.models.enums import TestState, TestResult
class Test(Base):
"""
Test model representing a security test for a MITRE ATT&CK technique.
Each test documents an attempt to validate coverage of a specific technique,
including the procedure, tools used, and outcome. V2 introduces dual
validation: Red Lead and Blue Lead must each approve independently.
"""
__tablename__ = "tests"
# ── Core fields ─────────────────────────────────────────────────
id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
technique_id = Column(UUID(as_uuid=True), ForeignKey("techniques.id"), nullable=False)
name = Column(String, nullable=False)
description = Column(Text, nullable=True)
platform = Column(String, nullable=True)
procedure_text = Column(Text, nullable=True)
tool_used = Column(String, nullable=True)
execution_date = Column(DateTime, nullable=True)
created_by = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
result = Column(Enum(TestResult, name="testresult"), nullable=True)
state = Column(Enum(TestState, name="teststate"), default=TestState.draft)
created_at = Column(DateTime, default=datetime.utcnow)
# ── Red Team fields ─────────────────────────────────────────────
red_summary = Column(Text, nullable=True)
attack_success = Column(Boolean, nullable=True)
red_validated_by = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
red_validated_at = Column(DateTime, nullable=True)
red_validation_status = Column(String, nullable=True) # pending / approved / rejected
red_validation_notes = Column(Text, nullable=True)
# ── Blue Team fields ────────────────────────────────────────────
blue_summary = Column(Text, nullable=True)
detection_result = Column(Enum(TestResult, name="testresult"), nullable=True)
blue_validated_by = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
blue_validated_at = Column(DateTime, nullable=True)
blue_validation_status = Column(String, nullable=True) # pending / approved / rejected
blue_validation_notes = Column(Text, nullable=True)
# ── Remediation fields ───────────────────────────────────────────
remediation_steps = Column(Text, nullable=True)
remediation_status = Column(String, nullable=True) # pending / in_progress / completed / not_applicable
remediation_assignee = Column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
# ── Re-test fields ────────────────────────────────────────────
retest_of = Column(UUID(as_uuid=True), ForeignKey("tests.id"), nullable=True)
retest_count = Column(Integer, default=0)
# ── Relationships ───────────────────────────────────────────────
technique = relationship("Technique", back_populates="tests")
evidences = relationship("Evidence", back_populates="test")
creator = relationship("User", foreign_keys=[created_by])
red_validator = relationship("User", foreign_keys=[red_validated_by])
blue_validator = relationship("User", foreign_keys=[blue_validated_by])
remediation_user = relationship("User", foreign_keys=[remediation_assignee])
original_test = relationship("Test", remote_side="Test.id", foreign_keys=[retest_of])
retests = relationship("Test", foreign_keys=[retest_of], back_populates="original_test")