"""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 @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=tactic, platforms=["linux"], ) db.add(technique) db.commit() db.refresh(technique) return technique def test_new_test_defaults_to_internal_use_only_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, name="Classification test", created_by=red_lead_user.id, ) db.add(test) db.commit() db.refresh(test) assert test.data_classification == "internal_use_only" 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 pii-tier tactic, so this should land on the baseline. assert resp.json()["data_classification"] == "internal_use_only" def test_admin_can_update_classification(client, db, admin_user, admin_token, red_lead_user): technique = _seed_technique(db) test = Test( technique_id=technique.id, name="Classify me", created_by=red_lead_user.id, state=TestState.draft, ) db.add(test) db.commit() response = client.patch( f"/api/v1/tests/{test.id}/classification", json={"data_classification": "pii"}, headers={"Authorization": f"Bearer {admin_token}"}, ) assert response.status_code == 200 assert response.json()["data_classification"] == "pii" db.refresh(test) assert test.data_classification == "pii" 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, name="Protected", created_by=red_lead_user.id, ) db.add(test) db.commit() response = client.patch( f"/api/v1/tests/{test.id}/classification", json={"data_classification": "pii"}, 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_internal_use_only(): assert determine_initial_classification(_technique_stub("execution")) == "internal_use_only" assert determine_initial_classification(_technique_stub(None)) == "internal_use_only" assert determine_initial_classification(None) == "internal_use_only" def test_determine_initial_classification_escalates_for_sensitive_tactics(): assert determine_initial_classification(_technique_stub("exfiltration")) == "pii" assert determine_initial_classification(_technique_stub("collection")) == "pii" assert determine_initial_classification(_technique_stub("credential-access")) == "pii" assert determine_initial_classification(_technique_stub("impact")) == "pii" assert determine_initial_classification(_technique_stub("Exfiltration")) == "pii"