diff --git a/backend/app/main.py b/backend/app/main.py index b10e57a..1358df2 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -43,6 +43,7 @@ from app.routers import techniques as techniques_router from app.routers import tests as tests_router from app.routers import evidence as evidence_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 metrics as metrics_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") # Call app.include_router() app.include_router(test_templates_router.router, prefix="/api/v1") +app.include_router(procedure_suggestions_router.router, prefix="/api/v1") # Call app.include_router() app.include_router(system_router.router, prefix="/api/v1") # Call app.include_router() diff --git a/backend/app/routers/procedure_suggestions.py b/backend/app/routers/procedure_suggestions.py new file mode 100644 index 0000000..ea0c8d0 --- /dev/null +++ b/backend/app/routers/procedure_suggestions.py @@ -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") diff --git a/backend/app/schemas/procedure_suggestion.py b/backend/app/schemas/procedure_suggestion.py new file mode 100644 index 0000000..8c04d28 --- /dev/null +++ b/backend/app/schemas/procedure_suggestion.py @@ -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) diff --git a/backend/app/services/procedure_suggestion_service.py b/backend/app/services/procedure_suggestion_service.py new file mode 100644 index 0000000..3974ef6 --- /dev/null +++ b/backend/app/services/procedure_suggestion_service.py @@ -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 diff --git a/backend/app/services/test_workflow_service.py b/backend/app/services/test_workflow_service.py index db868e4..b19f5c2 100644 --- a/backend/app/services/test_workflow_service.py +++ b/backend/app/services/test_workflow_service.py @@ -340,6 +340,12 @@ def submit_red_evidence(db: Session, test: Test, user: User) -> Test: except Exception as e: 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 @@ -577,6 +583,12 @@ def submit_blue_evidence(db: Session, test: Test, user: User) -> Test: except Exception as e: 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 diff --git a/backend/tests/test_procedure_suggestions.py b/backend/tests/test_procedure_suggestions.py new file mode 100644 index 0000000..3d8040c --- /dev/null +++ b/backend/tests/test_procedure_suggestions.py @@ -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