diff --git a/backend/alembic/versions/b059_data_classification_scheme.py b/backend/alembic/versions/b059_data_classification_scheme.py new file mode 100644 index 0000000..88c8c9b --- /dev/null +++ b/backend/alembic/versions/b059_data_classification_scheme.py @@ -0,0 +1,41 @@ +"""Rename data_classification values to match the org's real scheme. + +public -> public_release, internal -> general_use, sensitive -> confidential. +restricted is unchanged. Applies to tests, campaigns, and evidences tables +(all plain String(20) columns — no native enum type to alter). + +Revision ID: b059 +Revises: b058 +Create Date: 2026-07-07 +""" + +from alembic import op + +revision = "b059" +down_revision = "b058" +branch_labels = None +depends_on = None + +_TABLES = ("tests", "campaigns", "evidences") + +_RENAMES_UP = { + "public": "public_release", + "internal": "general_use", + "sensitive": "confidential", +} + +_RENAMES_DOWN = {v: k for k, v in _RENAMES_UP.items()} + + +def upgrade() -> None: + for table in _TABLES: + for old, new in _RENAMES_UP.items(): + op.execute(f"UPDATE {table} SET data_classification = '{new}' WHERE data_classification = '{old}'") + op.alter_column(table, "data_classification", server_default="confidential") + + +def downgrade() -> None: + for table in _TABLES: + for new, old in _RENAMES_DOWN.items(): + op.execute(f"UPDATE {table} SET data_classification = '{old}' WHERE data_classification = '{new}'") + op.alter_column(table, "data_classification", server_default="internal") diff --git a/backend/app/domain/enums.py b/backend/app/domain/enums.py index 1b2a5e8..f453635 100644 --- a/backend/app/domain/enums.py +++ b/backend/app/domain/enums.py @@ -91,13 +91,16 @@ class AttackSuccessResult(str, enum.Enum): # Define class DataClassification class DataClassification(str, enum.Enum): - """Data sensitivity classification levels for compliance and retention policies.""" + """Data sensitivity classification levels for compliance and retention policies. - # Assign public = "public" - public = "public" - # Assign internal = "internal" - internal = "internal" - # Assign sensitive = "sensitive" - sensitive = "sensitive" - # Assign restricted = "restricted" + Matches the organization's data protection scheme: + - public_release: approved for external/public release + - general_use: general internal information, no special restriction + - confidential: internal use only, or shared with trusted people + - restricted: trusted people, customer content, and/or PII + """ + + public_release = "public_release" + general_use = "general_use" + confidential = "confidential" restricted = "restricted" diff --git a/backend/app/models/campaign.py b/backend/app/models/campaign.py index f033494..c38b761 100644 --- a/backend/app/models/campaign.py +++ b/backend/app/models/campaign.py @@ -91,7 +91,7 @@ class Campaign(Base): # Assign created_at = Column(DateTime(timezone=True), server_default=func.now()) created_at = Column(DateTime(timezone=True), server_default=func.now()) # Assign data_classification = Column(String(20), nullable=False, server_default="internal") - data_classification = Column(String(20), nullable=False, server_default="internal") + data_classification = Column(String(20), nullable=False, server_default="confidential") # Recurring scheduling fields is_recurring = Column(Boolean, default=False) diff --git a/backend/app/models/evidence.py b/backend/app/models/evidence.py index df92165..87c9d32 100644 --- a/backend/app/models/evidence.py +++ b/backend/app/models/evidence.py @@ -51,7 +51,7 @@ class Evidence(Base): # Assign notes = Column(Text, nullable=True) notes = Column(Text, nullable=True) # Assign data_classification = Column(String(20), nullable=False, server_default="internal") - data_classification = Column(String(20), nullable=False, server_default="internal") + data_classification = Column(String(20), nullable=False, server_default="confidential") # Relationships test = relationship("Test", back_populates="evidences") diff --git a/backend/app/models/test.py b/backend/app/models/test.py index 03c685c..8695781 100644 --- a/backend/app/models/test.py +++ b/backend/app/models/test.py @@ -120,7 +120,7 @@ class Test(Base): # Assign retest_count = Column(Integer, default=0) retest_count = Column(Integer, default=0) # Assign data_classification = Column(String(20), nullable=False, server_default="internal") - data_classification = Column(String(20), nullable=False, server_default="internal") + data_classification = Column(String(20), nullable=False, server_default="confidential") # ── Relationships ─────────────────────────────────────────────── technique = relationship("Technique", back_populates="tests") diff --git a/backend/app/routers/tests.py b/backend/app/routers/tests.py index 258cca9..ffdb075 100644 --- a/backend/app/routers/tests.py +++ b/backend/app/routers/tests.py @@ -41,8 +41,8 @@ from sqlalchemy.orm import Session # Import get_db from app.database from app.database import get_db -# Import get_current_user, require_any_role, require_role from app.dependencies.auth -from app.dependencies.auth import get_current_user, require_any_role, require_role +# Import get_current_user, require_any_role from app.dependencies.auth +from app.dependencies.auth import get_current_user, require_any_role # Import UnitOfWork from app.domain.unit_of_work from app.domain.unit_of_work import UnitOfWork @@ -86,6 +86,7 @@ from app.services.webhook_service import dispatch_webhook from app.services.test_crud_service import ( create_test as crud_create_test, ) +from app.services.test_crud_service import determine_initial_classification # Import from app.services.test_crud_service from app.services.test_crud_service import ( @@ -553,9 +554,14 @@ def update_test_classification( # Entry: db db: Session = Depends(get_db), # Entry: current_user - current_user: User = Depends(require_role("admin")), + current_user: User = Depends(require_any_role( + "admin", "manager", "red_tech", "red_lead", "blue_tech", "blue_lead", + )), ) -> TestOut: - """Update the data classification label for a test (admin only). + """Update the data classification label for a test. + + The initial classification is a best-effort default based on the + technique's tactic — any test participant or admin can correct it. Args: test_id (uuid.UUID): Primary key of the test to classify. @@ -1787,6 +1793,7 @@ def import_rt( # validation_status pending so Blue Lead must confirm detection_result=detection_result, blue_validation_status=None, + data_classification=determine_initial_classification(technique), # Timing execution_date=exec_date_str, created_at=datetime.utcnow(), diff --git a/backend/app/schemas/test.py b/backend/app/schemas/test.py index f8e1781..e6bc423 100644 --- a/backend/app/schemas/test.py +++ b/backend/app/schemas/test.py @@ -37,7 +37,7 @@ class TestCreate(BaseModel): class TestClassificationUpdate(BaseModel): - """Admin-only payload for changing data classification.""" + """Payload for changing a test's data classification (any test participant or admin).""" # data_classification: DataClassification data_classification: DataClassification @@ -299,8 +299,8 @@ class TestOut(BaseModel): retest_of: uuid.UUID | None = None # Assign retest_count = 0 retest_count: int = 0 - # Assign data_classification = "internal" - data_classification: str = "internal" + # Assign data_classification = "confidential" + data_classification: str = "confidential" # Technique info (populated when joined) technique_mitre_id: str | None = None diff --git a/backend/app/seed_demo.py b/backend/app/seed_demo.py index e40522c..a08fc06 100644 --- a/backend/app/seed_demo.py +++ b/backend/app/seed_demo.py @@ -54,6 +54,7 @@ from app.models.test_template import TestTemplate # Import User from app.models.user from app.models.user import User +from app.services.test_crud_service import determine_initial_classification # Assign logger = logging.getLogger(__name__) logger = logging.getLogger(__name__) @@ -348,6 +349,7 @@ def _seed_tests(db: Session, users: list[User], techniques: list[Technique], cou state=state, # Keyword argument: created_at created_at=datetime.utcnow() - timedelta(days=random.randint(0, 90)), + data_classification=determine_initial_classification(technique), ) # Populate team fields based on state diff --git a/backend/app/services/attck_evaluations_service.py b/backend/app/services/attck_evaluations_service.py index 51dcb60..2da64aa 100644 --- a/backend/app/services/attck_evaluations_service.py +++ b/backend/app/services/attck_evaluations_service.py @@ -42,6 +42,7 @@ from app.models.test import Test from app.models.user import User from app.services.audit_service import log_action from app.services.status_service import recalculate_technique_status +from app.services.test_crud_service import determine_initial_classification logger = logging.getLogger(__name__) @@ -631,6 +632,7 @@ def import_evaluation_round( red_validated_by=current_user.id, red_validated_at=datetime.utcnow(), detection_result=detection_result, + data_classification=determine_initial_classification(technique), blue_validation_status=None, execution_date=datetime.utcnow(), created_at=datetime.utcnow(), diff --git a/backend/app/services/jira_service.py b/backend/app/services/jira_service.py index fd7e5ec..34835ed 100644 --- a/backend/app/services/jira_service.py +++ b/backend/app/services/jira_service.py @@ -79,6 +79,7 @@ JIRA_FIELD_ATTACK_SUCCESS = "customfield_11815" JIRA_FIELD_ATTACK_DETECTED = "customfield_11809" JIRA_FIELD_TIME_TO_DETECT = "customfield_11810" JIRA_FIELD_TIME_TO_CONTAIN = "customfield_11811" +JIRA_FIELD_DATA_SENSITIVITY = "customfield_11233" _ATTACK_SUCCESS_TO_JIRA = { "successful": "Yes", @@ -92,6 +93,13 @@ _DETECTION_TO_JIRA = { "partially_detected": "Partial", } +_DATA_CLASSIFICATION_TO_JIRA = { + "public_release": "Public Release", + "general_use": "General Use", + "confidential": "Confidential", + "restricted": "Restricted", +} + # Assign logger = logging.getLogger(__name__) logger = logging.getLogger(__name__) @@ -667,6 +675,10 @@ def auto_create_test_issue( JIRA_FIELD_SEVERITY: severity, } + data_sensitivity = _DATA_CLASSIFICATION_TO_JIRA.get(_enum_value(test.data_classification)) + if data_sensitivity: + fields[JIRA_FIELD_DATA_SENSITIVITY] = data_sensitivity + # Inherit campaign start date if available, otherwise use today from datetime import date as _date effective_start = campaign_start_date or _date.today() diff --git a/backend/app/services/test_crud_service.py b/backend/app/services/test_crud_service.py index 2d4688e..27550e3 100644 --- a/backend/app/services/test_crud_service.py +++ b/backend/app/services/test_crud_service.py @@ -18,7 +18,7 @@ from app.domain.errors import ( EntityNotFoundError, PermissionViolation, ) -from app.models.enums import TestState +from app.models.enums import DataClassification, TestState from app.models.technique import Technique from app.models.test import Test from app.models.test_template import TestTemplate @@ -26,6 +26,27 @@ from app.models.campaign import CampaignTest from app.models.audit import AuditLog from app.utils import escape_like +# Tactics whose test data is more likely to touch customer content, PII, or +# high-impact operational detail — these get an initial classification of +# 'restricted' rather than the 'confidential' baseline. This is a starting +# point only; the assigned operator or lead can always correct it. +_RESTRICTED_TACTICS = frozenset({ + "exfiltration", "collection", "credential-access", "impact", +}) + + +def determine_initial_classification(technique: Technique | None) -> str: + """Best-effort initial data classification for a new test. + + Every security test reveals internal attack/defense posture, so the + baseline is 'confidential' (internal use only) — never public or + general use by default. Escalated to 'restricted' when the technique's + tactic implies customer data, PII, or destructive impact. + """ + if technique and technique.tactic and technique.tactic.lower() in _RESTRICTED_TACTICS: + return DataClassification.restricted.value + return DataClassification.confidential.value + # Define function list_tests def list_tests( @@ -121,6 +142,8 @@ def create_test( # Raise EntityNotFoundError raise EntityNotFoundError("Technique", str(technique_id)) + data_classification = fields.pop("data_classification", None) or determine_initial_classification(technique) + # Assign test = Test( test = Test( # Keyword argument: technique_id @@ -130,6 +153,7 @@ def create_test( # Keyword argument: state state=TestState.draft, created_at=datetime.utcnow(), # explicit — DB column has no server default + data_classification=data_classification, **fields, ) # Stage new record(s) for database insertion @@ -223,6 +247,7 @@ def create_test_from_template( # Keyword argument: state state=TestState.draft, created_at=datetime.utcnow(), # explicit — DB column has no server default + data_classification=determine_initial_classification(technique), ) # Stage new record(s) for database insertion db.add(test) diff --git a/backend/tests/test_data_classification.py b/backend/tests/test_data_classification.py index 45e8b58..0c8274c 100644 --- a/backend/tests/test_data_classification.py +++ b/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"