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
@@ -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
|
||||
Reference in New Issue
Block a user