Files
Aegis/backend/app/services/test_template_service.py
T
kitos 82ac9c7014
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,campaigns): validate MITRE technique IDs, surface template suggestions first, manager can re-approve rejected campaigns
- TestTemplate/TemplateSuggestion creation and updates now reject any
  mitre_technique_id that doesn't match a real, already-synced MITRE
  ATT&CK technique. Frontend swaps the free-text ID input for a
  technique picker.
- Test Catalog now shows pending template suggestions before the
  catalog grid, with full field detail (platform, severity, tool,
  atomic ID, source URL, remediation) instead of just name/procedure.
- A manager can edit and directly re-approve a campaign they previously
  rejected, instead of needing the original lead to resubmit it.
2026-07-16 14:29:01 +02:00

297 lines
10 KiB
Python

"""Test template service — framework-agnostic CRUD and queries."""
# Enable future language features for compatibility
from __future__ import annotations
# Import uuid
import uuid
# Import func, or_ from sqlalchemy
from sqlalchemy import func, or_
# Import Session from sqlalchemy.orm
from sqlalchemy.orm import Session
# Import EntityNotFoundError from app.domain.errors
from app.domain.errors import BusinessRuleViolation, EntityNotFoundError
# Import TestTemplate from app.models.test_template
from app.models.test_template import TestTemplate
from app.models.technique import Technique
from app.models.test import Test
# Import escape_like from app.utils
from app.utils import escape_like
def _build_template_query(
db: Session,
*,
source: str | None = None,
platform: str | None = None,
severity: str | None = None,
mitre_technique_id: str | None = None,
search: str | None = None,
is_active: bool | None = None,
):
"""Build the filtered TestTemplate query shared by list_templates() and
count_templates() — a single source of truth so a paginated page of
results and the total count it's paginated against can never drift
apart (same pattern used for tests in test_crud_service.py)."""
query = db.query(TestTemplate)
if is_active is not None:
query = query.filter(TestTemplate.is_active == is_active)
if source:
query = query.filter(TestTemplate.source == source)
if platform:
query = query.filter(TestTemplate.platform.ilike(f"%{escape_like(platform)}%"))
if severity:
query = query.filter(TestTemplate.severity == severity)
if mitre_technique_id:
query = query.filter(TestTemplate.mitre_technique_id == mitre_technique_id)
if search:
pattern = f"%{escape_like(search)}%"
query = query.filter(
or_(
TestTemplate.name.ilike(pattern),
TestTemplate.description.ilike(pattern),
)
)
return query
# Define function list_templates
def list_templates(
# Entry: db
db: Session,
*,
# Entry: source
source: str | None = None,
# Entry: platform
platform: str | None = None,
# Entry: severity
severity: str | None = None,
# Entry: mitre_technique_id
mitre_technique_id: str | None = None,
# Entry: search
search: str | None = None,
# Entry: is_active
is_active: bool | None = None,
# Entry: offset
offset: int = 0,
# Entry: limit
limit: int = 50,
) -> list:
"""Return paginated, filterable list of test templates."""
query = _build_template_query(
db, source=source, platform=platform, severity=severity,
mitre_technique_id=mitre_technique_id, search=search, is_active=is_active,
)
templates = (
query
.order_by(TestTemplate.mitre_technique_id, TestTemplate.name)
.offset(offset)
.limit(limit)
.all()
)
# Attach existing_test_count per template — lets the catalog warn before
# creating a likely-duplicate test for a technique that already has one.
if templates:
mitre_ids = {t.mitre_technique_id for t in templates}
counts = dict(
db.query(Technique.mitre_id, func.count(Test.id))
.join(Test, Test.technique_id == Technique.id)
.filter(Technique.mitre_id.in_(mitre_ids))
.group_by(Technique.mitre_id)
.all()
)
for t in templates:
t.existing_test_count = counts.get(t.mitre_technique_id, 0)
return templates
def count_templates(
db: Session,
*,
source: str | None = None,
platform: str | None = None,
severity: str | None = None,
mitre_technique_id: str | None = None,
search: str | None = None,
is_active: bool | None = None,
) -> int:
"""Return the total count of templates matching the same filters as
list_templates() — lets the catalog UI show a true page count even
though the list endpoint caps how many rows it returns per request."""
query = _build_template_query(
db, source=source, platform=platform, severity=severity,
mitre_technique_id=mitre_technique_id, search=search, is_active=is_active,
)
return query.count()
# Define function get_template_stats
def get_template_stats(db: Session) -> dict:
"""Return catalog statistics: totals by source, platform, active/inactive."""
# Assign total = db.query(func.count(TestTemplate.id)).scalar() or 0
total = db.query(func.count(TestTemplate.id)).scalar() or 0
# Assign active = (
active = (
db.query(func.count(TestTemplate.id))
# Chain .filter() call
.filter(TestTemplate.is_active == True) # noqa: E712
# Chain .scalar() call
.scalar()
) or 0
# Assign inactive = total - active
inactive = total - active
# Assign source_rows = (
source_rows = (
db.query(TestTemplate.source, func.count(TestTemplate.id))
# Chain .filter() call
.filter(TestTemplate.is_active == True) # noqa: E712
# Chain .group_by() call
.group_by(TestTemplate.source)
# Chain .all() call
.all()
)
# Assign by_source = {source: cnt for source, cnt in source_rows}
by_source = {source: cnt for source, cnt in source_rows}
# Assign platform_rows = (
platform_rows = (
db.query(TestTemplate.platform, func.count(TestTemplate.id))
# Chain .filter() call
.filter(TestTemplate.is_active == True) # noqa: E712
# Chain .group_by() call
.group_by(TestTemplate.platform)
# Chain .all() call
.all()
)
# Assign by_platform = {(platform or "unspecified"): cnt for platform, cnt in platform_rows}
by_platform = {(platform or "unspecified"): cnt for platform, cnt in platform_rows}
# Return {
return {
# Literal argument value
"total": total,
# Literal argument value
"active": active,
# Literal argument value
"inactive": inactive,
# Literal argument value
"by_source": by_source,
# Literal argument value
"by_platform": by_platform,
}
# Define function bulk_activate
def bulk_activate(db: Session, *, activate: bool) -> int:
"""Set all templates to active or inactive. Returns count of affected. Does NOT commit."""
# Assign count = (
count = (
db.query(TestTemplate)
# Chain .filter() call
.filter(TestTemplate.is_active != activate)
# Chain .update() call
.update({TestTemplate.is_active: activate})
)
# Return count
return count
# Define function get_templates_by_technique
def get_templates_by_technique(db: Session, mitre_id: str) -> list:
"""Return all active templates mapped to a specific MITRE technique."""
# Return (
return (
db.query(TestTemplate)
# Chain .filter() call
.filter(
TestTemplate.mitre_technique_id == mitre_id,
TestTemplate.is_active == True, # noqa: E712
)
# Chain .order_by() call
.order_by(TestTemplate.name)
# Chain .all() call
.all()
)
# Define function get_template_or_raise
def get_template_or_raise(db: Session, template_id: uuid.UUID) -> TestTemplate:
"""Return a template by ID. Raises EntityNotFoundError if not found."""
# Assign template = db.query(TestTemplate).filter(TestTemplate.id == template_id).first()
template = db.query(TestTemplate).filter(TestTemplate.id == template_id).first()
# Check: template is None
if template is None:
# Raise EntityNotFoundError
raise EntityNotFoundError("Test template", str(template_id))
# Return template
return template
def validate_mitre_technique_id(db: Session, mitre_technique_id: str) -> None:
"""Raise BusinessRuleViolation unless *mitre_technique_id* matches a real,
already-synced MITRE ATT&CK technique — free text like "made up" or a
typo'd ID would otherwise silently create an orphaned template that can
never be found via technique lookups or coverage reporting."""
exists = db.query(Technique).filter(Technique.mitre_id == mitre_technique_id).first()
if not exists:
raise BusinessRuleViolation(
f"'{mitre_technique_id}' is not a known MITRE ATT&CK technique ID"
)
# Define function create_template
def create_template(db: Session, **fields: object) -> TestTemplate:
"""Create a test template from keyword args (e.g. payload.model_dump()). Does NOT commit."""
validate_mitre_technique_id(db, fields["mitre_technique_id"])
# Assign template = TestTemplate(**fields)
template = TestTemplate(**fields)
# Stage new record(s) for database insertion
db.add(template)
# Return template
return template
# Define function update_template
def update_template(db: Session, template_id: uuid.UUID, **fields: object) -> TestTemplate:
"""Update an existing template. Raises EntityNotFoundError if not found. Does NOT commit."""
if fields.get("mitre_technique_id"):
validate_mitre_technique_id(db, fields["mitre_technique_id"])
# Assign template = get_template_or_raise(db, template_id)
template = get_template_or_raise(db, template_id)
# Iterate over fields.items()
for field, value in fields.items():
# Check: hasattr(template, field)
if hasattr(template, field):
# Call setattr()
setattr(template, field, value)
# Return template
return template
# Define function toggle_template_active
def toggle_template_active(db: Session, template_id: uuid.UUID) -> TestTemplate:
"""Toggle template active/inactive. Does NOT commit."""
# Assign template = get_template_or_raise(db, template_id)
template = get_template_or_raise(db, template_id)
# Assign template.is_active = not template.is_active
template.is_active = not template.is_active
# Return template
return template
# Define function soft_delete_template
def soft_delete_template(db: Session, template_id: uuid.UUID) -> None:
"""Soft-delete a template by setting is_active=False. Does NOT commit."""
# Assign template = get_template_or_raise(db, template_id)
template = get_template_or_raise(db, template_id)
# Assign template.is_active = False
template.is_active = False