kitos
c99cc4946a
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.
311 lines
11 KiB
Python
311 lines
11 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 (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
|
|
|
|
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.
|
|
"""
|
|
# Assign tests = [
|
|
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):
|
|
# Assign results = [t.detection_result for t in tests if t.detection_result]
|
|
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):
|
|
# Assign self.status_global = TechniqueStatus.validated
|
|
self.status_global = TechniqueStatus.validated
|
|
# elif any(
|
|
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
|