feat(classification): org data classification scheme, tactic-based initial default, operator edit rights

backend/tests/test_data_classification.py

@@ -1,15 +1,31 @@
-"""Tests for data classification fields and admin updates."""
+"""Tests for data classification fields and edit permissions."""
+
+from unittest.mock import MagicMock
+
+import pytest
 
 from app.models.enums import TestState
 from app.models.test import Test
 from app.models.technique import Technique
+from app.services.test_crud_service import determine_initial_classification
 
 
-def _seed_technique(db) -> Technique:
+@pytest.fixture
+def technique(client, auth_headers):
+    """Create a technique for test association (mirrors test_tests.py)."""
+    response = client.post(
+        "/api/v1/techniques",
+        json={"mitre_id": "T9998", "name": "Classification Fixture Technique"},
+        headers=auth_headers,
+    )
+    return response.json()
+
+
+def _seed_technique(db, tactic="execution") -> Technique:
     technique = Technique(
         mitre_id="T9999",
         name="Test Technique",
-        tactic="test",
+        tactic=tactic,
         platforms=["linux"],
     )
     db.add(technique)
@@ -18,7 +34,9 @@ def _seed_technique(db) -> Technique:
     return technique
 
 
-def test_new_test_defaults_to_internal(db, red_lead_user):
+def test_new_test_defaults_to_confidential_via_db_default(db, red_lead_user):
+    """A Test() constructed without going through the create_test service
+    still gets a safe classification via the column's server_default."""
     technique = _seed_technique(db)
     test = Test(
         technique_id=technique.id,
@@ -28,7 +46,19 @@ def test_new_test_defaults_to_internal(db, red_lead_user):
     db.add(test)
     db.commit()
     db.refresh(test)
-    assert test.data_classification == "internal"
+    assert test.data_classification == "confidential"
+
+
+def test_create_test_endpoint_uses_tactic_heuristic(client, db, api, auth_headers, technique):
+    """Going through the real create-test flow applies the tactic-based heuristic."""
+    resp = api(
+        "post", "/api/v1/tests", auth_headers,
+        json={"technique_id": technique["id"], "name": "Heuristic test"},
+    )
+    assert resp.status_code == 201, resp.text
+    # The shared `technique` fixture (see conftest/test_tests.py) doesn't set
+    # a restricted-tier tactic, so this should land on the baseline.
+    assert resp.json()["data_classification"] == "confidential"
 
 
 def test_admin_can_update_classification(client, db, admin_user, admin_token, red_lead_user):
@@ -44,17 +74,51 @@ def test_admin_can_update_classification(client, db, admin_user, admin_token, re
 
     response = client.patch(
         f"/api/v1/tests/{test.id}/classification",
-        json={"data_classification": "sensitive"},
+        json={"data_classification": "restricted"},
         headers={"Authorization": f"Bearer {admin_token}"},
     )
     assert response.status_code == 200
-    assert response.json()["data_classification"] == "sensitive"
+    assert response.json()["data_classification"] == "restricted"
 
     db.refresh(test)
-    assert test.data_classification == "sensitive"
+    assert test.data_classification == "restricted"
 
 
-def test_non_admin_cannot_update_classification(client, db, admin_user, red_lead_token, red_lead_user):
+def test_operator_can_update_classification(client, db, admin_user, red_lead_token, red_lead_user):
+    """Any test participant (not just admin) can correct the initial classification."""
+    technique = _seed_technique(db)
+    test = Test(
+        technique_id=technique.id,
+        name="Operator-editable",
+        created_by=red_lead_user.id,
+    )
+    db.add(test)
+    db.commit()
+
+    response = client.patch(
+        f"/api/v1/tests/{test.id}/classification",
+        json={"data_classification": "general_use"},
+        headers={"Authorization": f"Bearer {red_lead_token}"},
+    )
+    assert response.status_code == 200
+    assert response.json()["data_classification"] == "general_use"
+
+
+def test_viewer_cannot_update_classification(client, db, admin_user, red_lead_user):
+    """Viewer is not a test participant and should still be forbidden."""
+    from app.auth import hash_password
+    from app.models.user import User
+
+    viewer = User(
+        username="viewer_classif", email="viewer_classif@test.com",
+        hashed_password=hash_password("x"), role="viewer", is_active=True,
+        must_change_password=False,
+    )
+    db.add(viewer)
+    db.commit()
+    login = client.post("/api/v1/auth/login", data={"username": "viewer_classif", "password": "x"})
+    viewer_token = login.json()["access_token"]
+
     technique = _seed_technique(db)
     test = Test(
         technique_id=technique.id,
@@ -67,6 +131,26 @@ def test_non_admin_cannot_update_classification(client, db, admin_user, red_lead
     response = client.patch(
         f"/api/v1/tests/{test.id}/classification",
         json={"data_classification": "restricted"},
-        headers={"Authorization": f"Bearer {red_lead_token}"},
+        headers={"Authorization": f"Bearer {viewer_token}"},
     )
     assert response.status_code == 403
+
+
+def _technique_stub(tactic):
+    t = MagicMock()
+    t.tactic = tactic
+    return t
+
+
+def test_determine_initial_classification_defaults_to_confidential():
+    assert determine_initial_classification(_technique_stub("execution")) == "confidential"
+    assert determine_initial_classification(_technique_stub(None)) == "confidential"
+    assert determine_initial_classification(None) == "confidential"
+
+
+def test_determine_initial_classification_escalates_for_sensitive_tactics():
+    assert determine_initial_classification(_technique_stub("exfiltration")) == "restricted"
+    assert determine_initial_classification(_technique_stub("collection")) == "restricted"
+    assert determine_initial_classification(_technique_stub("credential-access")) == "restricted"
+    assert determine_initial_classification(_technique_stub("impact")) == "restricted"
+    assert determine_initial_classification(_technique_stub("Exfiltration")) == "restricted"