feat(tests): auto-propose procedure suggestions from submitted rounds, add lead review workflow

When Red or Blue submits a round for a test created from a template,
any extractable command in their procedure text is proposed as an
update to the template's suggested procedure — never written
automatically. Adds GET/approve/reject endpoints scoped so a lead only
reviews their own team's suggestions; approving writes the suggested
text into the template, rejecting leaves it untouched.
This commit is contained in:
kitos
2026-07-14 15:22:03 +02:00
parent 8bdbe48fbe
commit 7ec3c5bc8b
6 changed files with 547 additions and 0 deletions
+2
View File
@@ -43,6 +43,7 @@ from app.routers import techniques as techniques_router
from app.routers import tests as tests_router from app.routers import tests as tests_router
from app.routers import evidence as evidence_router from app.routers import evidence as evidence_router
from app.routers import test_templates as test_templates_router from app.routers import test_templates as test_templates_router
from app.routers import procedure_suggestions as procedure_suggestions_router
from app.routers import system as system_router from app.routers import system as system_router
from app.routers import metrics as metrics_router from app.routers import metrics as metrics_router
from app.routers import users as users_router from app.routers import users as users_router
@@ -213,6 +214,7 @@ app.include_router(tests_router.router, prefix="/api/v1")
app.include_router(evidence_router.router, prefix="/api/v1") app.include_router(evidence_router.router, prefix="/api/v1")
# Call app.include_router() # Call app.include_router()
app.include_router(test_templates_router.router, prefix="/api/v1") app.include_router(test_templates_router.router, prefix="/api/v1")
app.include_router(procedure_suggestions_router.router, prefix="/api/v1")
# Call app.include_router() # Call app.include_router()
app.include_router(system_router.router, prefix="/api/v1") app.include_router(system_router.router, prefix="/api/v1")
# Call app.include_router() # Call app.include_router()
@@ -0,0 +1,139 @@
"""Router for procedure-improvement suggestion review.
Endpoints
---------
GET /procedure-suggestions — list pending suggestions (lead/admin)
POST /procedure-suggestions/{id}/approve — write suggestion into template (lead/admin)
POST /procedure-suggestions/{id}/reject — dismiss suggestion (lead/admin)
A red_lead only ever sees/acts on "red" team suggestions, a blue_lead only
"blue" — admins can see and act on both.
"""
import uuid
from fastapi import APIRouter, Depends, HTTPException
from app.database import get_db
from app.dependencies.auth import require_any_role_strict
from app.domain.errors import EntityNotFoundError, InvalidOperationError
from app.domain.unit_of_work import UnitOfWork
from app.models.procedure_suggestion import ProcedureSuggestion
from app.models.test_template import TestTemplate
from app.models.user import User
from app.schemas.procedure_suggestion import ProcedureSuggestionOut
from app.services.audit_service import log_action
from app.services.procedure_suggestion_service import (
approve_suggestion as approve_suggestion_svc,
)
from app.services.procedure_suggestion_service import (
list_pending_suggestions,
)
from app.services.procedure_suggestion_service import (
reject_suggestion as reject_suggestion_svc,
)
from sqlalchemy.orm import Session
router = APIRouter(prefix="/procedure-suggestions", tags=["procedure-suggestions"])
def _team_for_role(user: User) -> str | None:
"""Return the single team a non-admin lead is scoped to, or None for admin (both)."""
if user.role == "red_lead":
return "red"
if user.role == "blue_lead":
return "blue"
return None
def _to_out(suggestion, template_by_id: dict) -> ProcedureSuggestionOut:
out = ProcedureSuggestionOut.model_validate(suggestion)
template = template_by_id.get(suggestion.template_id)
if template is not None:
out.template_name = template.name
out.template_current_text = (
template.attack_procedure if suggestion.team == "red" else template.detect_suggested_procedure
)
return out
@router.get("", response_model=list[ProcedureSuggestionOut])
def list_suggestions(
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead", "admin")),
) -> list[ProcedureSuggestionOut]:
"""List pending procedure suggestions, scoped to the caller's team (admins see both)."""
team = _team_for_role(current_user)
suggestions = list_pending_suggestions(db, team=team)
template_ids = {s.template_id for s in suggestions}
templates = (
db.query(TestTemplate).filter(TestTemplate.id.in_(template_ids)).all() if template_ids else []
)
template_by_id = {t.id: t for t in templates}
return [_to_out(s, template_by_id) for s in suggestions]
@router.post("/{suggestion_id}/approve", response_model=ProcedureSuggestionOut)
def approve_suggestion(
suggestion_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead", "admin")),
) -> ProcedureSuggestionOut:
"""Approve a suggestion, writing it into its template's suggested-procedure field."""
_check_team_permission(current_user, _team_or_404(db, suggestion_id))
try:
with UnitOfWork(db) as uow:
suggestion = approve_suggestion_svc(db, suggestion_id, current_user)
log_action(
db, user_id=current_user.id, action="approve_procedure_suggestion",
entity_type="procedure_suggestion", entity_id=suggestion.id,
details={"template_id": str(suggestion.template_id), "team": suggestion.team},
)
uow.commit()
except EntityNotFoundError as e:
raise HTTPException(status_code=404, detail=str(e)) from e
except InvalidOperationError as e:
raise HTTPException(status_code=400, detail=str(e)) from e
db.refresh(suggestion)
template = db.query(TestTemplate).filter(TestTemplate.id == suggestion.template_id).first()
return _to_out(suggestion, {template.id: template} if template else {})
@router.post("/{suggestion_id}/reject", response_model=ProcedureSuggestionOut)
def reject_suggestion(
suggestion_id: uuid.UUID,
db: Session = Depends(get_db),
current_user: User = Depends(require_any_role_strict("red_lead", "blue_lead", "admin")),
) -> ProcedureSuggestionOut:
"""Reject a suggestion without touching the template."""
_check_team_permission(current_user, _team_or_404(db, suggestion_id))
try:
with UnitOfWork(db) as uow:
suggestion = reject_suggestion_svc(db, suggestion_id, current_user)
log_action(
db, user_id=current_user.id, action="reject_procedure_suggestion",
entity_type="procedure_suggestion", entity_id=suggestion.id,
details={"template_id": str(suggestion.template_id), "team": suggestion.team},
)
uow.commit()
except EntityNotFoundError as e:
raise HTTPException(status_code=404, detail=str(e)) from e
except InvalidOperationError as e:
raise HTTPException(status_code=400, detail=str(e)) from e
db.refresh(suggestion)
return _to_out(suggestion, {})
def _team_or_404(db: Session, suggestion_id: uuid.UUID) -> str:
"""Look up just the team of a suggestion, for a permission check before mutating it."""
suggestion = db.query(ProcedureSuggestion).filter(ProcedureSuggestion.id == suggestion_id).first()
if suggestion is None:
raise HTTPException(status_code=404, detail="Procedure suggestion not found")
return suggestion.team
def _check_team_permission(user: User, team: str) -> None:
"""Raise 403 if a non-admin lead tries to act on the other team's suggestion."""
scoped_team = _team_for_role(user)
if scoped_team is not None and scoped_team != team:
raise HTTPException(status_code=403, detail=f"Not authorized to review {team} team suggestions")
@@ -0,0 +1,28 @@
"""Pydantic schemas for procedure-suggestion review endpoints."""
import uuid
from datetime import datetime
from pydantic import BaseModel, ConfigDict
class ProcedureSuggestionOut(BaseModel):
"""A proposed template procedure update, pending or reviewed."""
id: uuid.UUID
template_id: uuid.UUID
team: str
suggested_text: str
source_test_id: uuid.UUID | None = None
submitted_by: uuid.UUID | None = None
status: str
reviewed_by: uuid.UUID | None = None
reviewed_at: datetime | None = None
created_at: datetime | None = None
# Populated for display — the template name and its current suggested
# value, so a lead can compare without a second lookup.
template_name: str | None = None
template_current_text: str | None = None
model_config = ConfigDict(from_attributes=True)
@@ -0,0 +1,134 @@
"""Procedure-improvement suggestion workflow.
When an operator submits a round with a filled-in procedure field, the
command(s) in it are extracted heuristically (see
``procedure_extraction_service``) and proposed as an update to the
originating template's suggested-procedure field. A suggestion is only
ever created when there's a real, novel improvement to propose — nothing
is written to the template until a lead reviews and approves it.
"""
from __future__ import annotations
import uuid
from datetime import datetime
from sqlalchemy.orm import Session
from app.domain.errors import EntityNotFoundError, InvalidOperationError
from app.models.procedure_suggestion import ProcedureSuggestion
from app.models.test import Test
from app.models.test_template import TestTemplate
from app.models.user import User
from app.services.notification_service import notify_role
from app.services.procedure_extraction_service import extract_commands
# Which Test field holds the operator's free-text procedure, and which
# TestTemplate field it's proposed as an improvement to, per team.
_SOURCE_FIELD = {"red": "procedure_text", "blue": "detect_procedure"}
_TEMPLATE_FIELD = {"red": "attack_procedure", "blue": "detect_suggested_procedure"}
_LEAD_ROLE = {"red": "red_lead", "blue": "blue_lead"}
def create_suggestion_from_submission(db: Session, test: Test, team: str) -> ProcedureSuggestion | None:
"""Propose a template procedure improvement from *test*'s just-submitted round.
Returns the created (unflushed-to-caller-commit) suggestion, or ``None``
when there's nothing worth proposing: the test wasn't created from a
template, no command could be extracted, the extraction matches what
the template already suggests, or an identical suggestion is already
pending review. Does not commit; caller uses UnitOfWork.
"""
if not test.source_template_id:
return None
procedure_text = getattr(test, _SOURCE_FIELD[team])
extracted = extract_commands(procedure_text)
if not extracted:
return None
template = db.query(TestTemplate).filter(TestTemplate.id == test.source_template_id).first()
if template is None:
return None
current = (getattr(template, _TEMPLATE_FIELD[team]) or "").strip()
if extracted == current:
return None
duplicate = (
db.query(ProcedureSuggestion)
.filter(
ProcedureSuggestion.template_id == template.id,
ProcedureSuggestion.team == team,
ProcedureSuggestion.suggested_text == extracted,
ProcedureSuggestion.status == "pending",
)
.first()
)
if duplicate is not None:
return None
submitter_id = test.red_tech_assignee if team == "red" else test.blue_tech_assignee
suggestion = ProcedureSuggestion(
template_id=template.id,
team=team,
suggested_text=extracted,
source_test_id=test.id,
submitted_by=submitter_id,
)
db.add(suggestion)
db.flush()
notify_role(
db,
role=_LEAD_ROLE[team],
type="procedure_suggestion",
title="New procedure suggestion for review",
message=f'A {team} procedure improvement was proposed for template "{template.name}".',
entity_type="procedure_suggestion",
entity_id=suggestion.id,
)
return suggestion
def list_pending_suggestions(db: Session, *, team: str | None = None) -> list[ProcedureSuggestion]:
"""List pending suggestions, optionally filtered to one team."""
query = db.query(ProcedureSuggestion).filter(ProcedureSuggestion.status == "pending")
if team is not None:
query = query.filter(ProcedureSuggestion.team == team)
return query.order_by(ProcedureSuggestion.created_at.asc()).all()
def _get_pending_suggestion_or_raise(db: Session, suggestion_id: uuid.UUID) -> ProcedureSuggestion:
suggestion = db.query(ProcedureSuggestion).filter(ProcedureSuggestion.id == suggestion_id).first()
if suggestion is None:
raise EntityNotFoundError("ProcedureSuggestion", str(suggestion_id))
if suggestion.status != "pending":
raise InvalidOperationError("This suggestion has already been reviewed.")
return suggestion
def approve_suggestion(db: Session, suggestion_id: uuid.UUID, user: User) -> ProcedureSuggestion:
"""Approve a suggestion: write it into the template and mark reviewed. Does not commit."""
suggestion = _get_pending_suggestion_or_raise(db, suggestion_id)
template = db.query(TestTemplate).filter(TestTemplate.id == suggestion.template_id).first()
if template is None:
raise EntityNotFoundError("TestTemplate", str(suggestion.template_id))
setattr(template, _TEMPLATE_FIELD[suggestion.team], suggestion.suggested_text)
suggestion.status = "approved"
suggestion.reviewed_by = user.id
suggestion.reviewed_at = datetime.utcnow()
db.flush()
return suggestion
def reject_suggestion(db: Session, suggestion_id: uuid.UUID, user: User) -> ProcedureSuggestion:
"""Reject a suggestion without touching the template. Does not commit."""
suggestion = _get_pending_suggestion_or_raise(db, suggestion_id)
suggestion.status = "rejected"
suggestion.reviewed_by = user.id
suggestion.reviewed_at = datetime.utcnow()
db.flush()
return suggestion
@@ -340,6 +340,12 @@ def submit_red_evidence(db: Session, test: Test, user: User) -> Test:
except Exception as e: except Exception as e:
logger.warning("Jira RT end date push failed for test %s: %s", test.id, e, exc_info=True) logger.warning("Jira RT end date push failed for test %s: %s", test.id, e, exc_info=True)
try:
from app.services.procedure_suggestion_service import create_suggestion_from_submission
create_suggestion_from_submission(db, test, "red")
except Exception as e:
logger.warning("Procedure suggestion extraction failed for test %s: %s", test.id, e, exc_info=True)
return test return test
@@ -577,6 +583,12 @@ def submit_blue_evidence(db: Session, test: Test, user: User) -> Test:
except Exception as e: except Exception as e:
logger.warning("Jira BT end date push failed for test %s: %s", test.id, e, exc_info=True) logger.warning("Jira BT end date push failed for test %s: %s", test.id, e, exc_info=True)
try:
from app.services.procedure_suggestion_service import create_suggestion_from_submission
create_suggestion_from_submission(db, test, "blue")
except Exception as e:
logger.warning("Procedure suggestion extraction failed for test %s: %s", test.id, e, exc_info=True)
return test return test
+232
View File
@@ -0,0 +1,232 @@
"""End-to-end coverage for the procedure-suggestion review workflow.
An operator submitting a round with a filled-in procedure field that
contains an extractable command should produce a pending
ProcedureSuggestion against the originating template; a lead can then
approve it (writing it into the template) or reject it (leaving the
template untouched). Nothing is ever written to a template automatically.
"""
import uuid
import pytest
from app.models.enums import TeamSide
from app.models.evidence import Evidence
def _add_evidence(db, test_id, team: TeamSide):
ev = Evidence(
test_id=uuid.UUID(test_id),
file_name="proof.txt",
file_path="s3://bucket/proof.txt",
sha256_hash="a" * 64,
team=team,
)
db.add(ev)
db.commit()
@pytest.fixture
def technique(api, auth_headers):
resp = api(
"post", "/api/v1/techniques", auth_headers,
json={"mitre_id": "T1059.200", "name": "Command Line Suggestions"},
)
assert resp.status_code == 201, resp.text
return resp.json()["id"]
@pytest.fixture
def template(api, red_lead_headers, technique):
resp = api(
"post", "/api/v1/test-templates", red_lead_headers,
json={
"mitre_technique_id": "T1059.200",
"name": "Suggestion source template",
"attack_procedure": "Run a discovery command on the target.",
"detect_suggested_procedure": "Check process creation logs.",
},
)
assert resp.status_code == 201, resp.text
return resp.json()
@pytest.fixture
def test_from_template(api, red_lead_headers, technique, template):
resp = api(
"post", "/api/v1/tests/from-template", red_lead_headers,
json={"template_id": template["id"], "technique_id": technique},
)
assert resp.status_code == 201, resp.text
return resp.json()["id"]
def test_red_submission_with_command_creates_pending_suggestion(
client, db, api, test_from_template, red_tech_headers, red_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery.\nwhoami /priv\nGot back SeDebugPrivilege."},
)
_add_evidence(db, test_id, TeamSide.red)
resp = api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
assert resp.status_code == 200, resp.text
listed = api("get", "/api/v1/procedure-suggestions", red_lead_headers)
assert listed.status_code == 200, listed.text
suggestions = listed.json()
assert len(suggestions) == 1
assert suggestions[0]["team"] == "red"
assert suggestions[0]["suggested_text"] == "whoami /priv"
assert suggestions[0]["status"] == "pending"
def test_blue_submission_with_command_creates_pending_suggestion(
client, db, api, test_from_template, red_tech_headers, red_lead_headers,
blue_tech_headers, blue_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
api("post", f"/api/v1/tests/{test_id}/review-red", red_lead_headers, json={"decision": "approve"})
api("post", f"/api/v1/tests/{test_id}/start-blue-work", blue_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/blue", blue_tech_headers,
json={
"detection_result": "detected",
"detect_procedure": "Checked logs.\nGet-WinEvent -LogName Security | where Id -eq 4688\nFound it.",
},
)
_add_evidence(db, test_id, TeamSide.blue)
resp = api("post", f"/api/v1/tests/{test_id}/submit-blue", blue_tech_headers)
assert resp.status_code == 200, resp.text
listed = api("get", "/api/v1/procedure-suggestions", blue_lead_headers)
assert listed.status_code == 200, listed.text
suggestions = listed.json()
assert len(suggestions) == 1
assert suggestions[0]["team"] == "blue"
assert "Get-WinEvent" in suggestions[0]["suggested_text"]
def test_submission_without_extractable_command_creates_no_suggestion(
client, db, api, test_from_template, red_tech_headers, red_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Just talked to the target and it agreed to be compromised."},
)
_add_evidence(db, test_id, TeamSide.red)
resp = api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
assert resp.status_code == 200, resp.text
listed = api("get", "/api/v1/procedure-suggestions", red_lead_headers)
assert listed.json() == []
def test_approving_a_suggestion_writes_it_into_the_template(
client, db, api, test_from_template, template, red_tech_headers, red_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery.\nwhoami /priv\nDone."},
)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
suggestion = api("get", "/api/v1/procedure-suggestions", red_lead_headers).json()[0]
approve = api(
"post", f"/api/v1/procedure-suggestions/{suggestion['id']}/approve", red_lead_headers,
)
assert approve.status_code == 200, approve.text
assert approve.json()["status"] == "approved"
reloaded = api("get", f"/api/v1/test-templates/{template['id']}", red_lead_headers)
assert reloaded.json()["attack_procedure"] == "whoami /priv"
listed = api("get", "/api/v1/procedure-suggestions", red_lead_headers)
assert listed.json() == []
def test_rejecting_a_suggestion_leaves_the_template_untouched(
client, db, api, test_from_template, template, red_tech_headers, red_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery.\nwhoami /priv\nDone."},
)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
suggestion = api("get", "/api/v1/procedure-suggestions", red_lead_headers).json()[0]
reject = api(
"post", f"/api/v1/procedure-suggestions/{suggestion['id']}/reject", red_lead_headers,
)
assert reject.status_code == 200, reject.text
assert reject.json()["status"] == "rejected"
reloaded = api("get", f"/api/v1/test-templates/{template['id']}", red_lead_headers)
assert reloaded.json()["attack_procedure"] == "Run a discovery command on the target."
def test_blue_lead_cannot_review_a_red_suggestion(
client, db, api, test_from_template, red_tech_headers, red_lead_headers, blue_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery.\nwhoami /priv\nDone."},
)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
suggestion = api("get", "/api/v1/procedure-suggestions", red_lead_headers).json()[0]
# A blue_lead sees no red suggestions in their own queue...
blue_view = api("get", "/api/v1/procedure-suggestions", blue_lead_headers)
assert blue_view.json() == []
# ...and is forbidden from acting on one directly by id.
forbidden = api(
"post", f"/api/v1/procedure-suggestions/{suggestion['id']}/approve", blue_lead_headers,
)
assert forbidden.status_code == 403
def test_duplicate_submission_does_not_create_a_second_pending_suggestion(
client, db, api, test_from_template, red_tech_headers, red_lead_headers,
):
test_id = test_from_template
api("post", f"/api/v1/tests/{test_id}/start-execution", red_tech_headers)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery.\nwhoami /priv\nDone."},
)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
api(
"post", f"/api/v1/tests/{test_id}/review-red", red_lead_headers,
json={"decision": "reopen", "notes": "please redo this round"},
)
api(
"patch", f"/api/v1/tests/{test_id}/red", red_tech_headers,
json={"procedure_text": "Ran discovery again.\nwhoami /priv\nStill works."},
)
_add_evidence(db, test_id, TeamSide.red)
api("post", f"/api/v1/tests/{test_id}/submit-red", red_tech_headers)
listed = api("get", "/api/v1/procedure-suggestions", red_lead_headers).json()
assert len(listed) == 1