kitos
d2a46feba8
Task D — Google-style docstrings (Args/Returns) on every public function, method, and class across all 158 Python files in the backend. Zero ruff D violations (pydocstyle Google convention). Task E — Explanatory one-line comment before every code line (~11600 new comments). ruff check passes clean after isort re-sort.
262 lines
8.9 KiB
Python
262 lines
8.9 KiB
Python
"""Pydantic schemas for Test endpoints."""
|
|
|
|
# Import uuid
|
|
import uuid
|
|
|
|
# Import datetime from datetime
|
|
from datetime import datetime
|
|
|
|
# Import BaseModel, ConfigDict from pydantic
|
|
from pydantic import BaseModel, ConfigDict
|
|
|
|
# Import DataClassification from app.domain.enums
|
|
from app.domain.enums import DataClassification
|
|
|
|
# Import TestResult, TestState from app.models.enums
|
|
from app.models.enums import TestResult, TestState
|
|
|
|
# ── Create ──────────────────────────────────────────────────────────
|
|
|
|
|
|
class TestCreate(BaseModel):
|
|
"""Payload for creating a new test."""
|
|
|
|
# technique_id: uuid.UUID
|
|
technique_id: uuid.UUID
|
|
# name: str
|
|
name: str
|
|
# Assign description = None
|
|
description: str | None = None
|
|
# Assign platform = None
|
|
platform: str | None = None
|
|
# Assign procedure_text = None
|
|
procedure_text: str | None = None
|
|
# Assign tool_used = None
|
|
tool_used: str | None = None
|
|
|
|
|
|
# ── Update (general) ───────────────────────────────────────────────
|
|
|
|
|
|
class TestClassificationUpdate(BaseModel):
|
|
"""Admin-only payload for changing data classification."""
|
|
|
|
# data_classification: DataClassification
|
|
data_classification: DataClassification
|
|
|
|
|
|
# Define class TestUpdate
|
|
class TestUpdate(BaseModel):
|
|
"""Payload for partially updating an existing test.
|
|
|
|
Every field is optional so callers send only what changed.
|
|
"""
|
|
|
|
# Assign name = None
|
|
name: str | None = None
|
|
# Assign description = None
|
|
description: str | None = None
|
|
# Assign platform = None
|
|
platform: str | None = None
|
|
# Assign procedure_text = None
|
|
procedure_text: str | None = None
|
|
# Assign tool_used = None
|
|
tool_used: str | None = None
|
|
# Assign result = None
|
|
result: TestResult | None = None
|
|
|
|
|
|
# ── Red Team update ────────────────────────────────────────────────
|
|
|
|
|
|
class TestRedUpdate(BaseModel):
|
|
"""Fields that Red Team fills in during the red_executing phase."""
|
|
|
|
# Assign name = None
|
|
name: str | None = None
|
|
# Assign description = None
|
|
description: str | None = None
|
|
# Assign procedure_text = None
|
|
procedure_text: str | None = None
|
|
# Assign tool_used = None
|
|
tool_used: str | None = None
|
|
# Assign attack_success = None
|
|
attack_success: bool | None = None
|
|
# Assign red_summary = None
|
|
red_summary: str | None = None
|
|
|
|
|
|
# ── Blue Team update ───────────────────────────────────────────────
|
|
|
|
|
|
class TestBlueUpdate(BaseModel):
|
|
"""Fields that Blue Team fills in during the blue_evaluating phase."""
|
|
|
|
# Assign detection_result = None
|
|
detection_result: TestResult | None = None
|
|
# Assign blue_summary = 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_status: str # "approved" or "rejected"
|
|
# Assign red_validation_notes = None
|
|
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_status: str # "approved" or "rejected"
|
|
# Assign blue_validation_notes = None
|
|
blue_validation_notes: str | None = None
|
|
|
|
|
|
# ── Remediation update ────────────────────────────────────────────
|
|
|
|
|
|
class TestRemediationUpdate(BaseModel):
|
|
"""Payload for updating remediation fields."""
|
|
|
|
# Assign remediation_steps = None
|
|
remediation_steps: str | None = None
|
|
# Assign remediation_status = None # pending / in_progress / completed / not_applicable
|
|
remediation_status: str | None = None # pending / in_progress / completed / not_applicable
|
|
# Assign remediation_assignee = None
|
|
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
|
|
result: TestResult
|
|
# Assign comments = None
|
|
comments: str | None = None
|
|
|
|
|
|
# ── Read (full) ─────────────────────────────────────────────────────
|
|
|
|
|
|
class TestOut(BaseModel):
|
|
"""Complete representation returned by the API."""
|
|
|
|
# id: uuid.UUID
|
|
id: uuid.UUID
|
|
# technique_id: uuid.UUID
|
|
technique_id: uuid.UUID
|
|
# name: str
|
|
name: str
|
|
# Assign description = None
|
|
description: str | None = None
|
|
# Assign platform = None
|
|
platform: str | None = None
|
|
# Assign procedure_text = None
|
|
procedure_text: str | None = None
|
|
# Assign tool_used = None
|
|
tool_used: str | None = None
|
|
# Assign execution_date = None
|
|
execution_date: datetime | None = None
|
|
# Assign created_by = None
|
|
created_by: uuid.UUID | None = None
|
|
# Assign result = None
|
|
result: TestResult | None = None
|
|
# Assign state = TestState.draft
|
|
state: TestState = TestState.draft
|
|
# Assign created_at = None
|
|
created_at: datetime | None = None
|
|
|
|
# Red Team fields
|
|
red_summary: str | None = None
|
|
# Assign attack_success = None
|
|
attack_success: bool | None = None
|
|
# Assign red_validated_by = None
|
|
red_validated_by: uuid.UUID | None = None
|
|
# Assign red_validated_at = None
|
|
red_validated_at: datetime | None = None
|
|
# Assign red_validation_status = None
|
|
red_validation_status: str | None = None
|
|
# Assign red_validation_notes = None
|
|
red_validation_notes: str | None = None
|
|
|
|
# Blue Team fields
|
|
blue_summary: str | None = None
|
|
# Assign detection_result = None
|
|
detection_result: TestResult | None = None
|
|
# Assign blue_validated_by = None
|
|
blue_validated_by: uuid.UUID | None = None
|
|
# Assign blue_validated_at = None
|
|
blue_validated_at: datetime | None = None
|
|
# Assign blue_validation_status = None
|
|
blue_validation_status: str | None = None
|
|
# Assign blue_validation_notes = None
|
|
blue_validation_notes: str | None = None
|
|
|
|
# Phase timing fields (for Tempo worklogs)
|
|
red_started_at: datetime | None = None
|
|
# Assign blue_started_at = None
|
|
blue_started_at: datetime | None = None
|
|
# Assign paused_at = None
|
|
paused_at: datetime | None = None
|
|
# Assign red_paused_seconds = 0
|
|
red_paused_seconds: int = 0
|
|
# Assign blue_paused_seconds = 0
|
|
blue_paused_seconds: int = 0
|
|
|
|
# Remediation fields
|
|
remediation_steps: str | None = None
|
|
# Assign remediation_status = None
|
|
remediation_status: str | None = None
|
|
# Assign remediation_assignee = None
|
|
remediation_assignee: uuid.UUID | None = None
|
|
|
|
# Re-test fields
|
|
retest_of: uuid.UUID | None = None
|
|
# Assign retest_count = 0
|
|
retest_count: int = 0
|
|
# Assign data_classification = "internal"
|
|
data_classification: str = "internal"
|
|
|
|
# Technique info (populated when joined)
|
|
technique_mitre_id: str | None = None
|
|
# Assign technique_name = None
|
|
technique_name: str | None = None
|
|
|
|
# Assign model_config = ConfigDict(from_attributes=True)
|
|
model_config = ConfigDict(from_attributes=True)
|
|
|
|
# Apply the @classmethod decorator
|
|
@classmethod
|
|
# Define function model_validate
|
|
def model_validate(cls, obj: object, **kwargs: object) -> "TestOut":
|
|
"""Populate technique fields from the ORM relationship before validation.
|
|
|
|
Args:
|
|
obj (object): The ORM model instance (or any compatible object) to validate.
|
|
**kwargs (object): Additional keyword arguments forwarded to the parent.
|
|
|
|
Returns:
|
|
TestOut: The validated schema instance with technique fields populated.
|
|
"""
|
|
# Check: hasattr(obj, "technique") and obj.technique is not None
|
|
if hasattr(obj, "technique") and obj.technique is not None:
|
|
# Assign obj.__dict__["technique_mitre_id"] = obj.technique.mitre_id
|
|
obj.__dict__["technique_mitre_id"] = obj.technique.mitre_id
|
|
# Assign obj.__dict__["technique_name"] = obj.technique.name
|
|
obj.__dict__["technique_name"] = obj.technique.name
|
|
# Return super().model_validate(obj, **kwargs)
|
|
return super().model_validate(obj, **kwargs)
|