Aegis/backend/app/schemas/test.py

"""Pydantic schemas for Test endpoints."""

import uuid
from datetime import datetime

from pydantic import BaseModel, ConfigDict

from app.domain.enums import DataClassification
from app.models.enums import TestResult, TestState


# ── Create ──────────────────────────────────────────────────────────


class TestCreate(BaseModel):
    """Payload for creating a new test."""

    technique_id: uuid.UUID
    name: str
    description: str | None = None
    platform: str | None = None
    procedure_text: str | None = None
    tool_used: str | None = None


# ── Update (general) ───────────────────────────────────────────────


class TestClassificationUpdate(BaseModel):
    """Admin-only payload for changing data classification."""

    data_classification: DataClassification


class TestUpdate(BaseModel):
    """Payload for partially updating an existing test.
    Every field is optional so callers send only what changed."""

    name: str | None = None
    description: str | None = None
    platform: str | None = None
    procedure_text: str | None = None
    tool_used: str | None = None
    result: TestResult | None = None


# ── Red Team update ────────────────────────────────────────────────


class TestRedUpdate(BaseModel):
    """Fields that Red Team fills in during the red_executing phase."""

    name: str | None = None
    description: str | None = None
    procedure_text: str | None = None
    tool_used: str | None = None
    attack_success: bool | None = None
    red_summary: str | None = None


# ── Blue Team update ───────────────────────────────────────────────


class TestBlueUpdate(BaseModel):
    """Fields that Blue Team fills in during the blue_evaluating phase."""

    detection_result: TestResult | None = None
    blue_summary: str | None = None


# ── Red Lead validation ────────────────────────────────────────────


class TestRedValidate(BaseModel):
    """Payload sent by Red Lead to approve/reject the red side."""

    red_validation_status: str   # "approved" or "rejected"
    red_validation_notes: str | None = None


# ── Blue Lead validation ───────────────────────────────────────────


class TestBlueValidate(BaseModel):
    """Payload sent by Blue Lead to approve/reject the blue side."""

    blue_validation_status: str   # "approved" or "rejected"
    blue_validation_notes: str | None = None


# ── Remediation update ────────────────────────────────────────────


class TestRemediationUpdate(BaseModel):
    """Payload for updating remediation fields."""

    remediation_steps: str | None = None
    remediation_status: str | None = None  # pending / in_progress / completed / not_applicable
    remediation_assignee: uuid.UUID | None = None


# ── Legacy validate (kept for backwards compat) ────────────────────


class TestValidate(BaseModel):
    """Payload sent by a reviewer to validate / reject a test."""

    result: TestResult
    comments: str | None = None


# ── Read (full) ─────────────────────────────────────────────────────


class TestOut(BaseModel):
    """Complete representation returned by the API."""

    id: uuid.UUID
    technique_id: uuid.UUID
    name: str
    description: str | None = None
    platform: str | None = None
    procedure_text: str | None = None
    tool_used: str | None = None
    execution_date: datetime | None = None
    created_by: uuid.UUID | None = None
    result: TestResult | None = None
    state: TestState = TestState.draft
    created_at: datetime | None = None

    # Red Team fields
    red_summary: str | None = None
    attack_success: bool | None = None
    red_validated_by: uuid.UUID | None = None
    red_validated_at: datetime | None = None
    red_validation_status: str | None = None
    red_validation_notes: str | None = None

    # Blue Team fields
    blue_summary: str | None = None
    detection_result: TestResult | None = None
    blue_validated_by: uuid.UUID | None = None
    blue_validated_at: datetime | None = None
    blue_validation_status: str | None = None
    blue_validation_notes: str | None = None

    # Phase timing fields (for Tempo worklogs)
    red_started_at: datetime | None = None
    blue_started_at: datetime | None = None
    blue_work_started_at: datetime | None = None
    paused_at: datetime | None = None
    red_paused_seconds: int = 0
    blue_paused_seconds: int = 0

    # Remediation fields
    remediation_steps: str | None = None
    remediation_status: str | None = None
    remediation_assignee: uuid.UUID | None = None

    # Re-test fields
    retest_of: uuid.UUID | None = None
    retest_count: int = 0
    data_classification: str = "internal"

    # Technique info (populated when joined)
    technique_mitre_id: str | None = None
    technique_name: str | None = None

    model_config = ConfigDict(from_attributes=True)

    @classmethod
    def model_validate(cls, obj, **kwargs):
        """Override to populate technique fields from the relationship."""
        if hasattr(obj, "technique") and obj.technique is not None:
            obj.__dict__["technique_mitre_id"] = obj.technique.mitre_id
            obj.__dict__["technique_name"] = obj.technique.name
        return super().model_validate(obj, **kwargs)