Files
Aegis/backend/app/routers/tests.py
Kitos 7af6be10be feat(phase-11): implement Red/Blue business logic services (T-106, T-107, T-108)
T-106: Create test_workflow_service.py with state-machine transitions for the complete test lifecycle (draft -> red_executing -> blue_evaluating -> in_review -> validated/rejected), dual validation by Red/Blue leads, and reopen capability with field cleanup.

T-107: Update status_service.py to use detection_result from Blue Team instead of legacy result field, and differentiate between partial progress (some validated) vs all-in-progress states.

T-108: Create atomic_import_service.py that downloads the Atomic Red Team repo as a ZIP (avoiding API rate limits), parses all atomics YAML files, and creates idempotent TestTemplate records mapped to MITRE techniques.

Includes validation tests for all three tasks (19 checks total).
2026-02-09 09:58:54 +01:00

280 lines
8.1 KiB
Python

"""CRUD router for security Tests."""
import uuid
from datetime import datetime
from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session, joinedload
from app.database import get_db
from app.dependencies.auth import get_current_user, require_role, require_any_role
from app.models.enums import TestState
from app.models.technique import Technique
from app.models.test import Test
from app.models.user import User
from app.schemas.test import TestCreate, TestOut, TestUpdate, TestValidate
from app.services.audit_service import log_action
from app.services.status_service import recalculate_technique_status
router = APIRouter(prefix="/tests", tags=["tests"])
# ---------------------------------------------------------------------------
# POST /tests — create (red_tech or admin)
# ---------------------------------------------------------------------------
@router.post(
"",
response_model=TestOut,
status_code=status.HTTP_201_CREATED,
)
def create_test(
payload: TestCreate,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_tech")),
):
"""Create a new test linked to an existing technique.
The ``created_by`` field is set automatically to the current user and
``state`` defaults to *draft*.
"""
# Verify the parent technique exists
technique = db.query(Technique).filter(Technique.id == payload.technique_id).first()
if technique is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"Technique with id '{payload.technique_id}' not found",
)
test = Test(
**payload.model_dump(),
created_by=current_user.id,
state=TestState.draft,
)
db.add(test)
db.commit()
db.refresh(test)
log_action(
db,
user_id=current_user.id,
action="create_test",
entity_type="test",
entity_id=test.id,
details={"name": test.name, "technique_id": str(test.technique_id)},
)
return test
# ---------------------------------------------------------------------------
# GET /tests/{id} — detail (with evidences)
# ---------------------------------------------------------------------------
@router.get("/{test_id}", response_model=TestOut)
def get_test(
test_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
"""Return full details for a single test, including its evidences."""
test = (
db.query(Test)
.options(joinedload(Test.evidences))
.filter(Test.id == test_id)
.first()
)
if test is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Test not found",
)
return test
# ---------------------------------------------------------------------------
# PATCH /tests/{id} — update (creator or admin, only in draft/rejected)
# ---------------------------------------------------------------------------
@router.patch("/{test_id}", response_model=TestOut)
def update_test(
test_id: uuid.UUID,
payload: TestUpdate,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
"""Update one or more fields of an existing test.
Only the original creator or an admin can update.
The test must be in ``draft`` or ``rejected`` state.
"""
test = db.query(Test).filter(Test.id == test_id).first()
if test is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Test not found",
)
# Ownership / admin check
if current_user.role != "admin" and test.created_by != current_user.id:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Not enough permissions",
)
# State guard
if test.state not in (TestState.draft, TestState.rejected):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=f"Cannot update a test in '{test.state.value}' state (must be draft or rejected)",
)
update_data = payload.model_dump(exclude_unset=True)
for field, value in update_data.items():
setattr(test, field, value)
db.commit()
db.refresh(test)
log_action(
db,
user_id=current_user.id,
action="update_test",
entity_type="test",
entity_id=test.id,
details={"updated_fields": list(update_data.keys())},
)
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/validate — validate (leads + admin)
# ---------------------------------------------------------------------------
@router.post("/{test_id}/validate", response_model=TestOut)
def validate_test(
test_id: uuid.UUID,
payload: TestValidate,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Validate the red or blue side of a test (dual validation).
Red Lead approves/rejects the red side; Blue Lead approves/rejects the
blue side. When *both* sides are approved the test state moves to
``validated``. If either side is rejected the state moves to ``rejected``.
"""
test = (
db.query(Test)
.options(joinedload(Test.technique))
.filter(Test.id == test_id)
.first()
)
if test is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Test not found",
)
now = datetime.utcnow()
if current_user.role in ("red_lead", "admin"):
test.red_validation_status = payload.result.value
test.red_validated_by = current_user.id
test.red_validated_at = now
side = "red"
elif current_user.role == "blue_lead":
test.blue_validation_status = payload.result.value
test.blue_validated_by = current_user.id
test.blue_validated_at = now
side = "blue"
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="Not enough permissions to validate",
)
# Store the overall result from the payload
test.result = payload.result
# Determine aggregate state
red_ok = test.red_validation_status == "approved"
blue_ok = test.blue_validation_status == "approved"
red_rej = test.red_validation_status == "rejected"
blue_rej = test.blue_validation_status == "rejected"
if red_ok and blue_ok:
test.state = TestState.validated
elif red_rej or blue_rej:
test.state = TestState.rejected
else:
test.state = TestState.in_review
db.commit()
db.refresh(test)
# Recalculate the parent technique's global status
technique = test.technique
recalculate_technique_status(db, technique)
log_action(
db,
user_id=current_user.id,
action="validate_test",
entity_type="test",
entity_id=test.id,
details={
"side": side,
"result": payload.result.value,
"technique_id": str(test.technique_id),
},
)
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/reject — reject (leads + admin)
# ---------------------------------------------------------------------------
@router.post("/{test_id}/reject", response_model=TestOut)
def reject_test(
test_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
):
"""Reject a test, setting its state to *rejected*."""
test = db.query(Test).filter(Test.id == test_id).first()
if test is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="Test not found",
)
test.state = TestState.rejected
db.commit()
db.refresh(test)
log_action(
db,
user_id=current_user.id,
action="reject_test",
entity_type="test",
entity_id=test.id,
details={"technique_id": str(test.technique_id)},
)
return test