Files
Aegis/backend/app/routers/tests.py
T

1589 lines
56 KiB
Python

"""CRUD router for security Tests — v2 with Red/Blue workflow.
Endpoints
---------
GET /tests — list with filters (state, technique_id)
POST /tests — create (red_tech, admin)
POST /tests/from-template — create from TestTemplate (red_tech, admin)
GET /tests/{id} — detail with split red/blue evidences
PATCH /tests/{id} — general update (draft/rejected only)
PATCH /tests/{id}/red — Red Team updates (draft, red_executing)
PATCH /tests/{id}/blue — Blue Team updates (blue_evaluating)
POST /tests/{id}/start-execution — draft → red_executing
POST /tests/{id}/submit-red — red_executing → blue_evaluating
POST /tests/{id}/start-blue-work — blue tech picks up (sets Tempo timer)
POST /tests/{id}/submit-blue — blue_evaluating → in_review
POST /tests/{id}/validate-red — Red Lead validates
POST /tests/{id}/validate-blue — Blue Lead validates
POST /tests/{id}/reopen — rejected → draft
GET /tests/{id}/timeline — audit-log history for this test
"""
import base64
import hashlib
import os
import uuid
from datetime import datetime
from typing import Any, Optional
# Import APIRouter, Depends, HTTPException, Query, Reque... from fastapi
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import BaseModel
from sqlalchemy.orm import Session
# Import get_db from app.database
from app.database import get_db
# Import get_current_user, require_any_role, require_role from app.dependencies.auth
from app.dependencies.auth import get_current_user, require_any_role, require_role
# Import UnitOfWork from app.domain.unit_of_work
from app.domain.unit_of_work import UnitOfWork
# Import limiter from app.limiter
from app.limiter import limiter
from app.models.enums import TestState, TestResult, TeamSide
from app.models.evidence import Evidence
from app.storage import upload_file
from app.models.technique import Technique
from app.models.test import Test
from app.models.user import User
# Import from app.schemas.test
from app.schemas.test import (
TestBlueUpdate,
TestBlueValidate,
TestClassificationUpdate,
TestCreate,
TestOut,
TestRedUpdate,
TestRedValidate,
TestRemediationUpdate,
TestUpdate,
)
# Import TestTemplateInstantiate from app.schemas.test_template
from app.schemas.test_template import TestTemplateInstantiate
# Import log_action from app.services.audit_service
from app.services.audit_service import log_action
# Import recalculate_technique_status from app.services.status_service
from app.services.status_service import recalculate_technique_status
from app.services.webhook_service import dispatch_webhook
from app.services.test_crud_service import (
create_test as crud_create_test,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
create_test_from_template as crud_create_from_template,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
get_test_detail as crud_get_test_detail,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
get_test_or_raise as crud_get_test_or_raise,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
get_test_timeline as crud_get_test_timeline,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
get_test_with_technique as crud_get_test_with_technique,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
list_tests as crud_list_tests,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
update_test as crud_update_test,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
update_test_blue as crud_update_test_blue,
)
# Import from app.services.test_crud_service
from app.services.test_crud_service import (
update_test_red as crud_update_test_red,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
start_execution as wf_start_execution,
submit_red_evidence as wf_submit_red,
submit_blue_evidence as wf_submit_blue,
start_blue_work as wf_start_blue_work,
validate_as_red_lead as wf_validate_red,
validate_as_blue_lead as wf_validate_blue,
reopen_test as wf_reopen,
handle_remediation_completed as wf_handle_remediation,
get_retest_chain as wf_get_retest_chain,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
handle_remediation_completed as wf_handle_remediation,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
pause_timer as wf_pause_timer,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
reopen_test as wf_reopen,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
resume_timer as wf_resume_timer,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
start_execution as wf_start_execution,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
submit_blue_evidence as wf_submit_blue,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
submit_red_evidence as wf_submit_red,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
validate_as_blue_lead as wf_validate_blue,
)
# Import from app.services.test_workflow_service
from app.services.test_workflow_service import (
validate_as_red_lead as wf_validate_red,
)
# Assign router = APIRouter(prefix="/tests", tags=["tests"])
router = APIRouter(prefix="/tests", tags=["tests"])
# ---------------------------------------------------------------------------
# GET /tests — list with filters
# ---------------------------------------------------------------------------
@router.get("", response_model=list[TestOut])
# Define function list_tests
def list_tests(
# Entry: state
state: Optional[str] = Query(None, description="Filter by test state"),
# Entry: technique_id
technique_id: Optional[uuid.UUID] = Query(None, description="Filter by technique"),
# Entry: platform
platform: Optional[str] = Query(None, description="Filter by platform"),
# Entry: created_by
created_by: Optional[uuid.UUID] = Query(None, description="Filter by creator"),
# Entry: pending_validation_side
pending_validation_side: Optional[str] = Query(
None, description="Filter in_review tests pending validation on 'red' or 'blue' side"
),
not_in_any_campaign: bool = Query(
False, description="Only return tests not linked to any campaign"
),
offset: int = Query(0, ge=0),
# Entry: limit
limit: int = Query(50, ge=1, le=200),
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> list:
"""Return a paginated list of tests, optionally filtered by state, technique, platform or creator.
Args:
state (Optional[str]): Filter by test state (e.g. ``draft``, ``validated``).
technique_id (Optional[uuid.UUID]): Filter tests belonging to a specific technique.
platform (Optional[str]): Filter by target platform (e.g. ``windows``, ``linux``).
created_by (Optional[uuid.UUID]): Filter by the UUID of the creator.
pending_validation_side (Optional[str]): Filter ``in_review`` tests pending validation
on ``'red'`` or ``'blue'`` side.
offset (int): Number of records to skip for pagination.
limit (int): Maximum number of records to return.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Serialised list of :class:`TestOut` objects matching the filters.
"""
# Return crud_list_tests(
return crud_list_tests(
db,
# Keyword argument: state
state=state,
# Keyword argument: technique_id
technique_id=technique_id,
# Keyword argument: platform
platform=platform,
# Keyword argument: created_by
created_by=created_by,
# Keyword argument: pending_validation_side
pending_validation_side=pending_validation_side,
not_in_any_campaign=not_in_any_campaign,
offset=offset,
# Keyword argument: limit
limit=limit,
)
# ---------------------------------------------------------------------------
# POST /tests — create (red_tech or admin)
# ---------------------------------------------------------------------------
@router.post(
# Literal argument value
"",
# Keyword argument: response_model
response_model=TestOut,
# Keyword argument: status_code
status_code=status.HTTP_201_CREATED,
)
# Apply the @limiter.limit decorator
@limiter.limit("30/minute")
# Define function create_test
def create_test(
# Entry: request
request: Request,
# Entry: payload
payload: TestCreate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> TestOut:
"""Create a new test linked to an existing technique.
``created_by`` is set automatically and ``state`` defaults to *draft*.
Args:
request (Request): FastAPI request object (used by the rate limiter).
payload (TestCreate): Fields for the new test, including ``technique_id``.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead creating the test.
Returns:
TestOut: The newly created test with all fields populated.
"""
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_create_test(
test = crud_create_test(
db,
# Keyword argument: technique_id
technique_id=payload.technique_id,
# Keyword argument: creator_id
creator_id=current_user.id,
**payload.model_dump(exclude={"technique_id"}),
)
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="create_test",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"name": test.name, "technique_id": str(test.technique_id)},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Auto-create Jira ticket (non-fatal — any failure is logged, not raised)
try:
from app.services.jira_service import auto_create_test_issue
auto_create_test_issue(db, test, current_user)
db.commit()
except Exception:
pass # jira_service already logs warnings internally
return test
# ---------------------------------------------------------------------------
# POST /tests/from-template — create from TestTemplate
# ---------------------------------------------------------------------------
@router.post(
# Literal argument value
"/from-template",
# Keyword argument: response_model
response_model=TestOut,
# Keyword argument: status_code
status_code=status.HTTP_201_CREATED,
)
# Apply the @limiter.limit decorator
@limiter.limit("30/minute")
# Define function create_test_from_template
def create_test_from_template(
# Entry: request
request: Request,
# Entry: payload
payload: TestTemplateInstantiate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> TestOut:
"""Instantiate a real Test from an existing TestTemplate.
The template's fields are copied into the new test as starting data.
Args:
request (Request): FastAPI request object (used by the rate limiter).
payload (TestTemplateInstantiate): Contains ``template_id`` and target ``technique_id``.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead creating the test.
Returns:
TestOut: The newly created test populated from the template.
"""
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_create_from_template(
test = crud_create_from_template(
db,
# Keyword argument: template_id
template_id=payload.template_id,
# Keyword argument: technique_id_or_mitre
technique_id_or_mitre=payload.technique_id,
# Keyword argument: creator_id
creator_id=current_user.id,
name_override=payload.name,
description_override=payload.description,
platform_override=payload.platform,
procedure_text_override=payload.procedure_text,
tool_used_override=payload.tool_used,
)
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="create_test_from_template",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={
# Literal argument value
"name": test.name,
# Literal argument value
"template_id": str(payload.template_id),
# Literal argument value
"technique_id": str(test.technique_id),
},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Auto-create Jira ticket (non-fatal)
try:
from app.services.jira_service import auto_create_test_issue
auto_create_test_issue(db, test, current_user)
db.commit()
except Exception:
pass
return test
# ---------------------------------------------------------------------------
# GET /tests/{id} — detail with evidences split by team
# ---------------------------------------------------------------------------
@router.get("/{test_id}", response_model=TestOut)
# Define function get_test
def get_test(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> TestOut:
"""Return full details for a single test, including its evidences.
Args:
test_id (uuid.UUID): Primary key of the test to retrieve.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
TestOut: Full test detail including split red/blue evidence lists.
"""
# Return crud_get_test_detail(db, test_id)
return crud_get_test_detail(db, test_id)
# ---------------------------------------------------------------------------
# PATCH /tests/{id} — general update (draft / rejected)
# ---------------------------------------------------------------------------
@router.patch("/{test_id}", response_model=TestOut)
# Define function update_test
def update_test(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestUpdate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> TestOut:
"""Update one or more fields of an existing test.
Only leads or admins can update general test fields.
The test must be in ``draft`` or ``rejected`` state.
Args:
test_id (uuid.UUID): Primary key of the test to update.
payload (TestUpdate): Partial update payload; only set fields are applied.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead performing the update.
Returns:
TestOut: The updated test with refreshed field values.
"""
# Assign update_data = payload.model_dump(exclude_unset=True)
update_data = payload.model_dump(exclude_unset=True)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_update_test(
test = crud_update_test(
db,
test_id,
# Keyword argument: updater_id
updater_id=current_user.id,
# Keyword argument: updater_role
updater_role=current_user.role,
**update_data,
)
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_test",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"updated_fields": list(update_data.keys())},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# PATCH /tests/{id}/classification — admin data classification
# ---------------------------------------------------------------------------
@router.patch("/{test_id}/classification", response_model=TestOut)
# Define function update_test_classification
def update_test_classification(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestClassificationUpdate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_role("admin")),
) -> TestOut:
"""Update the data classification label for a test (admin only).
Args:
test_id (uuid.UUID): Primary key of the test to classify.
payload (TestClassificationUpdate): Contains the new ``data_classification`` value.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated admin user.
Returns:
TestOut: The test with the updated ``data_classification`` field.
"""
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Assign test.data_classification = payload.data_classification.value
test.data_classification = payload.data_classification.value
# Flush changes to DB without committing the transaction
db.flush()
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_test_classification",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"data_classification": payload.data_classification.value},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# PATCH /tests/{id}/red — Red Team update (draft, red_executing)
# ---------------------------------------------------------------------------
@router.patch("/{test_id}/red", response_model=TestOut)
# Define function update_test_red
def update_test_red(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestRedUpdate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_tech", "red_lead")),
) -> TestOut:
"""Red Team updates their fields (allowed in ``draft`` and ``red_executing``).
Args:
test_id (uuid.UUID): Primary key of the test to update.
payload (TestRedUpdate): Red-team-specific fields to update.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_tech or red_lead.
Returns:
TestOut: The updated test with refreshed red-team field values.
"""
# Assign update_data = payload.model_dump(exclude_unset=True)
update_data = payload.model_dump(exclude_unset=True)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_update_test_red(db, test_id, **update_data)
test = crud_update_test_red(db, test_id, **update_data)
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_test_red",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"updated_fields": list(update_data.keys())},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# PATCH /tests/{id}/blue — Blue Team update (blue_evaluating only)
# ---------------------------------------------------------------------------
@router.patch("/{test_id}/blue", response_model=TestOut)
# Define function update_test_blue
def update_test_blue(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestBlueUpdate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("blue_tech", "blue_lead")),
) -> TestOut:
"""Blue Team updates their fields (allowed only in ``blue_evaluating``).
Args:
test_id (uuid.UUID): Primary key of the test to update.
payload (TestBlueUpdate): Blue-team-specific fields to update.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated blue_tech or blue_lead.
Returns:
TestOut: The updated test with refreshed blue-team field values.
"""
# Assign update_data = payload.model_dump(exclude_unset=True)
update_data = payload.model_dump(exclude_unset=True)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = crud_update_test_blue(db, test_id, **update_data)
test = crud_update_test_blue(db, test_id, **update_data)
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_test_blue",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"updated_fields": list(update_data.keys())},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/start-execution — draft → red_executing
# ---------------------------------------------------------------------------
@router.post("/{test_id}/start-execution", response_model=TestOut)
# Define function start_execution
def start_execution(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_tech", "red_lead")),
) -> TestOut:
"""Move a test from ``draft`` to ``red_executing``.
Args:
test_id (uuid.UUID): Primary key of the test to start.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_tech or red_lead initiating execution.
Returns:
TestOut: The updated test in ``red_executing`` state.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_start_execution(db, test, current_user)
test = wf_start_execution(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/submit-red — red_executing → blue_evaluating
# ---------------------------------------------------------------------------
@router.post("/{test_id}/submit-red", response_model=TestOut)
# Define function submit_red
def submit_red(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_tech", "red_lead")),
) -> TestOut:
"""Red Team finalises — move from ``red_executing`` to ``blue_evaluating``.
Args:
test_id (uuid.UUID): Primary key of the test to submit.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_tech or red_lead submitting red evidence.
Returns:
TestOut: The updated test in ``blue_evaluating`` state.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_submit_red(db, test, current_user)
test = wf_submit_red(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/submit-blue — blue_evaluating → in_review
# ---------------------------------------------------------------------------
@router.post("/{test_id}/submit-blue", response_model=TestOut)
# Define function submit_blue
def submit_blue(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("blue_tech", "blue_lead")),
) -> TestOut:
"""Blue Team finalises — move from ``blue_evaluating`` to ``in_review``.
Args:
test_id (uuid.UUID): Primary key of the test to submit.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated blue_tech or blue_lead submitting blue evidence.
Returns:
TestOut: The updated test in ``in_review`` state.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_submit_blue(db, test, current_user)
test = wf_submit_blue(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/start-blue-work — blue tech picks up test for evaluation
# ---------------------------------------------------------------------------
@router.post("/{test_id}/start-blue-work", response_model=TestOut)
def start_blue_work(
test_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("blue_tech", "blue_lead")),
):
"""Blue tech picks up the test to start evaluating. Sets the Tempo timer start."""
test = crud_get_test_or_raise(db, test_id)
with UnitOfWork(db) as uow:
test = wf_start_blue_work(db, test, current_user)
uow.commit()
db.refresh(test)
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/pause-timer — pause the active phase timer
# ---------------------------------------------------------------------------
@router.post("/{test_id}/pause-timer", response_model=TestOut)
# Define function pause_timer
def pause_timer(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_tech", "blue_tech", "red_lead", "blue_lead")),
) -> TestOut:
"""Pause the running timer for the current phase (red_executing or blue_evaluating).
Args:
test_id (uuid.UUID): Primary key of the test whose timer should be paused.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated team member in the active phase.
Returns:
TestOut: The updated test with the phase timer paused.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_pause_timer(db, test, current_user)
test = wf_pause_timer(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/resume-timer — resume a paused phase timer
# ---------------------------------------------------------------------------
@router.post("/{test_id}/resume-timer", response_model=TestOut)
# Define function resume_timer
def resume_timer(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_tech", "blue_tech", "red_lead", "blue_lead")),
) -> TestOut:
"""Resume the paused timer for the current phase.
Args:
test_id (uuid.UUID): Primary key of the test whose timer should be resumed.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated team member in the active phase.
Returns:
TestOut: The updated test with the phase timer running again.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_resume_timer(db, test, current_user)
test = wf_resume_timer(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/validate-red — Red Lead validates
# ---------------------------------------------------------------------------
@router.post("/{test_id}/validate-red", response_model=TestOut)
# Define function validate_red
def validate_red(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestRedValidate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead")),
) -> TestOut:
"""Red Lead approves or rejects the red side of a test.
Args:
test_id (uuid.UUID): Primary key of the test to validate.
payload (TestRedValidate): Validation status and optional notes from the Red Lead.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead performing the validation.
Returns:
TestOut: The updated test reflecting the red validation decision.
"""
# Assign test = crud_get_test_with_technique(db, test_id)
test = crud_get_test_with_technique(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_validate_red(
test = wf_validate_red(
db, test, current_user,
# Keyword argument: validation_status
validation_status=payload.red_validation_status,
# Keyword argument: notes
notes=payload.red_validation_notes,
)
# Check: test.state in (TestState.validated, TestState.rejected)
if test.state in (TestState.validated, TestState.rejected):
# Call recalculate_technique_status()
recalculate_technique_status(db, test.technique)
# Flag technique for review — coverage changed
if test.technique:
test.technique.review_required = True
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
if test.state == TestState.validated:
dispatch_webhook("test.validated", {"test_id": str(test.id), "technique_id": str(test.technique_id), "result": test.result.value if test.result else None})
elif test.state == TestState.rejected:
dispatch_webhook("test.rejected", {"test_id": str(test.id), "technique_id": str(test.technique_id)})
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/validate-blue — Blue Lead validates
# ---------------------------------------------------------------------------
@router.post("/{test_id}/validate-blue", response_model=TestOut)
# Define function validate_blue
def validate_blue(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestBlueValidate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("blue_lead")),
) -> TestOut:
"""Blue Lead approves or rejects the blue side of a test.
Args:
test_id (uuid.UUID): Primary key of the test to validate.
payload (TestBlueValidate): Validation status and optional notes from the Blue Lead.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated blue_lead performing the validation.
Returns:
TestOut: The updated test reflecting the blue validation decision.
"""
# Assign test = crud_get_test_with_technique(db, test_id)
test = crud_get_test_with_technique(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_validate_blue(
test = wf_validate_blue(
db, test, current_user,
# Keyword argument: validation_status
validation_status=payload.blue_validation_status,
# Keyword argument: notes
notes=payload.blue_validation_notes,
)
# Check: test.state in (TestState.validated, TestState.rejected)
if test.state in (TestState.validated, TestState.rejected):
# Call recalculate_technique_status()
recalculate_technique_status(db, test.technique)
# Flag technique for review — coverage changed
if test.technique:
test.technique.review_required = True
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
if test.state == TestState.validated:
dispatch_webhook("test.validated", {"test_id": str(test.id), "technique_id": str(test.technique_id), "result": test.result.value if test.result else None})
elif test.state == TestState.rejected:
dispatch_webhook("test.rejected", {"test_id": str(test.id), "technique_id": str(test.technique_id)})
return test
# ---------------------------------------------------------------------------
# POST /tests/{id}/reopen — rejected → draft
# ---------------------------------------------------------------------------
@router.post("/{test_id}/reopen", response_model=TestOut)
# Define function reopen
def reopen(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> TestOut:
"""Reopen a rejected test, moving it back to ``draft``.
Args:
test_id (uuid.UUID): Primary key of the rejected test to reopen.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead reopening the test.
Returns:
TestOut: The updated test in ``draft`` state.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Open context manager
with UnitOfWork(db) as uow:
# Assign test = wf_reopen(db, test, current_user)
test = wf_reopen(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# PATCH /tests/{id}/remediation — update remediation fields
# ---------------------------------------------------------------------------
@router.patch("/{test_id}/remediation", response_model=TestOut)
# Define function update_remediation
def update_remediation(
# Entry: test_id
test_id: uuid.UUID,
# Entry: payload
payload: TestRemediationUpdate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role("red_lead", "blue_lead")),
) -> TestOut:
"""Update remediation fields on a test.
When ``remediation_status`` transitions to ``'completed'``, an automatic
re-test is created (subject to ``MAX_RETEST_COUNT``).
Args:
test_id (uuid.UUID): Primary key of the test to update.
payload (TestRemediationUpdate): Remediation fields to update (status, notes, etc.).
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead updating remediation.
Returns:
TestOut: The updated test with refreshed remediation fields.
"""
# Assign test = crud_get_test_or_raise(db, test_id)
test = crud_get_test_or_raise(db, test_id)
# Assign old_remediation_status = test.remediation_status
old_remediation_status = test.remediation_status
# Assign update_data = payload.model_dump(exclude_unset=True)
update_data = payload.model_dump(exclude_unset=True)
# Iterate over update_data.items()
for field, value in update_data.items():
# Call setattr()
setattr(test, field, value)
# Open context manager
with UnitOfWork(db) as uow:
# Call log_action()
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="update_remediation",
# Keyword argument: entity_type
entity_type="test",
# Keyword argument: entity_id
entity_id=test.id,
# Keyword argument: details
details={"updated_fields": list(update_data.keys())},
)
# Assign new_status = update_data.get("remediation_status")
new_status = update_data.get("remediation_status")
# Check: new_status == "completed" and old_remediation_status != "completed"
if new_status == "completed" and old_remediation_status != "completed":
# Call wf_handle_remediation()
wf_handle_remediation(db, test, current_user)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(test)
# Return test
return test
# ---------------------------------------------------------------------------
# GET /tests/{id}/timeline — audit history for this test
# ---------------------------------------------------------------------------
@router.get("/{test_id}/timeline")
# Define function get_test_timeline
def get_test_timeline(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> list:
"""Return the chronological audit-log history for a test.
Args:
test_id (uuid.UUID): Primary key of the test whose timeline is requested.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Chronological list of audit-log entries for the test.
"""
# Return crud_get_test_timeline(db, test_id)
return crud_get_test_timeline(db, test_id)
# ---------------------------------------------------------------------------
# GET /tests/{id}/retest-chain — full retest chain
# ---------------------------------------------------------------------------
@router.get("/{test_id}/retest-chain")
# Define function get_retest_chain
def get_retest_chain(
# Entry: test_id
test_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> list:
"""Return the full chain of retests (original + all retests) for a test.
Args:
test_id (uuid.UUID): Primary key of any test in the retest chain.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Ordered list of dicts describing each test in the chain,
including state, result, remediation status, and retest metadata.
"""
# Assign chain = wf_get_retest_chain(db, test_id)
chain = wf_get_retest_chain(db, test_id)
# Check: not chain
if not chain:
# Raise HTTPException
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Test not found")
# Return [
return [
{
# Literal argument value
"id": str(t.id),
# Literal argument value
"name": t.name,
# Literal argument value
"state": t.state.value if t.state else None,
# Literal argument value
"retest_of": str(t.retest_of) if t.retest_of else None,
# Literal argument value
"retest_count": t.retest_count,
# Literal argument value
"result": t.result.value if t.result else None,
# Literal argument value
"detection_result": t.detection_result.value if t.detection_result else None,
# Literal argument value
"remediation_status": t.remediation_status,
# Literal argument value
"created_at": t.created_at.isoformat() if t.created_at else None,
}
for t in chain
]
# ---------------------------------------------------------------------------
# POST /tests/{id}/sync-tempo — manual Tempo sync for red execution worklog
# ---------------------------------------------------------------------------
@router.post("/{test_id}/sync-tempo")
def sync_tempo(
test_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
):
"""Manually sync this test's red team execution worklog(s) to Tempo.
Useful when the automatic sync failed at phase completion (e.g. Tempo
was not yet configured). Only red_team_execution worklogs are eligible.
Already-synced worklogs are skipped. Returns a summary of what happened.
"""
from datetime import datetime as _dt
from app.models.worklog import Worklog
from app.services.tempo_service import auto_log_test_worklog
from app.services.test_crud_service import get_test_or_raise as _get
test = _get(db, test_id)
worklogs = (
db.query(Worklog)
.filter(
Worklog.entity_type == "test",
Worklog.entity_id == test_id,
Worklog.activity_type == "red_team_execution",
)
.all()
)
if not worklogs:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="No red team execution worklog found for this test.",
)
results = []
for wl in worklogs:
if wl.tempo_synced:
results.append({"worklog_id": str(wl.id), "status": "already_synced"})
continue
try:
result = auto_log_test_worklog(
db=db,
test=test,
user=current_user,
activity_type=wl.activity_type,
duration_seconds=wl.duration_seconds,
)
if result and isinstance(result, dict):
wl.tempo_synced = _dt.utcnow()
wl.tempo_worklog_id = str(result.get("tempoWorklogId", ""))
db.commit()
results.append({"worklog_id": str(wl.id), "status": "synced"})
else:
results.append({
"worklog_id": str(wl.id),
"status": "skipped",
"detail": "Tempo not configured or conditions not met.",
})
except Exception as exc:
results.append({
"worklog_id": str(wl.id),
"status": "error",
"detail": str(exc),
})
return {"results": results}
# ---------------------------------------------------------------------------
# POST /tests/{id}/request-discussion — disputed: confirm vote + notify other lead
# ---------------------------------------------------------------------------
@router.post("/{test_id}/request-discussion")
def request_discussion(
test_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead", "blue_lead", "admin")),
):
"""Called when the approving lead confirms their vote in a disputed test.
Sends a notification to the other lead (who rejected) asking them to
discuss and resolve the conflict. The test remains in 'disputed' state.
"""
from app.models.enums import TestState as ModelTestState
from app.models.user import User as UserModel
from app.services.notification_service import create_notification
test = crud_get_test_or_raise(db, test_id)
if test.state.value != "disputed":
from app.domain.errors import BusinessRuleViolation
raise BusinessRuleViolation("Test is not in disputed state")
role = current_user.role
# Identify who the "other lead" is (the one who rejected)
if (role in ("red_lead", "admin")) and test.red_validation_status == "approved":
# Red approved, Blue rejected → notify Blue Lead who rejected
rejector_id = test.blue_validated_by
rejector_label = "Blue Lead"
requester_label = "Red Lead"
elif (role in ("blue_lead", "admin")) and test.blue_validation_status == "approved":
# Blue approved, Red rejected → notify Red Lead who rejected
rejector_id = test.red_validated_by
rejector_label = "Red Lead"
requester_label = "Blue Lead"
else:
from app.domain.errors import BusinessRuleViolation
raise BusinessRuleViolation(
"The conflict state is inconsistent — no approving lead found"
)
# Look up the rejecting lead's full info for the response
rejector = (
db.query(UserModel).filter(UserModel.id == rejector_id).first()
if rejector_id else None
)
rejector_name = rejector.username if rejector else rejector_label
rejector_email = getattr(rejector, "email", None) if rejector else None
# Notify the rejecting lead
if rejector_id:
try:
create_notification(
db,
user_id=rejector_id,
type="validation_conflict",
title="Discussion requested on disputed test",
message=(
f"{requester_label} ({current_user.username}) is confirming their approval "
f"of test '{test.name}' and wants to discuss your rejection with you. "
f"Please reach out to resolve the disagreement."
),
entity_type="test",
entity_id=str(test.id),
)
except Exception as e:
import logging
logging.getLogger(__name__).warning(
"Failed to send discussion notification: %s", e
)
log_action(
db,
user_id=current_user.id,
action="request_dispute_discussion",
entity_type="test",
entity_id=test.id,
details={"test_name": test.name, "rejector": rejector_name},
)
db.commit()
return {
"status": "notification_sent",
"message": f"Discussion request sent to {rejector_name}",
"rejector_username": rejector_name,
"rejector_email": rejector_email,
"rejector_role": rejector_label,
}
# ---------------------------------------------------------------------------
# POST /tests/import-rt — bulk import from a real Red Team engagement
# ---------------------------------------------------------------------------
_ALLOWED_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}
_MAX_EVIDENCE_BYTES = 10 * 1024 * 1024 # 10 MB decoded per image
class RTEvidenceEntry(BaseModel):
filename: str # e.g. "screenshot_edr.png"
data: str # base64-encoded image content
caption: Optional[str] = None # optional description shown as evidence notes
class RTTechniqueEntry(BaseModel):
mitre_id: str
result: str # "detected" | "not_detected" | "partially_detected"
attack_success: bool = True
platform: Optional[str] = None
notes: Optional[str] = None
evidence: list[RTEvidenceEntry] # REQUIRED — at least one image per technique
class RTImportPayload(BaseModel):
name: str # engagement name, e.g. "Red Team Q1 2024"
date: Optional[str] = None # ISO date string
description: Optional[str] = None
operator: Optional[str] = None # team / company that ran the RT
techniques: list[RTTechniqueEntry]
@router.post("/import-rt", status_code=status.HTTP_201_CREATED)
def import_rt(
payload: RTImportPayload,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role("red_lead")),
):
"""Import results from a real Red Team engagement.
Creates one Test record per technique in ``validated`` state (bypassing
the normal Red/Blue workflow) and immediately recalculates coverage metrics.
Requires ``red_lead`` or ``admin`` role.
"""
# Pre-validate: every technique must include at least one evidence image
for entry in payload.techniques:
if not entry.evidence:
raise HTTPException(
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
detail=(
f"Technique {entry.mitre_id} is missing evidence. "
"At least one screenshot or image is required per technique."
),
)
# Execution date from payload or now
exec_date_str = payload.date or datetime.utcnow().date().isoformat()
# Result string → TestResult enum
_result_map = {
"detected": TestResult.detected,
"not_detected": TestResult.not_detected,
"partially_detected": TestResult.partially_detected,
}
created: list[dict[str, Any]] = []
skipped: list[dict[str, str]] = []
affected_technique_ids: set = set()
with UnitOfWork(db) as uow:
for entry in payload.techniques:
# Find technique
technique = (
db.query(Technique)
.filter(Technique.mitre_id == entry.mitre_id.upper())
.first()
)
if technique is None:
skipped.append({"mitre_id": entry.mitre_id, "reason": "Technique not found"})
continue
detection_result = _result_map.get(entry.result)
if detection_result is None:
skipped.append({"mitre_id": entry.mitre_id, "reason": f"Unknown result value '{entry.result}'"})
continue
test_name = f"[RT] {payload.name}{technique.name}"
# Build red_summary from notes + engagement metadata
parts = []
if payload.operator:
parts.append(f"Operator: {payload.operator}")
parts.append(f"Engagement date: {exec_date_str}")
if entry.notes:
parts.append(f"\n{entry.notes}")
red_summary_text = "\n".join(parts)
# RT pre-validates the Red side (they ran it), but Blue Lead
# must still validate the detection result before it counts.
# State = in_review so it appears in the Blue Lead's validation queue.
test = Test(
technique_id=technique.id,
name=test_name,
description=payload.description,
platform=entry.platform,
procedure_text=entry.notes,
created_by=current_user.id,
state=TestState.in_review,
# Red team — approved by the RT operator
attack_success=entry.attack_success,
red_summary=red_summary_text,
red_validation_status="approved",
red_validated_by=current_user.id,
red_validated_at=datetime.utcnow(),
# Blue team — pre-fill the detection result but leave
# validation_status pending so Blue Lead must confirm
detection_result=detection_result,
blue_validation_status=None,
# Timing
execution_date=exec_date_str,
created_at=datetime.utcnow(),
)
db.add(test)
db.flush()
# ── Store evidence images ──────────────────────────────
evidence_count = 0
for ev in entry.evidence:
safe_name = os.path.basename(ev.filename) or "evidence.png"
ext = os.path.splitext(safe_name)[1].lower()
if ext not in _ALLOWED_IMAGE_EXTS:
# Skip non-image files silently (log warning)
continue
try:
img_bytes = base64.b64decode(ev.data)
except Exception:
continue # malformed base64 — skip
if len(img_bytes) > _MAX_EVIDENCE_BYTES:
continue # over size limit — skip
sha256 = hashlib.sha256(img_bytes).hexdigest()
key = f"{test.id}/{uuid.uuid4()}_{safe_name}"
try:
upload_file(img_bytes, key)
except Exception:
continue # storage error — skip but don't abort
evidence_obj = Evidence(
test_id=test.id,
file_name=safe_name,
file_path=key,
sha256_hash=sha256,
uploaded_by=current_user.id,
uploaded_at=datetime.utcnow(),
team=TeamSide.red,
notes=ev.caption,
)
db.add(evidence_obj)
evidence_count += 1
affected_technique_ids.add(technique.id)
created.append({
"mitre_id": entry.mitre_id,
"test_name": test_name,
"result": entry.result,
"attack_success": entry.attack_success,
"evidence_attached": evidence_count,
})
log_action(
db,
user_id=current_user.id,
action="rt_import_test",
entity_type="test",
entity_id=test.id,
details={"engagement": payload.name, "mitre_id": entry.mitre_id},
)
# Recalculate coverage for all affected techniques
for tech_id in affected_technique_ids:
tech = db.query(Technique).filter(Technique.id == tech_id).first()
if tech:
recalculate_technique_status(db, tech)
uow.commit()
return {
"created": len(created),
"skipped": len(skipped),
"items": created,
"warnings": skipped,
"engagement": payload.name,
}