Aegis/backend/tests/test_data_classification.py

"""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_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,
        name="Classification test",
        created_by=red_lead_user.id,
    )
    db.add(test)
    db.commit()
    db.refresh(test)
    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):
    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": "restricted"},
        headers={"Authorization": f"Bearer {admin_token}"},
    )
    assert response.status_code == 200
    assert response.json()["data_classification"] == "restricted"

    db.refresh(test)
    assert test.data_classification == "restricted"


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": "restricted"},
        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"