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

321 lines
12 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)
"""
# Enable future language features for compatibility
from __future__ import annotations
# Import uuid
import uuid
# Import dataclass, field from dataclasses
from dataclasses import dataclass, field
# Import datetime from datetime
from datetime import datetime
# Import TYPE_CHECKING from typing
from typing import TYPE_CHECKING
# Import TechniqueStatus, TestResult, TestState from app.domain.enums
from app.domain.enums import TechniqueStatus, TestResult, TestState
# Import MitreId from app.domain.value_objects.mitre_id
from app.domain.value_objects.mitre_id import MitreId
# Check: TYPE_CHECKING
if TYPE_CHECKING:
# Import Technique as TechniqueORM from app.models.technique
from app.models.technique import Technique as TechniqueORM
# Apply the @dataclass decorator
@dataclass(frozen=True)
# Define class _TestSnapshot
class _TestSnapshot:
"""Minimal read-only view of a test for status calculation."""
# state: TestState
state: TestState
# detection_result: str | None
detection_result: str | None
# Apply the @dataclass decorator
@dataclass
# Define class TechniqueEntity
class TechniqueEntity:
"""Pure domain representation of a MITRE ATT&CK technique."""
# id: uuid.UUID
id: uuid.UUID
# mitre_id: str
mitre_id: str
# name: str
name: str
# Assign tactic = None
tactic: str | None = None
# Assign description = None
description: str | None = None
# Assign platforms = field(default_factory=list)
platforms: list[str] = field(default_factory=list)
# Assign is_subtechnique = False
is_subtechnique: bool = False
# Assign parent_mitre_id = None
parent_mitre_id: str | None = None
# Assign status_global = TechniqueStatus.not_evaluated
status_global: TechniqueStatus = TechniqueStatus.not_evaluated
# Assign review_required = False
review_required: bool = False
# Assign last_review_date = None
last_review_date: datetime | None = None
# Assign mitre_version = None
mitre_version: str | None = None
# Assign mitre_last_modified = None
mitre_last_modified: datetime | None = None
# -- Factory -----------------------------------------------------------
@classmethod
# Define function create
def create(
cls,
*,
# Entry: mitre_id
mitre_id: str,
# Entry: name
name: str,
# Entry: tactic
tactic: str | None = None,
# Entry: description
description: str | None = None,
# Entry: platforms
platforms: list[str] | None = None,
) -> TechniqueEntity:
"""Create a new technique, validating the MITRE ID format.
Args:
mitre_id (str): MITRE ATT&CK identifier (e.g. ``"T1059"`` or ``"T1059.001"``).
name (str): Human-readable name of the technique.
tactic (str | None): MITRE tactic category the technique belongs to.
description (str | None): Optional free-text description.
platforms (list[str] | None): List of platform strings the technique applies to.
Returns:
TechniqueEntity: A new entity with a freshly generated UUID and
``status_global`` set to ``not_evaluated``.
"""
# Assign validated_id = MitreId(mitre_id)
validated_id = MitreId(mitre_id)
# Return cls(
return cls(
# Keyword argument: id
id=uuid.uuid4(),
# Keyword argument: mitre_id
mitre_id=validated_id.value,
# Keyword argument: name
name=name,
# Keyword argument: tactic
tactic=tactic,
# Keyword argument: description
description=description,
# Keyword argument: platforms
platforms=platforms or [],
# Keyword argument: is_subtechnique
is_subtechnique=validated_id.is_subtechnique,
# Keyword argument: parent_mitre_id
parent_mitre_id=validated_id.parent_id,
# Keyword argument: status_global
status_global=TechniqueStatus.not_evaluated,
)
# Apply the @classmethod decorator
@classmethod
# Define function from_orm
def from_orm(cls, model: TechniqueORM) -> TechniqueEntity:
"""Build a TechniqueEntity from a SQLAlchemy Technique model.
Args:
model (TechniqueORM): The ORM model instance to convert.
Returns:
TechniqueEntity: A fully populated domain entity reflecting the ORM state.
"""
# Assign raw_status = model.status_global
raw_status = model.status_global
# Check: raw_status is None
if raw_status is None:
# Assign status = TechniqueStatus.not_evaluated
status = TechniqueStatus.not_evaluated
# Alternative: isinstance(raw_status, TechniqueStatus)
elif isinstance(raw_status, TechniqueStatus):
# Assign status = raw_status
status = raw_status
# Fallback: handle remaining cases
else:
# Assign status = TechniqueStatus(raw_status)
status = TechniqueStatus(raw_status)
# Return cls(
return cls(
# Keyword argument: id
id=model.id,
# Keyword argument: mitre_id
mitre_id=model.mitre_id,
# Keyword argument: name
name=model.name,
# Keyword argument: tactic
tactic=model.tactic,
# Keyword argument: description
description=model.description,
# Keyword argument: platforms
platforms=model.platforms or [],
# Keyword argument: is_subtechnique
is_subtechnique=model.is_subtechnique or False,
# Keyword argument: parent_mitre_id
parent_mitre_id=model.parent_mitre_id,
# Keyword argument: status_global
status_global=status,
# Keyword argument: review_required
review_required=model.review_required or False,
# Keyword argument: last_review_date
last_review_date=model.last_review_date,
# Keyword argument: mitre_version
mitre_version=getattr(model, "mitre_version", None),
# Keyword argument: mitre_last_modified
mitre_last_modified=getattr(model, "mitre_last_modified", None),
)
# Define function apply_to
def apply_to(self, model: TechniqueORM) -> None:
"""Copy mutable fields back onto the ORM model.
Args:
model (TechniqueORM): The ORM model to update in-place.
Returns:
None
"""
# Assign model.status_global = self.status_global
model.status_global = self.status_global
# Assign model.review_required = self.review_required
model.review_required = self.review_required
# Assign model.last_review_date = self.last_review_date
model.last_review_date = self.last_review_date
# -- Business logic ----------------------------------------------------
def recalculate_status(
self,
# Entry: test_snapshots
test_snapshots: list[tuple[str, str | None]],
) -> TechniqueStatus:
"""Recompute ``status_global`` from a list of (state, detection_result) pairs.
Rules (v3):
1. No tests -> not_evaluated
2. All tests validated -> inspect detection results:
a. All detected AND ≥ 2 validated tests -> validated
b. All detected but only 1 validated test -> partial
(single test is not enough evidence for full coverage)
c. Any partially_detected -> partial
d. Otherwise (no detected results) -> not_covered
3. Some validated, others in intermediate states -> partial
4. All tests in intermediate states (draft/executing/evaluating/review/rejected)
-> in_progress
Minimum validated count for "validated": 2 tests.
With only 1 validated+detected test the technique is "partial" to
signal that more testing is recommended.
Args:
test_snapshots (list[tuple[str, str | None]]): Each element is a
``(state, detection_result)`` pair where *state* is a
:class:`TestState` value string and *detection_result* is a
:class:`TestResult` value string or ``None``.
Returns:
TechniqueStatus: The newly computed status, which is also stored on
the entity's ``status_global`` field.
"""
_MIN_VALIDATED_FOR_FULL = 2 # require ≥ N validated tests for "validated"
tests = [
_TestSnapshot(
# Keyword argument: state
state=s if isinstance(s, TestState) else TestState(s),
# Keyword argument: detection_result
detection_result=dr,
)
for s, dr in test_snapshots
]
# Check: not tests
if not tests:
# Assign self.status_global = TechniqueStatus.not_evaluated
self.status_global = TechniqueStatus.not_evaluated
# Alternative: all(t.state == TestState.validated for t in tests)
elif all(t.state == TestState.validated for t in tests):
validated_count = len(tests)
results = [t.detection_result for t in tests if t.detection_result]
# Check: results and all(r == TestResult.detected or r == "detected" for r i...
if results and all(r == TestResult.detected or r == "detected" for r in results):
# Need at least _MIN_VALIDATED_FOR_FULL tests for "validated"
if validated_count >= _MIN_VALIDATED_FOR_FULL:
self.status_global = TechniqueStatus.validated
else:
self.status_global = TechniqueStatus.partial
elif any(
# Keyword argument: r
r == TestResult.partially_detected or r == "partially_detected"
for r in results
):
# Assign self.status_global = TechniqueStatus.partial
self.status_global = TechniqueStatus.partial
# Fallback: handle remaining cases
else:
# Assign self.status_global = TechniqueStatus.not_covered
self.status_global = TechniqueStatus.not_covered
# Alternative: any(t.state == TestState.validated for t in tests)
elif any(t.state == TestState.validated for t in tests):
# Assign self.status_global = TechniqueStatus.partial
self.status_global = TechniqueStatus.partial
# Fallback: handle remaining cases
else:
# Assign self.status_global = TechniqueStatus.in_progress
self.status_global = TechniqueStatus.in_progress
# Return self.status_global
return self.status_global
# Define function mark_reviewed
def mark_reviewed(self) -> None:
"""Mark the technique as reviewed, clearing the review flag.
Returns:
None
"""
# Assign self.review_required = False
self.review_required = False
# Assign self.last_review_date = datetime.utcnow()
self.last_review_date = datetime.utcnow()
# Define function flag_for_review
def flag_for_review(self) -> None:
"""Flag the technique as needing review.
Returns:
None
"""
# Assign self.review_required = True
self.review_required = True