Files
Aegis/backend/app/routers/test_templates.py
T
kitos 0002601b50
Aegis CI / lint-and-test (push) Has been cancelled
Snyk Security Scan / Python vulnerabilities (backend) (push) Has been cancelled
Snyk Security Scan / npm vulnerabilities (frontend) (push) Has been cancelled
Snyk Security Scan / Docker image vulnerabilities (backend) (push) Has been cancelled
feat(templates): add pagination with total counts to catalog and campaign picker
Test Catalog's pagination showed only 'Page N' with no way to know
how many pages existed, and the campaign 'Add Test' modal's template
picker had no pagination at all — with 2800+ templates in the
catalog, its flat 100-row fetch made most templates unreachable
unless a search happened to match. Adds GET /test-templates/count
(open to any authenticated user, mirrors the list endpoint's filters)
and wires real 'Page N of M' + total-available counts into both.
2026-07-16 10:01:20 +02:00

561 lines
20 KiB
Python

"""CRUD router for TestTemplates — predefined test catalog.
Endpoints
---------
GET /test-templates — list with filters + pagination
GET /test-templates/stats — catalog statistics (admin)
GET /test-templates/{id} — detail
POST /test-templates — create custom (admin)
PATCH /test-templates/{id} — update (admin)
PATCH /test-templates/{id}/toggle-active — toggle active/inactive (admin)
DELETE /test-templates/{id} — soft delete (admin)
GET /test-templates/by-technique/{mitre_id} — templates for a MITRE technique
Filters (GET /test-templates)
-----------------------------
- source: atomic_red_team | mitre | custom
- platform: windows | linux | macos
- severity: low | medium | high | critical
- mitre_technique_id: filter by specific technique
- search: full-text search across name and description
- is_active: true | false (default only active)
- offset / limit: pagination (default limit=50)
"""
# Import uuid
import uuid
# Import Optional from typing
from typing import Optional
# Import APIRouter, Depends, Query, status from fastapi
from fastapi import APIRouter, Depends, Query, status
# Import Session from sqlalchemy.orm
from sqlalchemy.orm import Session
# Import get_db from app.database
from app.database import get_db
# Import get_current_user, require_any_role_strict from app.dependencies.auth
from app.dependencies.auth import get_current_user, require_any_role_strict
# Import UnitOfWork from app.domain.unit_of_work
from app.domain.unit_of_work import UnitOfWork
from app.models.technique import Technique
from app.models.user import User
# Import from app.schemas.test_template
from app.schemas.test_template import (
TestTemplateCreate,
TestTemplateOut,
TestTemplateSummary,
)
# Import log_action from app.services.audit_service
from app.services.audit_service import log_action
# Import from app.services.test_template_service
from app.services.test_template_service import (
bulk_activate,
count_templates,
get_template_or_raise,
get_template_stats,
list_templates,
soft_delete_template,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
create_template as create_template_svc,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
get_templates_by_technique as templates_by_technique,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
toggle_template_active as toggle_template_active_svc,
)
# Import from app.services.test_template_service
from app.services.test_template_service import (
update_template as update_template_svc,
)
# Assign router = APIRouter(prefix="/test-templates", tags=["test-templates"])
router = APIRouter(prefix="/test-templates", tags=["test-templates"])
# ---------------------------------------------------------------------------
# GET /test-templates — list with filters + pagination
# ---------------------------------------------------------------------------
@router.get("", response_model=list[TestTemplateSummary])
# Define function _list_templates_handler
def _list_templates_handler(
# Entry: source
source: Optional[str] = Query(None, description="Filter by source (atomic_red_team, mitre, custom)"),
# Entry: platform
platform: Optional[str] = Query(None, description="Filter by platform (windows, linux, macos)"),
# Entry: severity
severity: Optional[str] = Query(None, description="Filter by severity (low, medium, high, critical)"),
# Entry: mitre_technique_id
mitre_technique_id: Optional[str] = Query(None, description="Filter by MITRE technique ID"),
# Entry: search
search: Optional[str] = Query(None, description="Search in name and description"),
# Entry: is_active
is_active: Optional[bool] = Query(None, description="Filter by active status (true/false). Omit to return all."),
# Entry: offset
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, filterable list of test templates.
Args:
source (Optional[str]): Filter by source (``atomic_red_team``, ``mitre``, ``custom``).
platform (Optional[str]): Filter by platform (``windows``, ``linux``, ``macos``).
severity (Optional[str]): Filter by severity (``low``, ``medium``, ``high``, ``critical``).
mitre_technique_id (Optional[str]): Filter by MITRE technique ID string.
search (Optional[str]): Full-text search across name and description.
is_active (Optional[bool]): Filter by active status; omit to return all.
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:`TestTemplateSummary` objects.
"""
# Return list_templates(
return list_templates(
db,
# Keyword argument: source
source=source,
# Keyword argument: platform
platform=platform,
# Keyword argument: severity
severity=severity,
# Keyword argument: mitre_technique_id
mitre_technique_id=mitre_technique_id,
# Keyword argument: search
search=search,
# Keyword argument: is_active
is_active=is_active,
# Keyword argument: offset
offset=offset,
# Keyword argument: limit
limit=limit,
)
# ---------------------------------------------------------------------------
# GET /test-templates/count — total matching a filter, for pagination
# ---------------------------------------------------------------------------
@router.get("/count")
def _count_templates_handler(
source: Optional[str] = Query(None),
platform: Optional[str] = Query(None),
severity: Optional[str] = Query(None),
mitre_technique_id: Optional[str] = Query(None),
search: Optional[str] = Query(None),
is_active: Optional[bool] = Query(None),
db: Session = Depends(get_db),
current_user: User = Depends(get_current_user),
) -> dict:
"""Return the total count of templates matching the same filters as the
list endpoint, so the catalog UI can show a true page count."""
return {"total": count_templates(
db,
source=source,
platform=platform,
severity=severity,
mitre_technique_id=mitre_technique_id,
search=search,
is_active=is_active,
)}
# ---------------------------------------------------------------------------
# GET /test-templates/stats — catalog statistics (admin)
# ---------------------------------------------------------------------------
@router.get("/stats")
# Define function template_stats
def template_stats(
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> dict:
"""Return catalog statistics: active, by_source, by_platform.
Args:
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Counts of active templates broken down by source and platform.
"""
# Return get_template_stats(db)
return get_template_stats(db)
# ---------------------------------------------------------------------------
# PATCH /test-templates/bulk-activate — activate/deactivate all (admin)
# ---------------------------------------------------------------------------
@router.patch("/bulk-activate")
# Define function bulk_activate_templates
def bulk_activate_templates(
# Entry: activate
activate: bool = Query(True, description="True to activate all, False to deactivate all"),
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> dict:
"""Set all templates to active or inactive.
Args:
activate (bool): ``True`` to activate all templates, ``False`` to deactivate all.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Confirmation message with ``affected`` count and the applied ``is_active`` flag.
"""
# Assign count = bulk_activate(db, activate=activate)
count = bulk_activate(db, activate=activate)
# 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="bulk_activate_templates" if activate else "bulk_deactivate_templates",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=None,
# Keyword argument: details
details={"affected": count, "is_active": activate},
)
# Call uow.commit()
uow.commit()
# Return {
return {
# Literal argument value
"detail": f"{'Activated' if activate else 'Deactivated'} {count} templates",
# Literal argument value
"affected": count,
# Literal argument value
"is_active": activate,
}
# ---------------------------------------------------------------------------
# GET /test-templates/by-technique/{mitre_id}
# ---------------------------------------------------------------------------
@router.get("/by-technique/{mitre_id}", response_model=list[TestTemplateSummary])
# Define function _templates_by_technique_handler
def _templates_by_technique_handler(
# Entry: mitre_id
mitre_id: str,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> list:
"""Return all active templates mapped to a specific MITRE technique.
Args:
mitre_id (str): MITRE ATT&CK technique ID (e.g. ``T1059.001``).
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
list: Serialised list of :class:`TestTemplateSummary` objects for the technique.
"""
# Return templates_by_technique(db, mitre_id)
return templates_by_technique(db, mitre_id)
# ---------------------------------------------------------------------------
# GET /test-templates/{id} — detail
# ---------------------------------------------------------------------------
@router.get("/{template_id}", response_model=TestTemplateOut)
# Define function get_template
def get_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(get_current_user),
) -> TestTemplateOut:
"""Return full details for a single test template.
Args:
template_id (uuid.UUID): Primary key of the template to retrieve.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated user making the request.
Returns:
TestTemplateOut: Full template detail including all fields.
"""
# Return get_template_or_raise(db, template_id)
return get_template_or_raise(db, template_id)
# ---------------------------------------------------------------------------
# POST /test-templates — create (admin only)
# ---------------------------------------------------------------------------
@router.post(
# Literal argument value
"",
# Keyword argument: response_model
response_model=TestTemplateOut,
# Keyword argument: status_code
status_code=status.HTTP_201_CREATED,
)
# Define function create_template
def create_template(
# Entry: payload
payload: TestTemplateCreate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> TestTemplateOut:
"""Create a custom test template.
Args:
payload (TestTemplateCreate): All fields for the new template.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead creating the template.
Returns:
TestTemplateOut: The newly created template with all fields populated.
"""
# Assign template = create_template_svc(db, **payload.model_dump())
template = create_template_svc(db, **payload.model_dump())
# Open context manager
with UnitOfWork(db) as uow:
# Flag the associated technique for review — new template available
if template.mitre_technique_id:
technique = (
db.query(Technique)
.filter(Technique.mitre_id == template.mitre_technique_id)
.first()
)
if technique:
technique.review_required = True
log_action(
db,
# Keyword argument: user_id
user_id=current_user.id,
# Keyword argument: action
action="create_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={
# Literal argument value
"name": template.name,
# Literal argument value
"source": template.source,
# Literal argument value
"mitre_technique_id": template.mitre_technique_id,
},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
# ---------------------------------------------------------------------------
# PATCH /test-templates/{id} — update (admin only)
# ---------------------------------------------------------------------------
@router.patch("/{template_id}", response_model=TestTemplateOut)
# Define function update_template
def update_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: payload
payload: TestTemplateCreate,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> TestTemplateOut:
"""Update fields of an existing test template.
Args:
template_id (uuid.UUID): Primary key of the template to update.
payload (TestTemplateCreate): Fields to update; only set fields are applied.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead updating the template.
Returns:
TestTemplateOut: The updated template with refreshed field values.
"""
# Assign template = update_template_svc(db, template_id, **payload.model_dump(exclude_u...
template = update_template_svc(db, template_id, **payload.model_dump(exclude_unset=True))
# 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_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"updated_fields": list(payload.model_dump(exclude_unset=True).keys())},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
# ---------------------------------------------------------------------------
# PATCH /test-templates/{id}/toggle-active — toggle active/inactive (admin)
# ---------------------------------------------------------------------------
@router.patch("/{template_id}/toggle-active", response_model=TestTemplateOut)
# Define function toggle_template_active
def toggle_template_active(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> TestTemplateOut:
"""Toggle a template between active and inactive (is_active = not is_active).
Args:
template_id (uuid.UUID): Primary key of the template to toggle.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
TestTemplateOut: The template with the updated ``is_active`` flag.
"""
# Assign template = toggle_template_active_svc(db, template_id)
template = toggle_template_active_svc(db, template_id)
# 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="toggle_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"name": template.name, "is_active": template.is_active},
)
# Call uow.commit()
uow.commit()
# Reload ORM object attributes from the database
db.refresh(template)
# Return template
return template
# ---------------------------------------------------------------------------
# DELETE /test-templates/{id} — soft delete (admin only)
# ---------------------------------------------------------------------------
@router.delete("/{template_id}", status_code=status.HTTP_200_OK)
# Define function delete_template
def delete_template(
# Entry: template_id
template_id: uuid.UUID,
# Entry: db
db: Session = Depends(get_db),
# Entry: current_user
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead")),
) -> dict:
"""Soft-delete a test template by setting ``is_active=False``.
Args:
template_id (uuid.UUID): Primary key of the template to delete.
db (Session): SQLAlchemy database session.
current_user (User): Authenticated red_lead or blue_lead.
Returns:
dict: Confirmation message with key ``detail``.
"""
# Assign template = get_template_or_raise(db, template_id)
template = get_template_or_raise(db, template_id)
# Call soft_delete_template()
soft_delete_template(db, template_id)
# 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="delete_test_template",
# Keyword argument: entity_type
entity_type="test_template",
# Keyword argument: entity_id
entity_id=template.id,
# Keyword argument: details
details={"name": template.name},
)
# Call uow.commit()
uow.commit()
# Return {"detail": "Test template deactivated"}
return {"detail": "Test template deactivated"}