ben.carrasco/prowler
1
0 Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-10-18 01:14:16 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
9840fa640b8ecceb1526b5e6a2de0fe549422280
prowler/api/src/backend/tasks/tests/test_attack_paths_scan.py
T

2875 lines
101 KiB
Python
Raw Blame History

from contextlib import nullcontext
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock, call, patch
import pytest
from django_celery_results.models import TaskResult
from tasks.jobs.attack_paths import findings as findings_module
from tasks.jobs.attack_paths import indexes as indexes_module
from tasks.jobs.attack_paths import internet as internet_module
from tasks.jobs.attack_paths import sync as sync_module
from tasks.jobs.attack_paths.scan import run as attack_paths_run
from api.models import (
AttackPathsScan,
Finding,
Provider,
Resource,
ResourceFindingMapping,
Scan,
StateChoices,
StatusChoices,
Task,
)
from prowler.lib.check.models import Severity
@pytest.mark.django_db
class TestAttackPathsRun:
# Patching with decorators as we got a `SyntaxError: too many statically nested blocks` error if we use context managers
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph", return_value=0)
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
return_value="bolt://neo4j",
)
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_run_success_flow(
self,
mock_init_provider,
mock_get_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
mock_cartography_analysis,
mock_cartography_ontology,
mock_findings_indexes,
mock_findings_analysis,
mock_internet_analysis,
mock_sync_indexes,
mock_drop_subgraph,
mock_sync,
mock_starting,
mock_update_progress,
mock_finish,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_event_loop,
mock_drop_db,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
ingestion_result = {"organizations": "warning"}
ingestion_fn = MagicMock(return_value=ingestion_result)
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
side_effect=["db-scan-id", "tenant-db"],
) as mock_get_db_name,
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
) as mock_get_session,
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
) as mock_retrieve_scan,
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=ingestion_fn,
) as mock_get_ingestion,
):
result = attack_paths_run(str(tenant.id), str(scan.id), "task-123")
assert result == ingestion_result
mock_retrieve_scan.assert_called_once_with(str(tenant.id), str(scan.id))
mock_starting.assert_called_once()
config = mock_starting.call_args[0][1]
assert config.neo4j_database == "tenant-db"
mock_get_db_name.assert_has_calls(
[call(attack_paths_scan.id, temporary=True), call(provider.tenant_id)]
)
mock_create_db.assert_has_calls([call("db-scan-id"), call("tenant-db")])
mock_get_session.assert_has_calls([call("db-scan-id"), call("tenant-db")])
assert mock_cartography_indexes.call_count == 2
mock_findings_indexes.assert_has_calls([call(mock_session), call(mock_session)])
mock_sync_indexes.assert_called_once_with(mock_session)
# These use tmp_cartography_config (neo4j_database="db-scan-id")
mock_cartography_analysis.assert_called_once()
mock_cartography_ontology.assert_called_once()
mock_internet_analysis.assert_called_once()
mock_findings_analysis.assert_called_once()
mock_drop_subgraph.assert_called_once_with(
database="tenant-db",
provider_id=str(provider.id),
)
mock_sync.assert_called_once_with(
source_database="db-scan-id",
target_database="tenant-db",
tenant_id=str(provider.tenant_id),
provider_id=str(provider.id),
)
mock_get_ingestion.assert_called_once_with(provider.provider)
mock_event_loop.assert_called_once()
mock_update_progress.assert_any_call(attack_paths_scan, 1)
mock_update_progress.assert_any_call(attack_paths_scan, 2)
mock_update_progress.assert_any_call(attack_paths_scan, 95)
mock_update_progress.assert_any_call(attack_paths_scan, 97)
mock_update_progress.assert_any_call(attack_paths_scan, 98)
mock_update_progress.assert_any_call(attack_paths_scan, 99)
mock_finish.assert_called_once_with(
attack_paths_scan, StateChoices.COMPLETED, ingestion_result
)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
)
mock_set_graph_data_ready.assert_called_once_with(attack_paths_scan, True)
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Cartography failed: ingestion boom",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_run_failure_marks_scan_failed(
self,
mock_init_provider,
mock_get_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
mock_cartography_analysis,
mock_findings_indexes,
mock_internet_analysis,
mock_findings_analysis,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=ingestion_fn,
),
):
with pytest.raises(RuntimeError, match="ingestion boom"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
failure_args = mock_finish.call_args[0]
assert failure_args[0] is attack_paths_scan
assert failure_args[1] == StateChoices.FAILED
assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"}
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Cartography failed: ingestion boom",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_failure_before_gate_does_not_flip_graph_data_ready_true(
self,
mock_init_provider,
mock_get_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
mock_cartography_analysis,
mock_findings_indexes,
mock_internet_analysis,
mock_findings_analysis,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
"""Failure during ingestion (before set_provider_graph_data_ready(False))
must NOT flip graph_data_ready to True for providers that never had data."""
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=ingestion_fn,
),
):
with pytest.raises(RuntimeError, match="ingestion boom"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
# Gate was never applied, so recovery must not flip anything to True
mock_set_provider_graph_data_ready.assert_not_called()
mock_set_graph_data_ready.assert_not_called()
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Cartography failed: ingestion boom",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_database",
side_effect=ConnectionError("neo4j down"),
)
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
)
@patch("tasks.jobs.attack_paths.scan.graph_database.get_uri")
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_run_failure_marks_scan_failed_even_when_drop_database_fails(
self,
mock_init_provider,
mock_get_uri,
mock_get_db_name,
mock_create_db,
mock_cartography_indexes,
mock_cartography_analysis,
mock_findings_indexes,
mock_internet_analysis,
mock_findings_analysis,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=ingestion_fn,
),
):
with pytest.raises(RuntimeError, match="ingestion boom"):
attack_paths_run(str(tenant.id), str(scan.id), "task-789")
failure_args = mock_finish.call_args[0]
assert failure_args[0] is attack_paths_scan
assert failure_args[1] == StateChoices.FAILED
assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"}
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Attack Paths scan failed: drop failed",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
)
@patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_subgraph",
side_effect=RuntimeError("drop failed"),
)
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
return_value="bolt://neo4j",
)
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_failure_after_gate_before_drop_restores_graph_data_ready(
self,
mock_init_provider,
mock_get_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
mock_cartography_analysis,
mock_cartography_ontology,
mock_findings_indexes,
mock_findings_analysis,
mock_internet_analysis,
mock_sync_indexes,
mock_drop_subgraph,
mock_sync,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
side_effect=["db-scan-id", "tenant-db"],
),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=MagicMock(return_value={}),
),
):
with pytest.raises(RuntimeError, match="drop failed"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
assert mock_set_provider_graph_data_ready.call_args_list == [
call(attack_paths_scan, False),
call(attack_paths_scan, True),
]
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Attack Paths scan failed: sync failed",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
side_effect=RuntimeError("sync failed"),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph")
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
return_value="bolt://neo4j",
)
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_failure_after_drop_before_sync_leaves_graph_data_ready_false(
self,
mock_init_provider,
mock_get_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
mock_cartography_analysis,
mock_cartography_ontology,
mock_findings_indexes,
mock_findings_analysis,
mock_internet_analysis,
mock_sync_indexes,
mock_drop_subgraph,
mock_sync,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
side_effect=["db-scan-id", "tenant-db"],
),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=MagicMock(return_value={}),
),
):
with pytest.raises(RuntimeError, match="sync failed"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
# Only called with False (gate), never with True (no recovery for partial data)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
)
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Attack Paths scan failed: flag failed",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready",
side_effect=[RuntimeError("flag failed"), None],
)
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_subgraph")
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
return_value="bolt://neo4j",
)
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_failure_after_sync_restores_graph_data_ready(
self,
mock_init_provider,
mock_get_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
mock_cartography_analysis,
mock_cartography_ontology,
mock_findings_indexes,
mock_findings_analysis,
mock_internet_analysis,
mock_sync_indexes,
mock_drop_subgraph,
mock_sync,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
side_effect=["db-scan-id", "tenant-db"],
),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=MagicMock(return_value={}),
),
):
with pytest.raises(RuntimeError, match="flag failed"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
# sync completed: first call (normal path) raised, recovery retried and succeeded
assert mock_set_graph_data_ready.call_args_list == [
call(attack_paths_scan, True),
call(attack_paths_scan, True),
]
# set_provider_graph_data_ready only called once with False (the gate)
mock_set_provider_graph_data_ready.assert_called_once_with(
attack_paths_scan, False
)
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Attack Paths scan failed: drop failed",
)
@patch(
"tasks.jobs.attack_paths.scan.utils.call_within_event_loop",
side_effect=lambda fn, *a, **kw: fn(*a, **kw),
)
@patch("tasks.jobs.attack_paths.scan.graph_database.drop_database")
@patch("tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.set_provider_graph_data_ready")
@patch("tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress")
@patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan")
@patch(
"tasks.jobs.attack_paths.scan.sync.sync_graph",
return_value={"nodes": 0, "relationships": 0},
)
@patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_subgraph",
side_effect=RuntimeError("drop failed"),
)
@patch("tasks.jobs.attack_paths.scan.indexes.create_sync_indexes")
@patch("tasks.jobs.attack_paths.scan.internet.analysis")
@patch("tasks.jobs.attack_paths.scan.findings.analysis", return_value=(0, 0))
@patch("tasks.jobs.attack_paths.scan.indexes.create_findings_indexes")
@patch("tasks.jobs.attack_paths.scan.cartography_ontology.run")
@patch("tasks.jobs.attack_paths.scan.cartography_analysis.run")
@patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run")
@patch("tasks.jobs.attack_paths.scan.graph_database.clear_cache")
@patch("tasks.jobs.attack_paths.scan.graph_database.create_database")
@patch(
"tasks.jobs.attack_paths.scan.graph_database.get_uri",
return_value="bolt://neo4j",
)
@patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
)
@patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_recovery_failure_does_not_suppress_original_exception(
self,
mock_init_provider,
mock_get_uri,
mock_create_db,
mock_clear_cache,
mock_cartography_indexes,
mock_cartography_analysis,
mock_cartography_ontology,
mock_findings_indexes,
mock_findings_analysis,
mock_internet_analysis,
mock_sync_indexes,
mock_drop_subgraph,
mock_sync,
mock_starting,
mock_update_progress,
mock_set_provider_graph_data_ready,
mock_set_graph_data_ready,
mock_finish,
mock_drop_db,
mock_event_loop,
mock_stringify,
tenants_fixture,
providers_fixture,
scans_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.SCHEDULED,
)
# Recovery itself fails on the second call (True)
mock_set_provider_graph_data_ready.side_effect = [
None,
RuntimeError("recovery boom"),
]
mock_session = MagicMock()
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
with (
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
side_effect=["db-scan-id", "tenant-db"],
),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=MagicMock(return_value={}),
),
):
# Original exception propagates despite recovery failure
with pytest.raises(RuntimeError, match="drop failed"):
attack_paths_run(str(tenant.id), str(scan.id), "task-456")
def test_run_returns_early_for_unsupported_provider(self, tenants_fixture):
tenant = tenants_fixture[0]
provider = Provider.objects.create(
provider=Provider.ProviderChoices.GCP,
uid="gcp-account",
alias="gcp",
tenant_id=tenant.id,
)
scan = Scan.objects.create(
name="GCP Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
with (
patch(
"tasks.jobs.attack_paths.scan.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
),
patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(),
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=None,
) as mock_get_ingestion,
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan"
) as mock_retrieve,
):
mock_retrieve.return_value = None
result = attack_paths_run(str(tenant.id), str(scan.id), "task-789")
assert result == {
"global_error": "Provider gcp is not supported for Attack Paths scans"
}
mock_get_ingestion.assert_called_once_with(provider.provider)
mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
@pytest.mark.django_db
class TestFailAttackPathsScan:
def test_marks_executing_scan_as_failed(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
)
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
) as mock_retrieve,
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.drop_database"
) as mock_drop_db,
patch("tasks.jobs.attack_paths.db_utils.recover_graph_data_ready"),
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
expected_tmp_db = f"db-tmp-scan-{str(attack_paths_scan.id).lower()}"
mock_drop_db.assert_called_once_with(expected_tmp_db)
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.state == StateChoices.FAILED
assert attack_paths_scan.ingestion_exceptions == {
"global_error": "setup exploded"
}
def test_drops_temp_database_even_when_drop_fails(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
)
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.drop_database",
side_effect=Exception("Neo4j unreachable"),
),
patch("tasks.jobs.attack_paths.db_utils.recover_graph_data_ready"),
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.state == StateChoices.FAILED
def test_skips_already_failed_scan(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.FAILED,
)
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.drop_database"
) as mock_drop_db,
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
mock_drop_db.assert_called_once()
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.state == StateChoices.FAILED
def test_skips_when_no_scan_found(self, tenants_fixture):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
with patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=None,
):
fail_attack_paths_scan(str(tenant.id), "nonexistent", "setup exploded")
def test_fail_recovers_graph_data_ready_when_data_exists(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
)
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch("tasks.jobs.attack_paths.db_utils.graph_database.drop_database"),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.has_provider_data",
return_value=True,
),
patch(
"tasks.jobs.attack_paths.db_utils.set_provider_graph_data_ready"
) as mock_set_ready,
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "worker died")
mock_set_ready.assert_called_once_with(attack_paths_scan, True)
def test_fail_leaves_graph_data_ready_false_when_no_data(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import fail_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
)
with (
patch(
"tasks.jobs.attack_paths.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch("tasks.jobs.attack_paths.db_utils.graph_database.drop_database"),
patch(
"tasks.jobs.attack_paths.db_utils.graph_database.has_provider_data",
return_value=False,
),
patch(
"tasks.jobs.attack_paths.db_utils.set_provider_graph_data_ready"
) as mock_set_ready,
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "worker died")
mock_set_ready.assert_not_called()
def test_recover_graph_data_ready_never_raises(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import recover_graph_data_ready
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
)
with patch(
"tasks.jobs.attack_paths.db_utils.graph_database.has_provider_data",
side_effect=Exception("Neo4j unreachable"),
):
# Should not raise
recover_graph_data_ready(attack_paths_scan)
class TestAttackPathsScanRLSTaskOnFailure:
def test_on_failure_delegates_to_fail_attack_paths_scan(self):
from tasks.tasks import AttackPathsScanRLSTask
task = AttackPathsScanRLSTask()
with patch(
"tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan"
) as mock_fail:
task.on_failure(
exc=RuntimeError("boom"),
task_id="task-abc",
args=(),
kwargs={"tenant_id": "t-1", "scan_id": "s-1"},
_einfo=None,
)
mock_fail.assert_called_once_with("t-1", "s-1", "boom")
def test_on_failure_skips_when_missing_kwargs(self):
from tasks.tasks import AttackPathsScanRLSTask
task = AttackPathsScanRLSTask()
with patch(
"tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan"
) as mock_fail:
task.on_failure(
exc=RuntimeError("boom"),
task_id="task-abc",
args=(),
kwargs={},
_einfo=None,
)
mock_fail.assert_not_called()
@pytest.mark.django_db
class TestAttackPathsFindingsHelpers:
def test_create_findings_indexes_executes_all_statements(self):
mock_session = MagicMock()
with patch("tasks.jobs.attack_paths.indexes.run_write_query") as mock_run_write:
indexes_module.create_findings_indexes(mock_session)
from tasks.jobs.attack_paths.indexes import FINDINGS_INDEX_STATEMENTS
assert mock_run_write.call_count == len(FINDINGS_INDEX_STATEMENTS)
mock_run_write.assert_has_calls(
[call(mock_session, stmt) for stmt in FINDINGS_INDEX_STATEMENTS]
)
def test_load_findings_batches_requests(self, providers_fixture):
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
# Create a generator that yields two batches of dicts (pre-converted)
def findings_generator():
yield [{"id": "1", "resource_uid": "r-1"}]
yield [{"id": "2", "resource_uid": "r-2"}]
config = SimpleNamespace(update_tag=12345)
mock_session = MagicMock()
first_result = MagicMock()
first_result.single.return_value = {"merged_count": 1, "dropped_count": 0}
second_result = MagicMock()
second_result.single.return_value = {"merged_count": 0, "dropped_count": 1}
mock_session.run.side_effect = [first_result, second_result]
with (
patch(
"tasks.jobs.attack_paths.findings.get_node_uid_field",
return_value="arn",
),
patch(
"tasks.jobs.attack_paths.findings.get_provider_resource_label",
return_value="_AWSResource",
),
patch("tasks.jobs.attack_paths.findings.logger") as mock_logger,
):
findings_module.load_findings(
mock_session, findings_generator(), provider, config
)
assert mock_session.run.call_count == 2
for call_args in mock_session.run.call_args_list:
params = call_args.args[1]
assert params["last_updated"] == config.update_tag
assert "findings_data" in params
summary_log = next(
call_args.args[0]
for call_args in mock_logger.info.call_args_list
if call_args.args and "Finished loading" in call_args.args[0]
)
assert "edges_merged=1" in summary_log
assert "edges_dropped=1" in summary_log
def test_stream_findings_with_resources_returns_latest_scan_data(
self,
tenants_fixture,
providers_fixture,
):
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
resource = Resource.objects.create(
tenant_id=tenant.id,
provider=provider,
uid="resource-uid",
name="Resource",
region="us-east-1",
service="ec2",
type="instance",
)
older_scan = Scan.objects.create(
name="Older",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
old_finding = Finding.objects.create(
tenant_id=tenant.id,
uid="older-finding",
scan=older_scan,
delta=Finding.DeltaChoices.NEW,
status=StatusChoices.PASS,
status_extended="ok",
severity=Severity.low,
impact=Severity.low,
impact_extended="",
raw_result={},
check_id="check-old",
check_metadata={"checktitle": "Old"},
first_seen_at=older_scan.inserted_at,
)
ResourceFindingMapping.objects.create(
tenant_id=tenant.id,
resource=resource,
finding=old_finding,
)
latest_scan = Scan.objects.create(
name="Latest",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
finding = Finding.objects.create(
tenant_id=tenant.id,
uid="finding-uid",
scan=latest_scan,
delta=Finding.DeltaChoices.NEW,
status=StatusChoices.FAIL,
status_extended="failed",
severity=Severity.high,
impact=Severity.high,
impact_extended="",
raw_result={},
check_id="check-1",
check_metadata={"checktitle": "Check title"},
first_seen_at=latest_scan.inserted_at,
)
ResourceFindingMapping.objects.create(
tenant_id=tenant.id,
resource=resource,
finding=finding,
)
latest_scan.refresh_from_db()
with (
patch(
"tasks.jobs.attack_paths.findings.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
),
patch(
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
"default",
),
):
# Generator yields batches, collect all findings from all batches
findings_batches = findings_module.stream_findings_with_resources(
provider,
str(latest_scan.id),
)
findings_data = []
for batch in findings_batches:
findings_data.extend(batch)
assert len(findings_data) == 1
finding_result = findings_data[0]
assert finding_result["id"] == str(finding.id)
assert finding_result["resource_uid"] == resource.uid
assert finding_result["check_title"] == "Check title"
assert finding_result["scan_id"] == str(latest_scan.id)
def test_enrich_batch_with_resources_single_resource(
self,
tenants_fixture,
providers_fixture,
):
"""One finding + one resource = one output dict"""
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
resource = Resource.objects.create(
tenant_id=tenant.id,
provider=provider,
uid="resource-uid-1",
name="Resource 1",
region="us-east-1",
service="ec2",
type="instance",
)
scan = Scan.objects.create(
name="Test Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
finding = Finding.objects.create(
tenant_id=tenant.id,
uid="finding-uid",
scan=scan,
delta=Finding.DeltaChoices.NEW,
status=StatusChoices.FAIL,
status_extended="failed",
severity=Severity.high,
impact=Severity.high,
impact_extended="",
raw_result={},
check_id="check-1",
check_metadata={"checktitle": "Check title"},
first_seen_at=scan.inserted_at,
)
ResourceFindingMapping.objects.create(
tenant_id=tenant.id,
resource=resource,
finding=finding,
)
# Simulate the dict returned by .values()
finding_dict = {
"id": finding.id,
"uid": finding.uid,
"inserted_at": finding.inserted_at,
"updated_at": finding.updated_at,
"first_seen_at": finding.first_seen_at,
"scan_id": scan.id,
"delta": finding.delta,
"status": finding.status,
"status_extended": finding.status_extended,
"severity": finding.severity,
"check_id": finding.check_id,
"check_metadata__checktitle": finding.check_metadata["checktitle"],
"muted": finding.muted,
"muted_reason": finding.muted_reason,
}
# _enrich_batch_with_resources queries ResourceFindingMapping directly
# No RLS mock needed - test DB doesn't enforce RLS policies
with patch(
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
"default",
):
result = findings_module._enrich_batch_with_resources(
[finding_dict], str(tenant.id), lambda uid: f"short:{uid}"
)
assert len(result) == 1
assert result[0]["resource_uid"] == resource.uid
assert result[0]["resource_short_uid"] == f"short:{resource.uid}"
assert result[0]["id"] == str(finding.id)
assert result[0]["status"] == "FAIL"
def test_enrich_batch_with_resources_multiple_resources(
self,
tenants_fixture,
providers_fixture,
):
"""One finding + three resources = three output dicts"""
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
resources = []
for i in range(3):
resource = Resource.objects.create(
tenant_id=tenant.id,
provider=provider,
uid=f"resource-uid-{i}",
name=f"Resource {i}",
region="us-east-1",
service="ec2",
type="instance",
)
resources.append(resource)
scan = Scan.objects.create(
name="Test Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
finding = Finding.objects.create(
tenant_id=tenant.id,
uid="finding-uid",
scan=scan,
delta=Finding.DeltaChoices.NEW,
status=StatusChoices.FAIL,
status_extended="failed",
severity=Severity.high,
impact=Severity.high,
impact_extended="",
raw_result={},
check_id="check-1",
check_metadata={"checktitle": "Check title"},
first_seen_at=scan.inserted_at,
)
# Map finding to all 3 resources
for resource in resources:
ResourceFindingMapping.objects.create(
tenant_id=tenant.id,
resource=resource,
finding=finding,
)
finding_dict = {
"id": finding.id,
"uid": finding.uid,
"inserted_at": finding.inserted_at,
"updated_at": finding.updated_at,
"first_seen_at": finding.first_seen_at,
"scan_id": scan.id,
"delta": finding.delta,
"status": finding.status,
"status_extended": finding.status_extended,
"severity": finding.severity,
"check_id": finding.check_id,
"check_metadata__checktitle": finding.check_metadata["checktitle"],
"muted": finding.muted,
"muted_reason": finding.muted_reason,
}
# _enrich_batch_with_resources queries ResourceFindingMapping directly
# No RLS mock needed - test DB doesn't enforce RLS policies
with patch(
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
"default",
):
result = findings_module._enrich_batch_with_resources(
[finding_dict], str(tenant.id), lambda uid: uid
)
assert len(result) == 3
result_resource_uids = {r["resource_uid"] for r in result}
assert result_resource_uids == {r.uid for r in resources}
# All should have same finding data
for r in result:
assert r["id"] == str(finding.id)
assert r["status"] == "FAIL"
def test_enrich_batch_with_resources_no_resources_skips(
self,
tenants_fixture,
providers_fixture,
):
"""Finding without resources should be skipped"""
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = Scan.objects.create(
name="Test Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
finding = Finding.objects.create(
tenant_id=tenant.id,
uid="orphan-finding",
scan=scan,
delta=Finding.DeltaChoices.NEW,
status=StatusChoices.FAIL,
status_extended="failed",
severity=Severity.high,
impact=Severity.high,
impact_extended="",
raw_result={},
check_id="check-1",
check_metadata={"checktitle": "Check title"},
first_seen_at=scan.inserted_at,
)
# Note: No ResourceFindingMapping created
finding_dict = {
"id": finding.id,
"uid": finding.uid,
"inserted_at": finding.inserted_at,
"updated_at": finding.updated_at,
"first_seen_at": finding.first_seen_at,
"scan_id": scan.id,
"delta": finding.delta,
"status": finding.status,
"status_extended": finding.status_extended,
"severity": finding.severity,
"check_id": finding.check_id,
"check_metadata__checktitle": finding.check_metadata["checktitle"],
"muted": finding.muted,
"muted_reason": finding.muted_reason,
}
# Mock logger to verify no warning is emitted
with (
patch(
"tasks.jobs.attack_paths.findings.READ_REPLICA_ALIAS",
"default",
),
patch("tasks.jobs.attack_paths.findings.logger") as mock_logger,
):
result = findings_module._enrich_batch_with_resources(
[finding_dict], str(tenant.id), lambda uid: uid
)
assert len(result) == 0
mock_logger.warning.assert_not_called()
def test_generator_is_lazy(self, providers_fixture):
"""Generator should not execute queries until iterated"""
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan_id = "some-scan-id"
with patch("tasks.jobs.attack_paths.findings.rls_transaction") as mock_rls:
# Create generator but don't iterate
findings_module.stream_findings_with_resources(provider, scan_id)
# Nothing should be called yet
mock_rls.assert_not_called()
def test_load_findings_empty_generator(self, providers_fixture):
"""Empty generator should not call neo4j"""
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
mock_session = MagicMock()
config = SimpleNamespace(update_tag=12345)
def empty_gen():
return
yield # Make it a generator
with (
patch(
"tasks.jobs.attack_paths.findings.get_node_uid_field",
return_value="arn",
),
patch(
"tasks.jobs.attack_paths.findings.get_provider_resource_label",
return_value="_AWSResource",
),
):
findings_module.load_findings(mock_session, empty_gen(), provider, config)
mock_session.run.assert_not_called()
@pytest.mark.parametrize(
"uid, expected",
[
(
"arn:aws:ec2:us-east-1:552455647653:instance/i-05075b63eb51baacb",
"i-05075b63eb51baacb",
),
(
"arn:aws:ec2:us-east-1:123456789012:volume/vol-0abcd1234ef567890",
"vol-0abcd1234ef567890",
),
(
"arn:aws:ec2:us-east-1:123456789012:security-group/sg-0123abcd",
"sg-0123abcd",
),
("arn:aws:s3:::my-bucket-name", "my-bucket-name"),
("arn:aws:iam::123456789012:role/MyRole", "MyRole"),
(
"arn:aws:lambda:us-east-1:123456789012:function:my-function",
"my-function",
),
("i-05075b63eb51baacb", "i-05075b63eb51baacb"),
],
)
def test_extract_short_uid_aws_variants(self, uid, expected):
from tasks.jobs.attack_paths.aws import extract_short_uid
assert extract_short_uid(uid) == expected
def test_insert_finding_template_has_short_id_fallback(self):
from tasks.jobs.attack_paths.queries import (
INSERT_FINDING_TEMPLATE,
render_cypher_template,
)
rendered = render_cypher_template(
INSERT_FINDING_TEMPLATE,
{
"__NODE_UID_FIELD__": "arn",
"__RESOURCE_LABEL__": "_AWSResource",
},
)
assert (
"resource_by_uid:_AWSResource {arn: finding_data.resource_uid}" in rendered
)
assert "resource_by_id:_AWSResource {id: finding_data.resource_uid}" in rendered
assert (
"resource_by_short:_AWSResource {id: finding_data.resource_short_uid}"
in rendered
)
assert "head(collect(resource_by_short)) AS resource_by_short" in rendered
assert (
"COALESCE(resource_by_uid, resource_by_id, resource_by_short)" in rendered
)
assert "RETURN merged_count, dropped_count" in rendered
class TestAddResourceLabel:
def test_add_resource_label_applies_private_label(self):
mock_session = MagicMock()
first_result = MagicMock()
first_result.single.return_value = {"labeled_count": 5}
second_result = MagicMock()
second_result.single.return_value = {"labeled_count": 0}
mock_session.run.side_effect = [first_result, second_result]
total = findings_module.add_resource_label(mock_session, "aws", "123456789012")
assert total == 5
assert mock_session.run.call_count == 2
query = mock_session.run.call_args_list[0].args[0]
assert "_AWSResource" in query
assert "AWSResource" not in query.replace("_AWSResource", "")
def _make_session_ctx(session, call_order=None, name=None):
"""Create a mock context manager wrapping a mock session."""
ctx = MagicMock()
if call_order is not None and name is not None:
ctx.__enter__ = MagicMock(
side_effect=lambda: (call_order.append(f"{name}:enter"), session)[1]
)
ctx.__exit__ = MagicMock(
side_effect=lambda *a: (call_order.append(f"{name}:exit"), False)[1]
)
else:
ctx.__enter__ = MagicMock(return_value=session)
ctx.__exit__ = MagicMock(return_value=False)
return ctx
class TestSyncNodes:
def test_sync_nodes_adds_private_label(self):
row = {
"internal_id": 1,
"element_id": "elem-1",
"labels": ["SomeLabel"],
"props": {"key": "value"},
}
mock_source_1 = MagicMock()
mock_source_1.run.return_value = [row]
mock_target = MagicMock()
mock_source_2 = MagicMock()
mock_source_2.run.return_value = []
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(mock_source_1),
_make_session_ctx(mock_target),
_make_session_ctx(mock_source_2),
],
):
total = sync_module.sync_nodes(
"source-db", "target-db", "tenant-1", "prov-1"
)
assert total == 1
query = mock_target.run.call_args.args[0]
assert "_ProviderResource" in query
assert "_Tenant_tenant1" in query
assert "_Provider_prov1" in query
def test_sync_nodes_source_closes_before_target_opens(self):
row = {
"internal_id": 1,
"element_id": "elem-1",
"labels": ["SomeLabel"],
"props": {"key": "value"},
}
call_order = []
src_1 = MagicMock()
src_1.run.return_value = [row]
tgt = MagicMock()
src_2 = MagicMock()
src_2.run.return_value = []
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1, call_order, "source1"),
_make_session_ctx(tgt, call_order, "target"),
_make_session_ctx(src_2, call_order, "source2"),
],
):
sync_module.sync_nodes("src-db", "tgt-db", "t-1", "p-1")
assert call_order.index("source1:exit") < call_order.index("target:enter")
def test_sync_nodes_pagination_with_batch_size_1(self):
row_a = {
"internal_id": 1,
"element_id": "elem-1",
"labels": ["LabelA"],
"props": {"a": 1},
}
row_b = {
"internal_id": 2,
"element_id": "elem-2",
"labels": ["LabelB"],
"props": {"b": 2},
}
src_1 = MagicMock()
src_1.run.return_value = [row_a]
src_2 = MagicMock()
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
tgt_1 = MagicMock()
tgt_2 = MagicMock()
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(tgt_1),
_make_session_ctx(src_2),
_make_session_ctx(tgt_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
):
total = sync_module.sync_nodes("src", "tgt", "t-1", "p-1")
assert total == 2
assert src_1.run.call_args.args[1]["last_id"] == -1
assert src_2.run.call_args.args[1]["last_id"] == 1
def test_sync_nodes_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[_make_session_ctx(src)],
) as mock_get_session:
total = sync_module.sync_nodes("src", "tgt", "t-1", "p-1")
assert total == 0
assert mock_get_session.call_count == 1
class TestSyncRelationships:
def test_sync_relationships_source_closes_before_target_opens(self):
row = {
"internal_id": 1,
"rel_type": "HAS",
"start_element_id": "s-1",
"end_element_id": "e-1",
"props": {},
}
call_order = []
src_1 = MagicMock()
src_1.run.return_value = [row]
tgt = MagicMock()
src_2 = MagicMock()
src_2.run.return_value = []
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1, call_order, "source1"),
_make_session_ctx(tgt, call_order, "target"),
_make_session_ctx(src_2, call_order, "source2"),
],
):
sync_module.sync_relationships("src", "tgt", "p-1")
assert call_order.index("source1:exit") < call_order.index("target:enter")
def test_sync_relationships_pagination_with_batch_size_1(self):
row_a = {
"internal_id": 1,
"rel_type": "HAS",
"start_element_id": "s-1",
"end_element_id": "e-1",
"props": {"a": 1},
}
row_b = {
"internal_id": 2,
"rel_type": "CONNECTS",
"start_element_id": "s-2",
"end_element_id": "e-2",
"props": {"b": 2},
}
src_1 = MagicMock()
src_1.run.return_value = [row_a]
src_2 = MagicMock()
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
tgt_1 = MagicMock()
tgt_2 = MagicMock()
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(tgt_1),
_make_session_ctx(src_2),
_make_session_ctx(tgt_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
):
total = sync_module.sync_relationships("src", "tgt", "p-1")
assert total == 2
assert src_1.run.call_args.args[1]["last_id"] == -1
assert src_2.run.call_args.args[1]["last_id"] == 1
def test_sync_relationships_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[_make_session_ctx(src)],
) as mock_get_session:
total = sync_module.sync_relationships("src", "tgt", "p-1")
assert total == 0
assert mock_get_session.call_count == 1
class TestInternetAnalysis:
def _make_provider_and_config(self):
provider = MagicMock()
provider.provider = "aws"
provider.uid = "123456789012"
config = SimpleNamespace(update_tag=1234567890)
return provider, config
def test_analysis_creates_node_and_relationships(self):
"""Verify both Cypher statements are executed and relationship count returned."""
mock_session = MagicMock()
mock_result = MagicMock()
mock_result.single.return_value = {"relationships_merged": 3}
mock_session.run.side_effect = [None, mock_result]
provider, config = self._make_provider_and_config()
with patch(
"tasks.jobs.attack_paths.internet.get_root_node_label",
return_value="AWSAccount",
):
result = internet_module.analysis(mock_session, provider, config)
assert mock_session.run.call_count == 2
assert result == 3
def test_analysis_zero_exposed_resources(self):
"""When no resources are exposed, zero relationships are created."""
mock_session = MagicMock()
mock_result = MagicMock()
mock_result.single.return_value = {"relationships_merged": 0}
mock_session.run.side_effect = [None, mock_result]
provider, config = self._make_provider_and_config()
with patch(
"tasks.jobs.attack_paths.internet.get_root_node_label",
return_value="AWSAccount",
):
result = internet_module.analysis(mock_session, provider, config)
assert result == 0
@pytest.mark.django_db
class TestAttackPathsDbUtilsGraphDataReady:
"""Tests for db_utils functions related to graph_data_ready lifecycle."""
def test_create_attack_paths_scan_first_scan_defaults_to_false(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import create_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
attack_paths_scan = create_attack_paths_scan(
str(tenant.id), str(scan.id), provider.id
)
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is False
def test_create_attack_paths_scan_inherits_true_from_previous(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import create_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.COMPLETED,
graph_data_ready=True,
)
new_scan = Scan.objects.create(
name="New Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
attack_paths_scan = create_attack_paths_scan(
str(tenant.id), str(new_scan.id), provider.id
)
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is True
def test_create_attack_paths_scan_inherits_false_when_no_previous_ready(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import create_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.FAILED,
graph_data_ready=False,
)
new_scan = Scan.objects.create(
name="New Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
attack_paths_scan = create_attack_paths_scan(
str(tenant.id), str(new_scan.id), provider.id
)
assert attack_paths_scan is not None
assert attack_paths_scan.graph_data_ready is False
def test_set_graph_data_ready_updates_field(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import set_graph_data_ready
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
graph_data_ready=True,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
set_graph_data_ready(attack_paths_scan, False)
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.graph_data_ready is False
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
set_graph_data_ready(attack_paths_scan, True)
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.graph_data_ready is True
def test_finish_attack_paths_scan_does_not_modify_graph_data_ready(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import finish_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
graph_data_ready=True,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
finish_attack_paths_scan(attack_paths_scan, StateChoices.COMPLETED, {})
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.state == StateChoices.COMPLETED
assert attack_paths_scan.graph_data_ready is True
def test_finish_attack_paths_scan_preserves_graph_data_ready_on_failure(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import finish_attack_paths_scan
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan = scans_fixture[0]
scan.provider = provider
scan.save()
attack_paths_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
graph_data_ready=True,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
finish_attack_paths_scan(
attack_paths_scan,
StateChoices.FAILED,
{"global_error": "boom"},
)
attack_paths_scan.refresh_from_db()
assert attack_paths_scan.state == StateChoices.FAILED
assert attack_paths_scan.graph_data_ready is True
def test_set_provider_graph_data_ready_updates_all_scans_for_provider(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import set_provider_graph_data_ready
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
scan_a = scans_fixture[0]
scan_a.provider = provider
scan_a.save()
scan_b = Scan.objects.create(
name="Second Scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.AVAILABLE,
tenant_id=tenant.id,
)
old_ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan_a,
state=StateChoices.COMPLETED,
graph_data_ready=True,
)
new_ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan_b,
state=StateChoices.EXECUTING,
graph_data_ready=True,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
set_provider_graph_data_ready(new_ap_scan, False)
old_ap_scan.refresh_from_db()
new_ap_scan.refresh_from_db()
assert old_ap_scan.graph_data_ready is False
assert new_ap_scan.graph_data_ready is False
def test_set_provider_graph_data_ready_does_not_affect_other_providers(
self, tenants_fixture, providers_fixture, scans_fixture
):
from tasks.jobs.attack_paths.db_utils import set_provider_graph_data_ready
tenant = tenants_fixture[0]
provider_a = providers_fixture[0]
provider_a.provider = Provider.ProviderChoices.AWS
provider_a.save()
provider_b = providers_fixture[1]
provider_b.provider = Provider.ProviderChoices.AWS
provider_b.save()
scan_a = scans_fixture[0]
scan_a.provider = provider_a
scan_a.save()
scan_b = Scan.objects.create(
name="Scan for provider B",
provider=provider_b,
trigger=Scan.TriggerChoices.MANUAL,
state=StateChoices.COMPLETED,
tenant_id=tenant.id,
)
ap_scan_a = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider_a,
scan=scan_a,
state=StateChoices.EXECUTING,
graph_data_ready=True,
)
ap_scan_b = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider_b,
scan=scan_b,
state=StateChoices.COMPLETED,
graph_data_ready=True,
)
with patch(
"tasks.jobs.attack_paths.db_utils.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
):
set_provider_graph_data_ready(ap_scan_a, False)
ap_scan_a.refresh_from_db()
ap_scan_b.refresh_from_db()
assert ap_scan_a.graph_data_ready is False
assert ap_scan_b.graph_data_ready is True
@pytest.mark.django_db
class TestCleanupStaleAttackPathsScans:
def _create_executing_scan(
self, tenant, provider, scan=None, started_at=None, worker=None
):
"""Helper to create an EXECUTING AttackPathsScan with optional Task+TaskResult."""
ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=scan,
state=StateChoices.EXECUTING,
started_at=started_at or datetime.now(tz=timezone.utc),
)
task_result = None
if worker is not None:
task_result = TaskResult.objects.create(
task_id=str(ap_scan.id),
task_name="attack-paths-scan-perform",
status="STARTED",
worker=worker,
)
task = Task.objects.create(
id=task_result.task_id,
task_runner_task=task_result,
tenant_id=tenant.id,
)
ap_scan.task = task
ap_scan.save(update_fields=["task_id"])
return ap_scan, task_result
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
def test_cleans_up_scan_with_dead_worker(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
# Recent scan — should still be cleaned up because worker is dead
ap_scan, task_result = self._create_executing_scan(
tenant, provider, worker="dead-worker@host"
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 1
assert str(ap_scan.id) in result["scan_ids"]
mock_drop_db.assert_called_once()
mock_recover.assert_called_once()
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.FAILED
assert ap_scan.progress == 100
assert ap_scan.completed_at is not None
assert ap_scan.ingestion_exceptions == {
"global_error": "Worker dead — cleaned up by periodic task"
}
task_result.refresh_from_db()
assert task_result.status == "FAILURE"
assert task_result.date_done is not None
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=True)
def test_revokes_and_cleans_scan_exceeding_threshold_on_live_worker(
self,
mock_alive,
mock_revoke,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
old_start = datetime.now(tz=timezone.utc) - timedelta(hours=49)
ap_scan, task_result = self._create_executing_scan(
tenant, provider, started_at=old_start, worker="live-worker@host"
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 1
mock_revoke.assert_called_once_with(task_result)
mock_recover.assert_called_once()
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.FAILED
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=True)
def test_ignores_recent_executing_scans_on_live_worker(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
# Recent scan on live worker — should be skipped
self._create_executing_scan(tenant, provider, worker="live-worker@host")
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 0
mock_drop_db.assert_not_called()
mock_recover.assert_not_called()
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_ignores_completed_and_failed_scans(
self,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
state=StateChoices.COMPLETED,
)
AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
state=StateChoices.FAILED,
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 0
mock_drop_db.assert_not_called()
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch(
"tasks.jobs.attack_paths.cleanup.graph_database.drop_database",
side_effect=Exception("Neo4j unreachable"),
)
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
def test_handles_drop_database_failure_gracefully(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
self._create_executing_scan(tenant, provider, worker="dead-worker@host")
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 1
mock_drop_db.assert_called_once()
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
def test_cross_tenant_cleanup(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant1 = tenants_fixture[0]
tenant2 = tenants_fixture[1]
provider1 = providers_fixture[0]
provider1.provider = Provider.ProviderChoices.AWS
provider1.save()
provider2 = Provider.objects.create(
provider="aws",
uid="999888777666",
alias="aws_tenant2",
tenant_id=tenant2.id,
)
ap_scan1, _ = self._create_executing_scan(
tenant1, provider1, worker="dead-worker-1@host"
)
ap_scan2, _ = self._create_executing_scan(
tenant2, provider2, worker="dead-worker-2@host"
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 2
assert mock_recover.call_count == 2
ap_scan1.refresh_from_db()
ap_scan2.refresh_from_db()
assert ap_scan1.state == StateChoices.FAILED
assert ap_scan2.state == StateChoices.FAILED
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
def test_recovers_graph_data_ready_for_stale_scan(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
ap_scan, _ = self._create_executing_scan(
tenant, provider, worker="dead-worker@host"
)
cleanup_stale_attack_paths_scans()
mock_recover.assert_called_once()
recovered_scan = mock_recover.call_args[0][0]
assert recovered_scan.id == ap_scan.id
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
def test_fallback_to_time_heuristic_when_no_worker_field(
self,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
# Old scan with no Task/TaskResult
old_start = datetime.now(tz=timezone.utc) - timedelta(hours=49)
ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
state=StateChoices.EXECUTING,
started_at=old_start,
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 1
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.FAILED
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
def test_shared_worker_is_pinged_only_once(
self,
mock_alive,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
scans_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
# Two scans on the same dead worker
self._create_executing_scan(tenant, provider, worker="shared-worker@host")
self._create_executing_scan(tenant, provider, worker="shared-worker@host")
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 2
# Worker should be pinged exactly once — cache prevents second ping
mock_alive.assert_called_once_with("shared-worker@host")
# `SCHEDULED` state cleanup
def _create_scheduled_scan(
self,
tenant,
provider,
*,
age_minutes,
parent_state,
with_task=True,
):
"""Create a SCHEDULED AttackPathsScan with a parent Scan in `parent_state`.
`age_minutes` controls how far in the past `started_at` is set, so
callers can place rows safely past the cleanup cutoff.
"""
parent_scan = Scan.objects.create(
name="Parent Prowler scan",
provider=provider,
trigger=Scan.TriggerChoices.MANUAL,
state=parent_state,
tenant_id=tenant.id,
)
ap_scan = AttackPathsScan.objects.create(
tenant_id=tenant.id,
provider=provider,
scan=parent_scan,
state=StateChoices.SCHEDULED,
started_at=datetime.now(tz=timezone.utc) - timedelta(minutes=age_minutes),
)
task_result = None
if with_task:
task_result = TaskResult.objects.create(
task_id=str(ap_scan.id),
task_name="attack-paths-scan-perform",
status="PENDING",
)
task = Task.objects.create(
id=task_result.task_id,
task_runner_task=task_result,
tenant_id=tenant.id,
)
ap_scan.task = task
ap_scan.save(update_fields=["task_id"])
return ap_scan, task_result
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
def test_cleans_up_scheduled_scan_when_parent_is_terminal(
self,
mock_revoke,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
ap_scan, task_result = self._create_scheduled_scan(
tenant,
provider,
age_minutes=24 * 60 * 3, # 3 days, safely past any threshold
parent_state=StateChoices.FAILED,
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 1
assert str(ap_scan.id) in result["scan_ids"]
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.FAILED
assert ap_scan.progress == 100
assert ap_scan.completed_at is not None
assert ap_scan.ingestion_exceptions == {
"global_error": "Scan never started — cleaned up by periodic task"
}
# SCHEDULED revoke must NOT terminate a running worker
mock_revoke.assert_called_once()
assert mock_revoke.call_args.kwargs == {"terminate": False}
# Temp DB never created for SCHEDULED, so no drop attempted
mock_drop_db.assert_not_called()
# Tenant Neo4j data is untouched in this path
mock_recover.assert_not_called()
task_result.refresh_from_db()
assert task_result.status == "FAILURE"
assert task_result.date_done is not None
@patch("tasks.jobs.attack_paths.cleanup.recover_graph_data_ready")
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
@patch(
"tasks.jobs.attack_paths.cleanup.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
)
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
def test_skips_scheduled_scan_when_parent_still_in_flight(
self,
mock_revoke,
mock_drop_db,
mock_recover,
tenants_fixture,
providers_fixture,
):
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
tenant = tenants_fixture[0]
provider = providers_fixture[0]
provider.provider = Provider.ProviderChoices.AWS
provider.save()
ap_scan, _ = self._create_scheduled_scan(
tenant,
provider,
age_minutes=24 * 60 * 3,
parent_state=StateChoices.EXECUTING,
)
result = cleanup_stale_attack_paths_scans()
assert result["cleaned_up_count"] == 0
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.SCHEDULED
mock_revoke.assert_not_called()
Reference in new issue View Git Blame Copy Permalink