158 lines
5.3 KiB
Python
158 lines
5.3 KiB
Python
"""Pydantic schemas for Test endpoints."""
|
|
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
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 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
|
|
|
|
# Remediation fields
|
|
remediation_steps: str | None = None
|
|
remediation_status: str | None = None
|
|
remediation_assignee: uuid.UUID | None = None
|
|
|
|
# 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)
|