ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-24 13:01:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): mark attack paths scan as failed when celery task fails (#10065)

Browse Source
This commit is contained in:
Josema Camacho
2026-02-13 13:20:38 +01:00
committed by GitHub
parent 3defbcd386
commit 02b58d8a31
6 changed files with 283 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/backend/tasks/jobs/attack_paths/db_utils.py
+26 -1
View File
@@ -86,7 +86,11 @@ def finish_attack_paths_scan(
) -> None:
with rls_transaction(attack_paths_scan.tenant_id):
now = datetime.now(tz=timezone.utc)
duration = int((now - attack_paths_scan.started_at).total_seconds())
duration = (
int((now - attack_paths_scan.started_at).total_seconds())
if attack_paths_scan.started_at
else 0
)
attack_paths_scan.state = state
attack_paths_scan.progress = 100
@@ -144,3 +148,24 @@ def update_old_attack_paths_scan(
with rls_transaction(old_attack_paths_scan.tenant_id):
old_attack_paths_scan.is_graph_database_deleted = True
old_attack_paths_scan.save(update_fields=["is_graph_database_deleted"])
def fail_attack_paths_scan(
tenant_id: str,
scan_id: str,
error: str,
) -> None:
"""
Mark the `AttackPathsScan` row as `FAILED` unless it's already `COMPLETED` or `FAILED`.
Used as a safety net when the Celery task fails outside the job's own error handling.
"""
attack_paths_scan = retrieve_attack_paths_scan(tenant_id, scan_id)
if attack_paths_scan and attack_paths_scan.state not in (
StateChoices.COMPLETED,
StateChoices.FAILED,
):
finish_attack_paths_scan(
attack_paths_scan,
StateChoices.FAILED,
{"global_error": error},
)
api/src/backend/tasks/jobs/attack_paths/scan.py
+8 -2
View File
@@ -228,10 +228,16 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
except Exception as e:
exception_message = utils.stringify_exception(e, "Cartography failed")
logger.error(exception_message)
ingestion_exceptions["global_cartography_error"] = exception_message
ingestion_exceptions["global_error"] = exception_message
# Handling databases changes
graph_database.drop_database(tmp_cartography_config.neo4j_database)
try:
graph_database.drop_database(tmp_cartography_config.neo4j_database)
except Exception:
logger.exception(
f"Failed to drop temporary Neo4j database {tmp_cartography_config.neo4j_database} during cleanup"
)
db_utils.finish_attack_paths_scan(
attack_paths_scan, StateChoices.FAILED, ingestion_exceptions
)
api/src/backend/tasks/tasks.py
+19 -1
View File
@@ -10,6 +10,7 @@ from config.django.base import DJANGO_FINDINGS_BATCH_SIZE, DJANGO_TMP_OUTPUT_DIR
from django_celery_beat.models import PeriodicTask
from tasks.jobs.attack_paths import (
attack_paths_scan,
db_utils as attack_paths_db_utils,
can_provider_run_attack_paths_scan,
)
from tasks.jobs.backfill import (
@@ -359,8 +360,25 @@ def perform_scan_summary_task(tenant_id: str, scan_id: str):
return aggregate_findings(tenant_id=tenant_id, scan_id=scan_id)
class AttackPathsScanRLSTask(RLSTask):
"""
RLS task that marks the `AttackPathsScan` DB row as `FAILED` when the Celery task fails.
Covers failures that happen outside the job's own try/except (e.g. provider lookup,
SDK initialization, or Neo4j configuration errors during setup).
"""
def on_failure(self, exc, task_id, args, kwargs, _einfo):
tenant_id = kwargs.get("tenant_id")
scan_id = kwargs.get("scan_id")
if tenant_id and scan_id:
logger.error(f"Attack paths scan task {task_id} failed: {exc}")
attack_paths_db_utils.fail_attack_paths_scan(tenant_id, scan_id, str(exc))
@shared_task(
base=RLSTask,
base=AttackPathsScanRLSTask,
bind=True,
name="attack-paths-scan-perform",
queue="attack-paths-scans",
api/src/backend/tasks/tests/test_attack_paths_scan.py
+221 -3
View File
@@ -244,9 +244,91 @@ class TestAttackPathsRun:
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_cartography_error": "Cartography failed: ingestion boom"
}
assert failure_args[2] == {"global_error": "Cartography failed: ingestion boom"}
def test_run_failure_marks_scan_failed_even_when_drop_database_fails(
self, 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.rls_transaction",
new=lambda *args, **kwargs: nullcontext(),
),
patch(
"tasks.jobs.attack_paths.scan.initialize_prowler_provider",
return_value=MagicMock(_enabled_regions=["us-east-1"]),
),
patch("tasks.jobs.attack_paths.scan.graph_database.get_uri"),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_database_name",
return_value="db-scan-id",
),
patch("tasks.jobs.attack_paths.scan.graph_database.create_database"),
patch(
"tasks.jobs.attack_paths.scan.graph_database.get_session",
return_value=session_ctx,
),
patch("tasks.jobs.attack_paths.scan.cartography_create_indexes.run"),
patch("tasks.jobs.attack_paths.scan.cartography_analysis.run"),
patch("tasks.jobs.attack_paths.scan.findings.create_findings_indexes"),
patch("tasks.jobs.attack_paths.scan.internet.analysis"),
patch("tasks.jobs.attack_paths.scan.findings.analysis"),
patch(
"tasks.jobs.attack_paths.scan.db_utils.retrieve_attack_paths_scan",
return_value=attack_paths_scan,
),
patch("tasks.jobs.attack_paths.scan.db_utils.starting_attack_paths_scan"),
patch(
"tasks.jobs.attack_paths.scan.db_utils.update_attack_paths_scan_progress"
),
patch(
"tasks.jobs.attack_paths.scan.db_utils.finish_attack_paths_scan"
) as mock_finish,
patch(
"tasks.jobs.attack_paths.scan.graph_database.drop_database",
side_effect=ConnectionError("neo4j down"),
),
patch(
"tasks.jobs.attack_paths.scan.get_cartography_ingestion_function",
return_value=ingestion_fn,
),
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.utils.stringify_exception",
return_value="Cartography failed: ingestion boom",
),
):
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"}
def test_run_returns_early_for_unsupported_provider(self, tenants_fixture):
tenant = tenants_fixture[0]
@@ -291,6 +373,142 @@ class TestAttackPathsRun:
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.finish_attack_paths_scan"
) as mock_finish,
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
mock_retrieve.assert_called_once_with(str(tenant.id), str(scan.id))
mock_finish.assert_called_once_with(
attack_paths_scan,
StateChoices.FAILED,
{"global_error": "setup exploded"},
)
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.finish_attack_paths_scan"
) as mock_finish,
):
fail_attack_paths_scan(str(tenant.id), str(scan.id), "setup exploded")
mock_finish.assert_not_called()
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,
),
patch(
"tasks.jobs.attack_paths.db_utils.finish_attack_paths_scan"
) as mock_finish,
):
fail_attack_paths_scan(str(tenant.id), "nonexistent", "setup exploded")
mock_finish.assert_not_called()
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):