refactor(techniques): wire TechniqueRepository into techniques router replacing direct db.query() with repo pattern, domain exceptions, and UnitOfWork

This commit is contained in:
2026-02-19 15:13:52 +01:00
parent 0b65f51d1c
commit 2b6d9090c9
4 changed files with 97 additions and 79 deletions

View File

@@ -46,6 +46,8 @@ class TechniqueEntity:
status_global: TechniqueStatus = TechniqueStatus.not_evaluated status_global: TechniqueStatus = TechniqueStatus.not_evaluated
review_required: bool = False review_required: bool = False
last_review_date: datetime | None = None last_review_date: datetime | None = None
mitre_version: str | None = None
mitre_last_modified: datetime | None = None
# -- Factory ----------------------------------------------------------- # -- Factory -----------------------------------------------------------
@@ -94,6 +96,8 @@ class TechniqueEntity:
status_global=status, status_global=status,
review_required=model.review_required or False, review_required=model.review_required or False,
last_review_date=model.last_review_date, 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: def apply_to(self, model: Any) -> None:

View File

@@ -163,6 +163,8 @@ class SATechniqueRepository:
existing.platforms = technique.platforms existing.platforms = technique.platforms
existing.is_subtechnique = technique.is_subtechnique existing.is_subtechnique = technique.is_subtechnique
existing.parent_mitre_id = technique.parent_mitre_id existing.parent_mitre_id = technique.parent_mitre_id
existing.mitre_version = technique.mitre_version
existing.mitre_last_modified = technique.mitre_last_modified
self._session.flush() self._session.flush()
return TechniqueMapper.to_entity(existing) return TechniqueMapper.to_entity(existing)
else: else:
@@ -178,6 +180,8 @@ class SATechniqueRepository:
status_global=technique.status_global, status_global=technique.status_global,
review_required=technique.review_required, review_required=technique.review_required,
last_review_date=technique.last_review_date, last_review_date=technique.last_review_date,
mitre_version=technique.mitre_version,
mitre_last_modified=technique.mitre_last_modified,
) )
self._session.add(model) self._session.add(model)
self._session.flush() self._session.flush()

View File

@@ -1,13 +1,23 @@
"""CRUD router for MITRE ATT&CK Techniques.""" """CRUD router for MITRE ATT&CK Techniques.
from datetime import datetime Uses the TechniqueRepository for data access and domain exceptions
for error signaling. The error_handler middleware maps domain
exceptions to HTTP responses automatically.
"""
from fastapi import APIRouter, Depends, HTTPException, Query, status from fastapi import APIRouter, Depends, Query, status
from sqlalchemy.orm import Session, joinedload from sqlalchemy.orm import Session, joinedload
from app.database import get_db from app.database import get_db
from app.dependencies.auth import get_current_user, require_role, require_any_role from app.dependencies.auth import get_current_user, require_role, require_any_role
from app.models.enums import TechniqueStatus from app.dependencies.repositories import get_technique_repository
from app.domain.entities.technique import TechniqueEntity
from app.domain.errors import DuplicateEntityError, EntityNotFoundError
from app.domain.enums import TechniqueStatus
from app.domain.unit_of_work import UnitOfWork
from app.infrastructure.persistence.repositories.sa_technique_repository import (
SATechniqueRepository,
)
from app.models.technique import Technique from app.models.technique import Technique
from app.models.user import User from app.models.user import User
from app.schemas.technique import ( from app.schemas.technique import (
@@ -34,24 +44,19 @@ def list_techniques(
None, alias="status", description="Filter by global status" None, alias="status", description="Filter by global status"
), ),
review_required: bool | None = Query(None, description="Filter by review flag"), review_required: bool | None = Query(None, description="Filter by review flag"),
db: Session = Depends(get_db), repo: SATechniqueRepository = Depends(get_technique_repository),
current_user: User = Depends(get_current_user), current_user: User = Depends(get_current_user),
): ):
"""Return a lightweight list of techniques, optionally filtered.""" """Return a lightweight list of techniques, optionally filtered."""
query = db.query(Technique) return repo.list_all(
tactic=tactic,
if tactic is not None: status=status_global,
query = query.filter(Technique.tactic == tactic) review_required=review_required,
if status_global is not None: )
query = query.filter(Technique.status_global == status_global)
if review_required is not None:
query = query.filter(Technique.review_required == review_required)
return query.order_by(Technique.mitre_id).all()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# GET /techniques/{mitre_id} — detail (with tests) # GET /techniques/{mitre_id} — detail (with tests + D3FEND)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -70,12 +75,8 @@ def get_technique(
) )
if technique is None: if technique is None:
raise HTTPException( raise EntityNotFoundError("Technique", mitre_id)
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Technique {mitre_id} not found",
)
# Build response dict manually to include D3FEND defenses
defenses = get_defenses_for_technique(db, technique.id) defenses = get_defenses_for_technique(db, technique.id)
return { return {
@@ -120,34 +121,35 @@ def get_technique(
def create_technique( def create_technique(
payload: TechniqueCreate, payload: TechniqueCreate,
db: Session = Depends(get_db), db: Session = Depends(get_db),
repo: SATechniqueRepository = Depends(get_technique_repository),
current_user: User = Depends(require_role("admin")), current_user: User = Depends(require_role("admin")),
): ):
"""Create a new technique manually.""" """Create a new technique manually."""
# Ensure mitre_id is unique if repo.exists_by_mitre_id(payload.mitre_id):
existing = ( raise DuplicateEntityError("Technique", "mitre_id", payload.mitre_id)
db.query(Technique).filter(Technique.mitre_id == payload.mitre_id).first()
)
if existing is not None:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail=f"Technique with mitre_id '{payload.mitre_id}' already exists",
)
technique = Technique(**payload.model_dump()) entity = TechniqueEntity.create(
db.add(technique) mitre_id=payload.mitre_id,
db.commit() name=payload.name,
db.refresh(technique) description=payload.description,
tactic=payload.tactic,
platforms=payload.platforms,
)
with UnitOfWork(db) as uow:
saved = repo.save(entity)
uow.commit()
log_action( log_action(
db, db,
user_id=current_user.id, user_id=current_user.id,
action="create_technique", action="create_technique",
entity_type="technique", entity_type="technique",
entity_id=technique.id, entity_id=saved.id,
details={"mitre_id": technique.mitre_id, "name": technique.name}, details={"mitre_id": saved.mitre_id, "name": saved.name},
) )
return technique return saved
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -160,36 +162,32 @@ def update_technique(
mitre_id: str, mitre_id: str,
payload: TechniqueUpdate, payload: TechniqueUpdate,
db: Session = Depends(get_db), db: Session = Depends(get_db),
repo: SATechniqueRepository = Depends(get_technique_repository),
current_user: User = Depends(require_role("admin")), current_user: User = Depends(require_role("admin")),
): ):
"""Update one or more fields of an existing technique.""" """Update one or more fields of an existing technique."""
technique = ( entity = repo.find_by_mitre_id(mitre_id)
db.query(Technique).filter(Technique.mitre_id == mitre_id).first() if entity is None:
) raise EntityNotFoundError("Technique", mitre_id)
if technique is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Technique {mitre_id} not found",
)
update_data = payload.model_dump(exclude_unset=True) update_data = payload.model_dump(exclude_unset=True)
for field, value in update_data.items(): for field, value in update_data.items():
setattr(technique, field, value) setattr(entity, field, value)
db.commit() with UnitOfWork(db) as uow:
db.refresh(technique) saved = repo.save(entity)
uow.commit()
log_action( log_action(
db, db,
user_id=current_user.id, user_id=current_user.id,
action="update_technique", action="update_technique",
entity_type="technique", entity_type="technique",
entity_id=technique.id, entity_id=saved.id,
details={"mitre_id": mitre_id, "updated_fields": list(update_data.keys())}, details={"mitre_id": mitre_id, "updated_fields": list(update_data.keys())},
) )
return technique return saved
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -201,6 +199,7 @@ def update_technique(
def review_technique( def review_technique(
mitre_id: str, mitre_id: str,
db: Session = Depends(get_db), db: Session = Depends(get_db),
repo: SATechniqueRepository = Depends(get_technique_repository),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")), current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
): ):
"""Mark a technique as reviewed. """Mark a technique as reviewed.
@@ -208,29 +207,23 @@ def review_technique(
Sets ``review_required`` to *False* and records the current timestamp Sets ``review_required`` to *False* and records the current timestamp
in ``last_review_date``. in ``last_review_date``.
""" """
technique = ( entity = repo.find_by_mitre_id(mitre_id)
db.query(Technique).filter(Technique.mitre_id == mitre_id).first() if entity is None:
) raise EntityNotFoundError("Technique", mitre_id)
if technique is None: entity.mark_reviewed()
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Technique {mitre_id} not found",
)
technique.review_required = False with UnitOfWork(db) as uow:
technique.last_review_date = datetime.utcnow() saved = repo.save(entity)
uow.commit()
db.commit()
db.refresh(technique)
log_action( log_action(
db, db,
user_id=current_user.id, user_id=current_user.id,
action="review_technique", action="review_technique",
entity_type="technique", entity_type="technique",
entity_id=technique.id, entity_id=saved.id,
details={"mitre_id": mitre_id}, details={"mitre_id": mitre_id},
) )
return technique return saved

View File

@@ -1,26 +1,43 @@
# Aegis — Task Tracker # Aegis — Architectural Refactoring Task Tracker
## In Progress ## Tier 1 — Quick Wins
- [ ] Clean Architecture foundation: domain enums, value objects, entities, repository ports + implementations - [ ] QW-1: Wire existing repos into `techniques.py` router
- [ ] QW-2: Fix `audit_service` to follow UoW (no direct `db.commit()`)
- [ ] QW-3: Consolidate `status_service` with `TechniqueEntity.recalculate_status()`
- [ ] QW-4: Remove remaining `HTTPException` from services
## Completed ## Tier 2 — Service Extraction (fat routers → thin routers + services)
- [ ] SE-1: Extract reports service from `reports.py`
- [ ] SE-2: Extract metrics service from `metrics.py`
- [ ] SE-3: Extract compliance service from `compliance.py`
- [ ] SE-4: Extract detection_rules service from `detection_rules.py`
- [ ] SE-5: Extract threat_actors service from `threat_actors.py`
## Tier 3 — Architectural Fixes
- [ ] AF-1: Persist scoring weights in DB (replace mutable `settings`)
- [ ] AF-2: Slim `tests.py` router (CRUD to repo/service)
- [ ] AF-3: Slim `evidence.py` router (permissions to domain)
- [ ] AF-4: Slim `campaigns.py` router (CRUD to service)
## Tier 4 — Polish
- [ ] P-1: Structured JSON logging
- [ ] P-2: Create architecture skill file for future agents
## Completed (prior sessions)
- [x] Domain exceptions hierarchy (domain/errors.py) - [x] Domain exceptions hierarchy (domain/errors.py)
- [x] TestEntity with state machine (domain/test_entity.py) - [x] TestEntity with state machine (domain/test_entity.py)
- [x] TechniqueEntity (domain/entities/technique.py)
- [x] Value objects: MitreId, ScoringWeights
- [x] Unit of Work (domain/unit_of_work.py) - [x] Unit of Work (domain/unit_of_work.py)
- [x] Error handler middleware (middleware/error_handler.py) - [x] Error handler middleware (middleware/error_handler.py)
- [x] Redis-backed token blacklist (auth.py) - [x] Redis-backed token blacklist (auth.py)
- [x] CI pipeline (.github/workflows/ci.yml) - [x] CI pipeline (.github/workflows/ci.yml)
- [x] Heatmap service extracted (services/heatmap_service.py) - [x] Heatmap service extracted (services/heatmap_service.py)
- [x] Scoring bulk queries (bulk_technique_scores) - [x] Scoring bulk queries (bulk_technique_scores)
- [x] Architecture skill file (.cursor/rules/aegis-architecture.md) - [x] Repository ports + implementations (Technique, Test)
- [x] Agent validation script (scripts/agent_validate_backend.sh) - [x] Agent validation script (scripts/agent_validate_backend.sh)
## Backlog
- [ ] Application layer use cases
- [ ] Migrate fat routers to use repositories
- [ ] Scoring config persistence (DB instead of mutable settings)
- [ ] Structured JSON logging
- [ ] Frontend type generation from OpenAPI