Files
Aegis/backend/app/domain/entities/technique.py

165 lines
5.6 KiB
Python

"""TechniqueEntity — pure domain object for a MITRE ATT&CK technique.
Owns the status recalculation logic that was previously in
``status_service.py``. Has **no** dependency on FastAPI, SQLAlchemy,
or any infrastructure concern.
Usage::
entity = TechniqueEntity.from_orm(technique_orm_model)
entity.recalculate_status(test_states_and_results)
entity.mark_reviewed()
entity.apply_to(technique_orm_model)
"""
from __future__ import annotations
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
from app.domain.enums import TechniqueStatus, TestResult, TestState
from app.domain.value_objects.mitre_id import MitreId
@dataclass(frozen=True)
class _TestSnapshot:
"""Minimal read-only view of a test for status calculation."""
state: TestState
detection_result: str | None
@dataclass
class TechniqueEntity:
"""Pure domain representation of a MITRE ATT&CK technique."""
id: uuid.UUID
mitre_id: str
name: str
tactic: str | None = None
description: str | None = None
platforms: list[str] = field(default_factory=list)
is_subtechnique: bool = False
parent_mitre_id: str | None = None
status_global: TechniqueStatus = TechniqueStatus.not_evaluated
review_required: bool = False
last_review_date: datetime | None = None
mitre_version: str | None = None
mitre_last_modified: datetime | None = None
# -- Factory -----------------------------------------------------------
@classmethod
def create(
cls,
*,
mitre_id: str,
name: str,
tactic: str | None = None,
description: str | None = None,
platforms: list[str] | None = None,
) -> TechniqueEntity:
"""Create a new technique, validating the MITRE ID format."""
validated_id = MitreId(mitre_id)
return cls(
id=uuid.uuid4(),
mitre_id=validated_id.value,
name=name,
tactic=tactic,
description=description,
platforms=platforms or [],
is_subtechnique=validated_id.is_subtechnique,
parent_mitre_id=validated_id.parent_id,
status_global=TechniqueStatus.not_evaluated,
)
@classmethod
def from_orm(cls, model: Any) -> TechniqueEntity:
"""Build a TechniqueEntity from a SQLAlchemy Technique model."""
raw_status = model.status_global
if raw_status is None:
status = TechniqueStatus.not_evaluated
elif isinstance(raw_status, TechniqueStatus):
status = raw_status
else:
status = TechniqueStatus(raw_status)
return cls(
id=model.id,
mitre_id=model.mitre_id,
name=model.name,
tactic=model.tactic,
description=model.description,
platforms=model.platforms or [],
is_subtechnique=model.is_subtechnique or False,
parent_mitre_id=model.parent_mitre_id,
status_global=status,
review_required=model.review_required or False,
last_review_date=model.last_review_date,
mitre_version=getattr(model, "mitre_version", None),
mitre_last_modified=getattr(model, "mitre_last_modified", None),
)
def apply_to(self, model: Any) -> None:
"""Copy mutable fields back onto the ORM model."""
model.status_global = self.status_global
model.review_required = self.review_required
model.last_review_date = self.last_review_date
# -- Business logic ----------------------------------------------------
def recalculate_status(
self,
test_snapshots: list[tuple[str, str | None]],
) -> TechniqueStatus:
"""Recompute ``status_global`` from a list of (state, detection_result) pairs.
Rules (v2):
1. No tests -> not_evaluated
2. All validated -> inspect detection results:
- All detected -> validated
- Any partially_detected -> partial
- Otherwise -> not_covered
3. Some validated, others in progress -> partial
4. All in intermediate states -> in_progress
Returns the new status (also set on the entity).
"""
tests = [
_TestSnapshot(
state=s if isinstance(s, TestState) else TestState(s),
detection_result=dr,
)
for s, dr in test_snapshots
]
if not tests:
self.status_global = TechniqueStatus.not_evaluated
elif all(t.state == TestState.validated for t in tests):
results = [t.detection_result for t in tests if t.detection_result]
if results and all(r == TestResult.detected or r == "detected" for r in results):
self.status_global = TechniqueStatus.validated
elif any(
r == TestResult.partially_detected or r == "partially_detected"
for r in results
):
self.status_global = TechniqueStatus.partial
else:
self.status_global = TechniqueStatus.not_covered
elif any(t.state == TestState.validated for t in tests):
self.status_global = TechniqueStatus.partial
else:
self.status_global = TechniqueStatus.in_progress
return self.status_global
def mark_reviewed(self) -> None:
"""Mark the technique as reviewed, clearing the review flag."""
self.review_required = False
self.last_review_date = datetime.utcnow()
def flag_for_review(self) -> None:
"""Flag the technique as needing review."""
self.review_required = True