mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-07-23 04:21:52 +00:00
fix(api): prevent false Attack Paths stale cleanup (#11986)
This commit is contained in:
Josema Camacho
@@ -0,0 +1 @@
|
||||
`attack-paths-cleanup-stale-scans` now retries worker pings and checks recent scan activity before failing scans and removing temporary databases
|
||||
@@ -152,10 +152,10 @@ def execute_custom_query(
|
||||
scan: AttackPathsScan,
|
||||
) -> dict[str, Any]:
|
||||
# Defense-in-depth for custom queries:
|
||||
# 1. `neo4j.READ_ACCESS` — prevents mutations at the driver level
|
||||
# 2. `inject_provider_label()` — regex-based label injection scopes node patterns
|
||||
# 3. `_serialize_graph()` — post-query filter drops nodes without the provider label
|
||||
# 4. `USING QUERY:TIMEOUTMILLISECONDS` on Neptune — server-side runaway cutoff
|
||||
# 1. `neo4j.READ_ACCESS` - prevents mutations at the driver level
|
||||
# 2. `inject_provider_label()` - regex-based label injection scopes node patterns
|
||||
# 3. `_serialize_graph()` - post-query filter drops nodes without the provider label
|
||||
# 4. `USING QUERY:TIMEOUTMILLISECONDS` on Neptune - server-side runaway cutoff
|
||||
#
|
||||
# Layer 2 is best-effort (regex can't fully parse Cypher);
|
||||
# layer 3 is the safety net that guarantees provider isolation.
|
||||
|
||||
@@ -308,6 +308,9 @@ CSRF_COOKIE_SECURE = True
|
||||
SESSION_COOKIE_SECURE = True
|
||||
|
||||
# Attack Paths
|
||||
ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES = env.int(
|
||||
"ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES", 30
|
||||
)
|
||||
ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES = env.int(
|
||||
"ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES", 2880
|
||||
) # 48h
|
||||
|
||||
@@ -1,40 +1,50 @@
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from functools import partial
|
||||
|
||||
from api.attack_paths import database as graph_database
|
||||
from api.db_router import MainRouter
|
||||
from api.db_utils import rls_transaction
|
||||
from api.models import AttackPathsScan, StateChoices
|
||||
from celery import states
|
||||
from celery import current_app, states
|
||||
from celery.utils.log import get_task_logger
|
||||
from config.django.base import ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES
|
||||
from config.django.base import (
|
||||
ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES,
|
||||
ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES,
|
||||
)
|
||||
from django.db import DatabaseError
|
||||
from django.db.transaction import on_commit
|
||||
from tasks.jobs.attack_paths.db_utils import (
|
||||
mark_scan_finished,
|
||||
recover_graph_data_ready,
|
||||
)
|
||||
from tasks.jobs.orphan_recovery import is_worker_alive as _is_worker_alive
|
||||
from tasks.jobs.orphan_recovery import revoke_task as _revoke_task
|
||||
|
||||
logger = get_task_logger(__name__)
|
||||
|
||||
WORKER_PING_BASE_TIMEOUT_SECONDS = 5
|
||||
WORKER_PING_MAX_ATTEMPTS = 3
|
||||
|
||||
|
||||
def cleanup_stale_attack_paths_scans() -> dict:
|
||||
"""
|
||||
Mark stale `AttackPathsScan` rows as `FAILED`.
|
||||
|
||||
Covers two stuck-state scenarios:
|
||||
1. `EXECUTING` scans whose workers are dead, or that have exceeded the
|
||||
stale threshold while alive.
|
||||
2. `SCHEDULED` scans that never made it to a worker — parent scan
|
||||
1. `EXECUTING` scans whose workers are unresponsive and whose rows have
|
||||
stopped receiving progress updates, or that exceeded the stale threshold.
|
||||
2. `SCHEDULED` scans that never made it to a worker - parent scan
|
||||
crashed before dispatch, broker lost the message, etc. Detected by
|
||||
age plus the parent `Scan` no longer being in flight.
|
||||
"""
|
||||
threshold = timedelta(minutes=ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES)
|
||||
now = datetime.now(tz=UTC)
|
||||
cutoff = now - threshold
|
||||
stale_cutoff = now - timedelta(minutes=ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES)
|
||||
inactivity_cutoff = now - timedelta(
|
||||
minutes=ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES
|
||||
)
|
||||
|
||||
cleaned_up: list[str] = []
|
||||
cleaned_up.extend(_cleanup_stale_executing_scans(cutoff))
|
||||
cleaned_up.extend(_cleanup_stale_scheduled_scans(cutoff))
|
||||
cleaned_up.extend(_cleanup_stale_executing_scans(stale_cutoff, inactivity_cutoff))
|
||||
cleaned_up.extend(_cleanup_stale_scheduled_scans(stale_cutoff))
|
||||
|
||||
logger.info(
|
||||
f"Stale `AttackPathsScan` cleanup: {len(cleaned_up)} scan(s) cleaned up"
|
||||
@@ -42,13 +52,57 @@ def cleanup_stale_attack_paths_scans() -> dict:
|
||||
return {"cleaned_up_count": len(cleaned_up), "scan_ids": cleaned_up}
|
||||
|
||||
|
||||
def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]:
|
||||
def _ping_workers(workers: set[str]) -> tuple[set[str], set[str] | None]:
|
||||
"""Ping worker destinations in parallel and retry only missing workers.
|
||||
|
||||
The second tuple item is `None` when the final ping attempt raises. In that
|
||||
case the pending workers have unknown liveness and their scans must be kept.
|
||||
"""
|
||||
pending = set(workers)
|
||||
responsive: set[str] = set()
|
||||
|
||||
for attempt in range(WORKER_PING_MAX_ATTEMPTS):
|
||||
if not pending:
|
||||
return responsive, set()
|
||||
|
||||
timeout = WORKER_PING_BASE_TIMEOUT_SECONDS * 2**attempt
|
||||
try:
|
||||
response = current_app.control.inspect(
|
||||
destination=sorted(pending), timeout=timeout
|
||||
).ping()
|
||||
except Exception:
|
||||
attempts_remaining = WORKER_PING_MAX_ATTEMPTS - attempt - 1
|
||||
if attempts_remaining:
|
||||
logger.warning(
|
||||
f"Attack Paths worker ping attempt {attempt + 1} failed; "
|
||||
f"retrying pending workers with {attempts_remaining} "
|
||||
"attempt(s) remaining",
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
|
||||
logger.exception(
|
||||
"Attack Paths worker ping attempts exhausted; preserving scans "
|
||||
"for workers with unknown liveness"
|
||||
)
|
||||
return responsive, None
|
||||
|
||||
responded = pending.intersection((response or {}).keys())
|
||||
responsive.update(responded)
|
||||
pending.difference_update(responded)
|
||||
|
||||
return responsive, pending
|
||||
|
||||
|
||||
def _cleanup_stale_executing_scans(
|
||||
stale_cutoff: datetime, inactivity_cutoff: datetime
|
||||
) -> list[str]:
|
||||
"""
|
||||
Two-pass detection for `EXECUTING` scans:
|
||||
1. If `TaskResult.worker` exists, ping the worker.
|
||||
- Dead worker: cleanup immediately (any age).
|
||||
- Alive + past threshold: revoke the task, then cleanup.
|
||||
- Alive + within threshold: skip.
|
||||
1. Ping all recorded workers in parallel with bounded retries.
|
||||
- Responsive + past stale threshold: cleanup.
|
||||
- Unresponsive + past inactivity threshold: cleanup.
|
||||
- Unknown after a final ping exception: preserve.
|
||||
2. If no worker field: fall back to time-based heuristic only.
|
||||
"""
|
||||
executing_scans = list(
|
||||
@@ -57,14 +111,13 @@ def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]:
|
||||
.select_related("task__task_runner_task")
|
||||
)
|
||||
|
||||
# Cache worker liveness so each worker is pinged at most once
|
||||
workers = {
|
||||
tr.worker
|
||||
for scan in executing_scans
|
||||
if (tr := getattr(scan.task, "task_runner_task", None) if scan.task else None)
|
||||
and tr.worker
|
||||
}
|
||||
worker_alive = {w: _is_worker_alive(w) for w in workers}
|
||||
responsive_workers, unresponsive_workers = _ping_workers(workers)
|
||||
|
||||
cleaned_up: list[str] = []
|
||||
|
||||
@@ -75,27 +128,50 @@ def _cleanup_stale_executing_scans(cutoff: datetime) -> list[str]:
|
||||
worker = task_result.worker if task_result else None
|
||||
|
||||
if worker:
|
||||
alive = worker_alive.get(worker, True)
|
||||
|
||||
if alive:
|
||||
if scan.started_at and scan.started_at >= cutoff:
|
||||
if worker in responsive_workers:
|
||||
if scan.started_at is None or scan.started_at >= stale_cutoff:
|
||||
continue
|
||||
|
||||
# Alive but stale — revoke before cleanup
|
||||
_revoke_task(task_result)
|
||||
reason = "Scan exceeded stale threshold — cleaned up by periodic task"
|
||||
reason = "Scan exceeded stale threshold - cleaned up by periodic task"
|
||||
recheck_activity_cutoff = None
|
||||
elif unresponsive_workers is None or worker not in unresponsive_workers:
|
||||
logger.info(
|
||||
f"Preserving scan {scan.id}: worker {worker} liveness is "
|
||||
f"unknown (progress={scan.progress}, updated_at={scan.updated_at})"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
reason = "Worker dead — cleaned up by periodic task"
|
||||
if scan.updated_at >= inactivity_cutoff:
|
||||
logger.info(
|
||||
f"Preserving scan {scan.id}: worker {worker} is unresponsive "
|
||||
f"but activity is recent (progress={scan.progress}, "
|
||||
f"updated_at={scan.updated_at})"
|
||||
)
|
||||
continue
|
||||
|
||||
reason = (
|
||||
"Worker unresponsive and scan inactive for "
|
||||
f"{ATTACK_PATHS_SCAN_INACTIVITY_THRESHOLD_MINUTES} minutes - "
|
||||
"cleaned up by periodic task"
|
||||
)
|
||||
recheck_activity_cutoff = inactivity_cutoff
|
||||
else:
|
||||
# No worker recorded, time-based heuristic only
|
||||
if scan.started_at and scan.started_at >= cutoff:
|
||||
if scan.started_at is None or scan.started_at >= stale_cutoff:
|
||||
continue
|
||||
reason = (
|
||||
"No worker recorded, scan exceeded stale threshold — "
|
||||
"No worker recorded, scan exceeded stale threshold - "
|
||||
"cleaned up by periodic task"
|
||||
)
|
||||
recheck_activity_cutoff = None
|
||||
|
||||
if _cleanup_scan(scan, task_result, reason):
|
||||
if _cleanup_scan(
|
||||
scan,
|
||||
task_result,
|
||||
reason,
|
||||
revoke=worker is not None,
|
||||
inactivity_cutoff=recheck_activity_cutoff,
|
||||
):
|
||||
cleaned_up.append(str(scan.id))
|
||||
|
||||
return cleaned_up
|
||||
@@ -112,10 +188,9 @@ def _cleanup_stale_scheduled_scans(cutoff: datetime) -> list[str]:
|
||||
avoids cleaning up rows whose parent Prowler scan is legitimately still
|
||||
running.
|
||||
|
||||
For each match: revoke the queued task (best-effort; harmless if already
|
||||
consumed), atomically flip to `FAILED`, and mark the `TaskResult`. The
|
||||
temp Neo4j database is never created while `SCHEDULED`, so no drop is
|
||||
needed.
|
||||
For each match: lock and recheck the row, mark the scan and `TaskResult` as
|
||||
failed, then revoke the queued task after the transaction commits. The temp
|
||||
Neo4j database is never created while `SCHEDULED`, so no drop is needed.
|
||||
"""
|
||||
scheduled_scans = list(
|
||||
AttackPathsScan.all_objects.using(MainRouter.admin_db)
|
||||
@@ -141,42 +216,54 @@ def _cleanup_stale_scheduled_scans(cutoff: datetime) -> list[str]:
|
||||
task_result = (
|
||||
getattr(scan.task, "task_runner_task", None) if scan.task else None
|
||||
)
|
||||
if task_result:
|
||||
_revoke_task(task_result, terminate=False)
|
||||
|
||||
reason = "Scan never started — cleaned up by periodic task"
|
||||
reason = "Scan never started - cleaned up by periodic task"
|
||||
if _cleanup_scheduled_scan(scan, task_result, reason):
|
||||
cleaned_up.append(str(scan.id))
|
||||
|
||||
return cleaned_up
|
||||
|
||||
|
||||
def _cleanup_scan(scan, task_result, reason: str) -> bool:
|
||||
def _cleanup_scan(
|
||||
scan,
|
||||
task_result,
|
||||
reason: str,
|
||||
*,
|
||||
revoke: bool = False,
|
||||
inactivity_cutoff: datetime | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Clean up a single stale `AttackPathsScan`:
|
||||
drop temp DB, mark `FAILED`, update `TaskResult`, recover `graph_data_ready`.
|
||||
lock and recheck, mark `FAILED`, revoke after commit, drop the temp DB, and
|
||||
recover graph readiness.
|
||||
|
||||
Returns `True` if the scan was actually cleaned up, `False` if skipped.
|
||||
"""
|
||||
scan_id_str = str(scan.id)
|
||||
|
||||
# Drop temp Neo4j database
|
||||
try:
|
||||
fresh_scan = _finalize_failed_scan(
|
||||
scan,
|
||||
StateChoices.EXECUTING,
|
||||
reason,
|
||||
task_result=task_result,
|
||||
revoke=revoke,
|
||||
inactivity_cutoff=inactivity_cutoff,
|
||||
)
|
||||
except DatabaseError:
|
||||
logger.exception(
|
||||
f"Failed to mark stale Attack Paths scan {scan_id_str} as failed"
|
||||
)
|
||||
return False
|
||||
|
||||
if fresh_scan is None:
|
||||
return False
|
||||
|
||||
tmp_db_name = graph_database.get_database_name(scan.id, temporary=True)
|
||||
try:
|
||||
graph_database.drop_database(tmp_db_name)
|
||||
except Exception:
|
||||
logger.exception(f"Failed to drop temp database {tmp_db_name}")
|
||||
|
||||
fresh_scan = _finalize_failed_scan(scan, StateChoices.EXECUTING, reason)
|
||||
if fresh_scan is None:
|
||||
return False
|
||||
|
||||
# Mark `TaskResult` as `FAILURE` (not RLS-protected, outside lock)
|
||||
if task_result:
|
||||
task_result.status = states.FAILURE
|
||||
task_result.date_done = datetime.now(tz=UTC)
|
||||
task_result.save(update_fields=["status", "date_done"])
|
||||
|
||||
recover_graph_data_ready(fresh_scan)
|
||||
|
||||
logger.info(f"Cleaned up stale scan {scan_id_str}: {reason}")
|
||||
@@ -187,31 +274,49 @@ def _cleanup_scheduled_scan(scan, task_result, reason: str) -> bool:
|
||||
"""
|
||||
Clean up a `SCHEDULED` scan that never reached a worker.
|
||||
|
||||
Skips the temp Neo4j drop — the database is only created once the worker
|
||||
Skips the temp Neo4j drop - the database is only created once the worker
|
||||
enters `EXECUTING`, so dropping it here just produces noisy log output.
|
||||
|
||||
Returns `True` if the scan was actually cleaned up, `False` if skipped.
|
||||
"""
|
||||
scan_id_str = str(scan.id)
|
||||
|
||||
fresh_scan = _finalize_failed_scan(scan, StateChoices.SCHEDULED, reason)
|
||||
if fresh_scan is None:
|
||||
try:
|
||||
fresh_scan = _finalize_failed_scan(
|
||||
scan,
|
||||
StateChoices.SCHEDULED,
|
||||
reason,
|
||||
task_result=task_result,
|
||||
revoke=task_result is not None,
|
||||
terminate=False,
|
||||
)
|
||||
except DatabaseError:
|
||||
logger.exception(
|
||||
f"Failed to mark scheduled Attack Paths scan {scan_id_str} as failed"
|
||||
)
|
||||
return False
|
||||
|
||||
if task_result:
|
||||
task_result.status = states.FAILURE
|
||||
task_result.date_done = datetime.now(tz=UTC)
|
||||
task_result.save(update_fields=["status", "date_done"])
|
||||
if fresh_scan is None:
|
||||
return False
|
||||
|
||||
logger.info(f"Cleaned up scheduled scan {scan_id_str}: {reason}")
|
||||
return True
|
||||
|
||||
|
||||
def _finalize_failed_scan(scan, expected_state: str, reason: str):
|
||||
def _finalize_failed_scan(
|
||||
scan,
|
||||
expected_state: str,
|
||||
reason: str,
|
||||
*,
|
||||
task_result=None,
|
||||
revoke: bool = False,
|
||||
terminate: bool = True,
|
||||
inactivity_cutoff: datetime | None = None,
|
||||
):
|
||||
"""
|
||||
Atomically lock the row, verify it's still in `expected_state`, and
|
||||
mark it `FAILED`. Returns the locked row on success, `None` if the
|
||||
row is gone or has already moved on.
|
||||
Atomically lock the row, verify it's still eligible, and mark it `FAILED`.
|
||||
If requested, register revocation after commit. Returns the locked row on
|
||||
success, `None` if the row is gone or has already moved on.
|
||||
"""
|
||||
scan_id_str = str(scan.id)
|
||||
with rls_transaction(str(scan.tenant_id)):
|
||||
@@ -225,6 +330,23 @@ def _finalize_failed_scan(scan, expected_state: str, reason: str):
|
||||
logger.info(f"Scan {scan_id_str} is now {fresh_scan.state}, skipping")
|
||||
return None
|
||||
|
||||
if inactivity_cutoff is not None and fresh_scan.updated_at >= inactivity_cutoff:
|
||||
logger.info(
|
||||
f"Scan {scan_id_str} received activity during worker checks, skipping"
|
||||
)
|
||||
return None
|
||||
|
||||
mark_scan_finished(fresh_scan, StateChoices.FAILED, {"global_error": reason})
|
||||
|
||||
if task_result:
|
||||
task_result.status = states.FAILURE
|
||||
task_result.date_done = datetime.now(tz=UTC)
|
||||
task_result.save(update_fields=["status", "date_done"])
|
||||
|
||||
if revoke and task_result:
|
||||
on_commit(
|
||||
partial(_revoke_task, task_result, terminate=terminate),
|
||||
using=fresh_scan._state.db,
|
||||
)
|
||||
|
||||
return fresh_scan
|
||||
|
||||
@@ -126,14 +126,17 @@ def starting_attack_paths_scan(
|
||||
if locked.state != StateChoices.SCHEDULED:
|
||||
return False
|
||||
|
||||
now = datetime.now(tz=UTC)
|
||||
locked.state = StateChoices.EXECUTING
|
||||
locked.started_at = datetime.now(tz=UTC)
|
||||
locked.started_at = now
|
||||
locked.updated_at = now
|
||||
locked.update_tag = cartography_config.update_tag
|
||||
locked.save(update_fields=["state", "started_at", "update_tag"])
|
||||
locked.save(update_fields=["state", "started_at", "updated_at", "update_tag"])
|
||||
|
||||
# Keep the in-memory object the caller is holding in sync.
|
||||
attack_paths_scan.state = locked.state
|
||||
attack_paths_scan.started_at = locked.started_at
|
||||
attack_paths_scan.updated_at = locked.updated_at
|
||||
attack_paths_scan.update_tag = locked.update_tag
|
||||
return True
|
||||
|
||||
@@ -181,7 +184,8 @@ def update_attack_paths_scan_progress(
|
||||
) -> None:
|
||||
with rls_transaction(attack_paths_scan.tenant_id):
|
||||
attack_paths_scan.progress = progress
|
||||
attack_paths_scan.save(update_fields=["progress"])
|
||||
attack_paths_scan.updated_at = datetime.now(tz=UTC)
|
||||
attack_paths_scan.save(update_fields=["progress", "updated_at"])
|
||||
|
||||
|
||||
def set_graph_data_ready(
|
||||
|
||||
@@ -172,11 +172,9 @@ def reconcile_orphans(
|
||||
window_hours: int = 6,
|
||||
dry_run: bool = False,
|
||||
) -> dict:
|
||||
"""Run the full orphan sweep under a single-flight advisory lock.
|
||||
"""Run the orphan task sweep under a single-flight advisory lock.
|
||||
|
||||
Recovers any orphaned in-flight task and delegates attack-paths scans that
|
||||
never reached a worker to their existing stale-cleanup. Returns a summary;
|
||||
a no-op (lock not won) is reported too.
|
||||
Returns a recovery summary. A no-op is reported when the lock is not acquired.
|
||||
"""
|
||||
with advisory_lock() as acquired:
|
||||
if not acquired:
|
||||
@@ -200,11 +198,6 @@ def reconcile_orphans(
|
||||
logger.info("Orphan task recovery disabled by feature flag")
|
||||
result = {"recovered": [], "failed": [], "skipped": [], "enabled": False}
|
||||
|
||||
if not dry_run:
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
result["attack_paths"] = cleanup_stale_attack_paths_scans()
|
||||
|
||||
return {"acquired": True, **result}
|
||||
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ from api.models import (
|
||||
StatusChoices,
|
||||
Task,
|
||||
)
|
||||
from django.db import DEFAULT_DB_ALIAS
|
||||
from django.db import DEFAULT_DB_ALIAS, DatabaseError
|
||||
from django_celery_results.models import TaskResult
|
||||
from prowler.lib.check.models import Severity
|
||||
from tasks.jobs.attack_paths import findings as findings_module
|
||||
@@ -2686,10 +2686,194 @@ class TestAttackPathsDbUtilsGraphDataReady:
|
||||
assert ap_scan_b.graph_data_ready is True
|
||||
|
||||
|
||||
class TestAttackPathsWorkerPing:
|
||||
@patch("tasks.jobs.attack_paths.cleanup.current_app")
|
||||
def test_pings_workers_in_parallel_and_retries_only_missing(self, mock_app):
|
||||
from tasks.jobs.attack_paths.cleanup import _ping_workers
|
||||
|
||||
first_ping = MagicMock(return_value={"worker-a@host": {"ok": "pong"}})
|
||||
second_ping = MagicMock(return_value={"worker-b@host": {"ok": "pong"}})
|
||||
third_ping = MagicMock(return_value={"worker-c@host": {"ok": "pong"}})
|
||||
mock_app.control.inspect.side_effect = [
|
||||
MagicMock(ping=first_ping),
|
||||
MagicMock(ping=second_ping),
|
||||
MagicMock(ping=third_ping),
|
||||
]
|
||||
|
||||
responsive, unresponsive = _ping_workers(
|
||||
{"worker-c@host", "worker-a@host", "worker-b@host"}
|
||||
)
|
||||
|
||||
assert responsive == {
|
||||
"worker-a@host",
|
||||
"worker-b@host",
|
||||
"worker-c@host",
|
||||
}
|
||||
assert unresponsive == set()
|
||||
assert mock_app.control.inspect.call_args_list == [
|
||||
call(
|
||||
destination=["worker-a@host", "worker-b@host", "worker-c@host"],
|
||||
timeout=5,
|
||||
),
|
||||
call(destination=["worker-b@host", "worker-c@host"], timeout=10),
|
||||
call(destination=["worker-c@host"], timeout=20),
|
||||
]
|
||||
|
||||
@patch("tasks.jobs.attack_paths.cleanup.logger")
|
||||
@patch("tasks.jobs.attack_paths.cleanup.current_app")
|
||||
def test_retries_intermediate_ping_exceptions(self, mock_app, mock_logger):
|
||||
from tasks.jobs.attack_paths.cleanup import _ping_workers
|
||||
|
||||
mock_app.control.inspect.side_effect = [
|
||||
MagicMock(ping=MagicMock(side_effect=ConnectionError("first"))),
|
||||
MagicMock(ping=MagicMock(side_effect=ConnectionError("second"))),
|
||||
MagicMock(ping=MagicMock(return_value={})),
|
||||
]
|
||||
|
||||
responsive, unresponsive = _ping_workers({"worker@host"})
|
||||
|
||||
assert responsive == set()
|
||||
assert unresponsive == {"worker@host"}
|
||||
assert mock_logger.warning.call_count == 2
|
||||
assert all(
|
||||
warning.kwargs["exc_info"] is True
|
||||
for warning in mock_logger.warning.call_args_list
|
||||
)
|
||||
mock_logger.exception.assert_not_called()
|
||||
|
||||
@patch("tasks.jobs.attack_paths.cleanup.logger")
|
||||
@patch("tasks.jobs.attack_paths.cleanup.current_app")
|
||||
def test_final_ping_exception_leaves_pending_workers_unknown(
|
||||
self, mock_app, mock_logger
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import _ping_workers
|
||||
|
||||
mock_app.control.inspect.side_effect = [
|
||||
MagicMock(ping=MagicMock(return_value={"worker-a@host": {"ok": "pong"}})),
|
||||
MagicMock(ping=MagicMock(return_value={})),
|
||||
MagicMock(ping=MagicMock(side_effect=ConnectionError("final"))),
|
||||
]
|
||||
|
||||
responsive, unresponsive = _ping_workers({"worker-a@host", "worker-b@host"})
|
||||
|
||||
assert responsive == {"worker-a@host"}
|
||||
assert unresponsive is None
|
||||
mock_logger.exception.assert_called_once()
|
||||
|
||||
@patch("tasks.jobs.attack_paths.cleanup.logger")
|
||||
@patch("tasks.jobs.attack_paths.cleanup.current_app")
|
||||
def test_worker_can_respond_after_an_intermediate_exception(
|
||||
self, mock_app, mock_logger
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import _ping_workers
|
||||
|
||||
mock_app.control.inspect.side_effect = [
|
||||
MagicMock(ping=MagicMock(side_effect=ConnectionError("first"))),
|
||||
MagicMock(ping=MagicMock(side_effect=ConnectionError("second"))),
|
||||
MagicMock(ping=MagicMock(return_value={"worker@host": {"ok": "pong"}})),
|
||||
]
|
||||
|
||||
responsive, unresponsive = _ping_workers({"worker@host"})
|
||||
|
||||
assert responsive == {"worker@host"}
|
||||
assert unresponsive == set()
|
||||
assert mock_logger.warning.call_count == 2
|
||||
mock_logger.exception.assert_not_called()
|
||||
|
||||
|
||||
class TestAttackPathsCleanupTask:
|
||||
@patch(
|
||||
"tasks.tasks.cleanup_stale_attack_paths_scans",
|
||||
return_value={"cleaned_up_count": 1, "scan_ids": ["scan-id"]},
|
||||
)
|
||||
def test_hourly_task_invokes_attack_paths_cleanup(self, mock_cleanup):
|
||||
from tasks.tasks import cleanup_stale_attack_paths_scans_task
|
||||
|
||||
result = cleanup_stale_attack_paths_scans_task.run()
|
||||
|
||||
assert result == {"cleaned_up_count": 1, "scan_ids": ["scan-id"]}
|
||||
mock_cleanup.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestAttackPathsDbUtilsActivity:
|
||||
@patch(
|
||||
"tasks.jobs.attack_paths.db_utils.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
def test_starting_scan_refreshes_updated_at(
|
||||
self, tenants_fixture, aws_provider, scans_fixture
|
||||
):
|
||||
from tasks.jobs.attack_paths.db_utils import starting_attack_paths_scan
|
||||
|
||||
old_updated_at = datetime.now(tz=UTC) - timedelta(hours=1)
|
||||
attack_paths_scan = AttackPathsScan.objects.create(
|
||||
tenant_id=tenants_fixture[0].id,
|
||||
provider=aws_provider,
|
||||
scan=scans_fixture[0],
|
||||
state=StateChoices.SCHEDULED,
|
||||
)
|
||||
AttackPathsScan.objects.filter(id=attack_paths_scan.id).update(
|
||||
updated_at=old_updated_at
|
||||
)
|
||||
attack_paths_scan.refresh_from_db()
|
||||
|
||||
started = starting_attack_paths_scan(
|
||||
attack_paths_scan, SimpleNamespace(update_tag=123)
|
||||
)
|
||||
|
||||
assert attack_paths_scan.updated_at > old_updated_at
|
||||
attack_paths_scan.refresh_from_db()
|
||||
assert started is True
|
||||
assert attack_paths_scan.updated_at > old_updated_at
|
||||
|
||||
@patch(
|
||||
"tasks.jobs.attack_paths.db_utils.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
def test_progress_update_refreshes_updated_at(
|
||||
self, tenants_fixture, aws_provider, scans_fixture
|
||||
):
|
||||
from tasks.jobs.attack_paths.db_utils import update_attack_paths_scan_progress
|
||||
|
||||
old_updated_at = datetime.now(tz=UTC) - timedelta(hours=1)
|
||||
attack_paths_scan = AttackPathsScan.objects.create(
|
||||
tenant_id=tenants_fixture[0].id,
|
||||
provider=aws_provider,
|
||||
scan=scans_fixture[0],
|
||||
state=StateChoices.EXECUTING,
|
||||
)
|
||||
AttackPathsScan.objects.filter(id=attack_paths_scan.id).update(
|
||||
updated_at=old_updated_at
|
||||
)
|
||||
attack_paths_scan.refresh_from_db()
|
||||
|
||||
update_attack_paths_scan_progress(attack_paths_scan, 42)
|
||||
|
||||
assert attack_paths_scan.updated_at > old_updated_at
|
||||
attack_paths_scan.refresh_from_db()
|
||||
assert attack_paths_scan.progress == 42
|
||||
assert attack_paths_scan.updated_at > old_updated_at
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
class TestCleanupStaleAttackPathsScans:
|
||||
@pytest.fixture(autouse=True)
|
||||
def execute_on_commit_callbacks(self):
|
||||
with patch(
|
||||
"tasks.jobs.attack_paths.cleanup.on_commit",
|
||||
side_effect=lambda callback, **kwargs: callback(),
|
||||
):
|
||||
yield
|
||||
|
||||
def _create_executing_scan(
|
||||
self, tenant, provider, scan=None, started_at=None, worker=None
|
||||
self,
|
||||
tenant,
|
||||
provider,
|
||||
scan=None,
|
||||
started_at=None,
|
||||
updated_at=None,
|
||||
worker=None,
|
||||
):
|
||||
"""Helper to create an EXECUTING AttackPathsScan with optional Task+TaskResult."""
|
||||
ap_scan = AttackPathsScan.objects.create(
|
||||
@@ -2716,18 +2900,66 @@ class TestCleanupStaleAttackPathsScans:
|
||||
ap_scan.task = task
|
||||
ap_scan.save(update_fields=["task_id"])
|
||||
|
||||
if updated_at is not None:
|
||||
AttackPathsScan.objects.filter(id=ap_scan.id).update(updated_at=updated_at)
|
||||
ap_scan.updated_at = updated_at
|
||||
|
||||
return ap_scan, task_result
|
||||
|
||||
@patch(
|
||||
"tasks.jobs.attack_paths.cleanup.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
def test_defers_revoke_until_scan_failure_is_persisted(
|
||||
self,
|
||||
mock_revoke,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import _finalize_failed_scan
|
||||
|
||||
ap_scan, task_result = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
|
||||
with patch("tasks.jobs.attack_paths.cleanup.on_commit") as mock_on_commit:
|
||||
finalized_scan = _finalize_failed_scan(
|
||||
ap_scan,
|
||||
StateChoices.EXECUTING,
|
||||
"Cleanup reason",
|
||||
task_result=task_result,
|
||||
revoke=True,
|
||||
)
|
||||
|
||||
assert finalized_scan is not None
|
||||
ap_scan.refresh_from_db()
|
||||
task_result.refresh_from_db()
|
||||
assert ap_scan.state == StateChoices.FAILED
|
||||
assert task_result.status == "FAILURE"
|
||||
mock_revoke.assert_not_called()
|
||||
mock_on_commit.assert_called_once()
|
||||
assert mock_on_commit.call_args.kwargs == {"using": DEFAULT_DB_ALIAS}
|
||||
|
||||
callback = mock_on_commit.call_args.args[0]
|
||||
callback()
|
||||
|
||||
mock_revoke.assert_called_once_with(task_result, terminate=True)
|
||||
|
||||
@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(
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_cleans_up_inactive_scan_with_unresponsive_worker(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -2735,19 +2967,39 @@ class TestCleanupStaleAttackPathsScans:
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
from tasks.jobs.attack_paths.db_utils import mark_scan_finished
|
||||
|
||||
tenant = tenants_fixture[0]
|
||||
provider = aws_provider
|
||||
|
||||
# Recent scan — should still be cleaned up because worker is dead
|
||||
updated_at = datetime.now(tz=UTC) - timedelta(minutes=31)
|
||||
ap_scan, task_result = self._create_executing_scan(
|
||||
tenant, provider, worker="dead-worker@host"
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=updated_at,
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"unresponsive-worker@host"})
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
with patch(
|
||||
"tasks.jobs.attack_paths.cleanup.mark_scan_finished",
|
||||
wraps=mark_scan_finished,
|
||||
) as mock_mark_failed:
|
||||
call_order = MagicMock()
|
||||
call_order.attach_mock(mock_revoke, "revoke")
|
||||
call_order.attach_mock(mock_mark_failed, "mark_failed")
|
||||
call_order.attach_mock(mock_drop_db, "drop_database")
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 1
|
||||
assert str(ap_scan.id) in result["scan_ids"]
|
||||
assert [entry[0] for entry in call_order.mock_calls] == [
|
||||
"mark_failed",
|
||||
"revoke",
|
||||
"drop_database",
|
||||
]
|
||||
mock_revoke.assert_called_once_with(task_result, terminate=True)
|
||||
mock_drop_db.assert_called_once()
|
||||
mock_recover.assert_called_once()
|
||||
|
||||
@@ -2756,7 +3008,10 @@ class TestCleanupStaleAttackPathsScans:
|
||||
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"
|
||||
"global_error": (
|
||||
"Worker unresponsive and scan inactive for 30 minutes - "
|
||||
"cleaned up by periodic task"
|
||||
)
|
||||
}
|
||||
|
||||
task_result.refresh_from_db()
|
||||
@@ -2770,10 +3025,10 @@ class TestCleanupStaleAttackPathsScans:
|
||||
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(
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_revokes_and_cleans_scan_exceeding_threshold_on_responsive_worker(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
@@ -2790,11 +3045,12 @@ class TestCleanupStaleAttackPathsScans:
|
||||
ap_scan, task_result = self._create_executing_scan(
|
||||
tenant, provider, started_at=old_start, worker="live-worker@host"
|
||||
)
|
||||
mock_ping.return_value = ({"live-worker@host"}, set())
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 1
|
||||
mock_revoke.assert_called_once_with(task_result)
|
||||
mock_revoke.assert_called_once_with(task_result, terminate=True)
|
||||
mock_recover.assert_called_once()
|
||||
|
||||
ap_scan.refresh_from_db()
|
||||
@@ -2806,10 +3062,10 @@ class TestCleanupStaleAttackPathsScans:
|
||||
"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(
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_ignores_recent_executing_scans_on_responsive_worker(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -2821,8 +3077,8 @@ class TestCleanupStaleAttackPathsScans:
|
||||
tenant = tenants_fixture[0]
|
||||
provider = aws_provider
|
||||
|
||||
# Recent scan on live worker — should be skipped
|
||||
self._create_executing_scan(tenant, provider, worker="live-worker@host")
|
||||
mock_ping.return_value = ({"live-worker@host"}, set())
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
@@ -2830,6 +3086,130 @@ class TestCleanupStaleAttackPathsScans:
|
||||
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(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_preserves_recent_scan_on_unresponsive_worker(
|
||||
self,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(minutes=29),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"unresponsive-worker@host"})
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
mock_revoke.assert_not_called()
|
||||
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(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers", return_value=(set(), None))
|
||||
def test_final_ping_exception_preserves_pending_worker_scan(
|
||||
self,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
ap_scan, _ = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(hours=1),
|
||||
worker="unknown-worker@host",
|
||||
)
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
mock_ping.assert_called_once_with({"unknown-worker@host"})
|
||||
ap_scan.refresh_from_db()
|
||||
assert ap_scan.state == StateChoices.EXECUTING
|
||||
mock_revoke.assert_not_called()
|
||||
mock_drop_db.assert_not_called()
|
||||
mock_recover.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("inactive_seconds", "should_clean"),
|
||||
[(29 * 60 + 59, False), (30 * 60, False), (30 * 60 + 1, True)],
|
||||
)
|
||||
@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._ping_workers")
|
||||
def test_inactivity_boundary_is_strict(
|
||||
self,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
inactive_seconds,
|
||||
should_clean,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
now = datetime.now(tz=UTC)
|
||||
ap_scan, task_result = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=now - timedelta(seconds=inactive_seconds),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"unresponsive-worker@host"})
|
||||
|
||||
with patch("tasks.jobs.attack_paths.cleanup.datetime") as mock_datetime:
|
||||
mock_datetime.now.return_value = now
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == int(should_clean)
|
||||
ap_scan.refresh_from_db()
|
||||
expected_state = StateChoices.FAILED if should_clean else StateChoices.EXECUTING
|
||||
assert ap_scan.state == expected_state
|
||||
if should_clean:
|
||||
mock_revoke.assert_called_once_with(task_result, terminate=True)
|
||||
mock_drop_db.assert_called_once()
|
||||
mock_recover.assert_called_once()
|
||||
else:
|
||||
mock_revoke.assert_not_called()
|
||||
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(
|
||||
@@ -2868,16 +3248,20 @@ class TestCleanupStaleAttackPathsScans:
|
||||
@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"),
|
||||
side_effect=[Exception("Neo4j unreachable"), None],
|
||||
)
|
||||
@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(
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
@patch("tasks.jobs.attack_paths.cleanup.logger")
|
||||
def test_neo4j_failure_leaves_scan_failed_and_continues(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_logger,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -2889,12 +3273,75 @@ class TestCleanupStaleAttackPathsScans:
|
||||
tenant = tenants_fixture[0]
|
||||
provider = aws_provider
|
||||
|
||||
self._create_executing_scan(tenant, provider, worker="dead-worker@host")
|
||||
updated_at = datetime.now(tz=UTC) - timedelta(minutes=31)
|
||||
ap_scan_1, _ = self._create_executing_scan(
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=updated_at,
|
||||
worker="unresponsive-worker-1@host",
|
||||
)
|
||||
ap_scan_2, _ = self._create_executing_scan(
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=updated_at,
|
||||
worker="unresponsive-worker-2@host",
|
||||
)
|
||||
mock_ping.return_value = (
|
||||
set(),
|
||||
{"unresponsive-worker-1@host", "unresponsive-worker-2@host"},
|
||||
)
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 1
|
||||
mock_drop_db.assert_called_once()
|
||||
assert result["cleaned_up_count"] == 2
|
||||
assert mock_revoke.call_count == 2
|
||||
assert mock_drop_db.call_count == 2
|
||||
mock_logger.exception.assert_called_once()
|
||||
ap_scan_1.refresh_from_db()
|
||||
ap_scan_2.refresh_from_db()
|
||||
assert ap_scan_1.state == StateChoices.FAILED
|
||||
assert ap_scan_2.state == StateChoices.FAILED
|
||||
|
||||
@patch("tasks.jobs.attack_paths.cleanup.graph_database.drop_database")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
@patch(
|
||||
"tasks.jobs.attack_paths.cleanup.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
@patch(
|
||||
"tasks.jobs.attack_paths.cleanup.mark_scan_finished",
|
||||
side_effect=DatabaseError("PostgreSQL unavailable"),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup.logger")
|
||||
def test_postgresql_failure_prevents_revoke_and_neo4j_deletion(
|
||||
self,
|
||||
mock_logger,
|
||||
mock_mark_failed,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(minutes=31),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"unresponsive-worker@host"})
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
mock_mark_failed.assert_called_once()
|
||||
mock_logger.exception.assert_called_once()
|
||||
mock_revoke.assert_not_called()
|
||||
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")
|
||||
@@ -2902,10 +3349,12 @@ class TestCleanupStaleAttackPathsScans:
|
||||
"tasks.jobs.attack_paths.cleanup.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_cross_tenant_cleanup(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -2924,16 +3373,28 @@ class TestCleanupStaleAttackPathsScans:
|
||||
tenant_id=tenant2.id,
|
||||
)
|
||||
|
||||
updated_at = datetime.now(tz=UTC) - timedelta(minutes=31)
|
||||
ap_scan1, _ = self._create_executing_scan(
|
||||
tenant1, provider1, worker="dead-worker-1@host"
|
||||
tenant1,
|
||||
provider1,
|
||||
updated_at=updated_at,
|
||||
worker="unresponsive-worker-1@host",
|
||||
)
|
||||
ap_scan2, _ = self._create_executing_scan(
|
||||
tenant2, provider2, worker="dead-worker-2@host"
|
||||
tenant2,
|
||||
provider2,
|
||||
updated_at=updated_at,
|
||||
worker="unresponsive-worker-2@host",
|
||||
)
|
||||
mock_ping.return_value = (
|
||||
set(),
|
||||
{"unresponsive-worker-1@host", "unresponsive-worker-2@host"},
|
||||
)
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 2
|
||||
assert mock_revoke.call_count == 2
|
||||
assert mock_recover.call_count == 2
|
||||
|
||||
ap_scan1.refresh_from_db()
|
||||
@@ -2947,10 +3408,12 @@ class TestCleanupStaleAttackPathsScans:
|
||||
"tasks.jobs.attack_paths.cleanup.rls_transaction",
|
||||
new=lambda *args, **kwargs: nullcontext(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._is_worker_alive", return_value=False)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_recovers_graph_data_ready_for_stale_scan(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -2963,11 +3426,16 @@ class TestCleanupStaleAttackPathsScans:
|
||||
provider = aws_provider
|
||||
|
||||
ap_scan, _ = self._create_executing_scan(
|
||||
tenant, provider, worker="dead-worker@host"
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(minutes=31),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"unresponsive-worker@host"})
|
||||
|
||||
cleanup_stale_attack_paths_scans()
|
||||
|
||||
mock_revoke.assert_called_once()
|
||||
mock_recover.assert_called_once()
|
||||
recovered_scan = mock_recover.call_args[0][0]
|
||||
assert recovered_scan.id == ap_scan.id
|
||||
@@ -3013,10 +3481,57 @@ class TestCleanupStaleAttackPathsScans:
|
||||
"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(
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_preserves_scans_without_a_started_at_timestamp(
|
||||
self,
|
||||
mock_alive,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
responsive_scan, _ = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
worker="responsive-worker@host",
|
||||
)
|
||||
AttackPathsScan.objects.filter(id=responsive_scan.id).update(started_at=None)
|
||||
no_worker_scan = AttackPathsScan.objects.create(
|
||||
tenant_id=tenants_fixture[0].id,
|
||||
provider=aws_provider,
|
||||
state=StateChoices.EXECUTING,
|
||||
started_at=None,
|
||||
)
|
||||
mock_ping.return_value = ({"responsive-worker@host"}, set())
|
||||
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
responsive_scan.refresh_from_db()
|
||||
no_worker_scan.refresh_from_db()
|
||||
assert responsive_scan.state == StateChoices.EXECUTING
|
||||
assert no_worker_scan.state == StateChoices.EXECUTING
|
||||
mock_revoke.assert_not_called()
|
||||
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(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
@patch("tasks.jobs.attack_paths.cleanup._ping_workers")
|
||||
def test_shared_worker_is_collected_only_once(
|
||||
self,
|
||||
mock_ping,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
@@ -3028,15 +3543,114 @@ class TestCleanupStaleAttackPathsScans:
|
||||
tenant = tenants_fixture[0]
|
||||
provider = aws_provider
|
||||
|
||||
# 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")
|
||||
updated_at = datetime.now(tz=UTC) - timedelta(minutes=31)
|
||||
self._create_executing_scan(
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=updated_at,
|
||||
worker="shared-worker@host",
|
||||
)
|
||||
self._create_executing_scan(
|
||||
tenant,
|
||||
provider,
|
||||
updated_at=updated_at,
|
||||
worker="shared-worker@host",
|
||||
)
|
||||
mock_ping.return_value = (set(), {"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")
|
||||
assert mock_revoke.call_count == 2
|
||||
mock_ping.assert_called_once_with({"shared-worker@host"})
|
||||
|
||||
@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_locked_recheck_preserves_scan_with_new_activity(
|
||||
self,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
ap_scan, _ = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(minutes=31),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
|
||||
def record_activity(_workers):
|
||||
AttackPathsScan.objects.filter(id=ap_scan.id).update(
|
||||
updated_at=datetime.now(tz=UTC)
|
||||
)
|
||||
return set(), {"unresponsive-worker@host"}
|
||||
|
||||
with patch(
|
||||
"tasks.jobs.attack_paths.cleanup._ping_workers",
|
||||
side_effect=record_activity,
|
||||
):
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
ap_scan.refresh_from_db()
|
||||
assert ap_scan.state == StateChoices.EXECUTING
|
||||
mock_revoke.assert_not_called()
|
||||
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(),
|
||||
)
|
||||
@patch("tasks.jobs.attack_paths.cleanup._revoke_task")
|
||||
def test_locked_recheck_preserves_scan_that_changed_state(
|
||||
self,
|
||||
mock_revoke,
|
||||
mock_drop_db,
|
||||
mock_recover,
|
||||
tenants_fixture,
|
||||
aws_provider,
|
||||
scans_fixture,
|
||||
):
|
||||
from tasks.jobs.attack_paths.cleanup import cleanup_stale_attack_paths_scans
|
||||
|
||||
ap_scan, _ = self._create_executing_scan(
|
||||
tenants_fixture[0],
|
||||
aws_provider,
|
||||
updated_at=datetime.now(tz=UTC) - timedelta(minutes=31),
|
||||
worker="unresponsive-worker@host",
|
||||
)
|
||||
|
||||
def complete_scan(_workers):
|
||||
AttackPathsScan.objects.filter(id=ap_scan.id).update(
|
||||
state=StateChoices.COMPLETED
|
||||
)
|
||||
return set(), {"unresponsive-worker@host"}
|
||||
|
||||
with patch(
|
||||
"tasks.jobs.attack_paths.cleanup._ping_workers",
|
||||
side_effect=complete_scan,
|
||||
):
|
||||
result = cleanup_stale_attack_paths_scans()
|
||||
|
||||
assert result["cleaned_up_count"] == 0
|
||||
ap_scan.refresh_from_db()
|
||||
assert ap_scan.state == StateChoices.COMPLETED
|
||||
mock_revoke.assert_not_called()
|
||||
mock_drop_db.assert_not_called()
|
||||
mock_recover.assert_not_called()
|
||||
|
||||
# `SCHEDULED` state cleanup
|
||||
def _create_scheduled_scan(
|
||||
@@ -3123,7 +3737,7 @@ class TestCleanupStaleAttackPathsScans:
|
||||
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"
|
||||
"global_error": "Scan never started - cleaned up by periodic task"
|
||||
}
|
||||
|
||||
# SCHEDULED revoke must NOT terminate a running worker
|
||||
|
||||
@@ -7,6 +7,7 @@ from celery import states
|
||||
from django.test import override_settings
|
||||
from django_celery_results.models import TaskResult
|
||||
from tasks.jobs.orphan_recovery import (
|
||||
_SKIP_RECOVERY,
|
||||
_decode_celery_field,
|
||||
_reconcile_task_results,
|
||||
_recovery_attempt_count,
|
||||
@@ -14,6 +15,7 @@ from tasks.jobs.orphan_recovery import (
|
||||
is_worker_alive,
|
||||
reconcile_orphans,
|
||||
reenqueueable_tasks,
|
||||
revoke_task,
|
||||
)
|
||||
|
||||
|
||||
@@ -180,10 +182,18 @@ class TestReconcileTaskResults:
|
||||
assert tr.task_id in result["failed"]
|
||||
mock_count.assert_not_called()
|
||||
|
||||
def test_scan_task_is_skipped_entirely(self, tenants_fixture):
|
||||
@pytest.mark.parametrize(
|
||||
"task_name",
|
||||
[
|
||||
"scan-perform",
|
||||
"attack-paths-scan-perform",
|
||||
"attack-paths-cleanup-stale-scans",
|
||||
],
|
||||
)
|
||||
def test_scan_task_is_skipped_entirely(self, tenants_fixture, task_name):
|
||||
"""Scan tasks are excluded from recovery: the watchdog never touches them."""
|
||||
tr = _orphan_result(
|
||||
name="scan-perform",
|
||||
name=task_name,
|
||||
kwargs={
|
||||
"tenant_id": str(tenants_fixture[0].id),
|
||||
"scan_id": str(uuid4()),
|
||||
@@ -339,6 +349,15 @@ class TestOrphanRecoveryHelpers:
|
||||
):
|
||||
assert is_worker_alive("w@h") is False
|
||||
|
||||
def test_revoke_task_terminates_with_sigterm_by_default(self):
|
||||
task_result = MagicMock(task_id="task-id")
|
||||
with patch(
|
||||
"tasks.jobs.orphan_recovery.current_app.control.revoke"
|
||||
) as mock_revoke:
|
||||
revoke_task(task_result)
|
||||
|
||||
mock_revoke.assert_called_once_with("task-id", terminate=True, signal="SIGTERM")
|
||||
|
||||
def test_recovery_attempt_count_increments(self):
|
||||
# Unique signature so the Valkey counter starts fresh for this test.
|
||||
kwargs_repr = repr({"probe": str(uuid4())})
|
||||
@@ -350,6 +369,12 @@ class TestOrphanRecoveryHelpers:
|
||||
|
||||
|
||||
class TestRecoveryFeatureFlags:
|
||||
def test_attack_paths_tasks_are_excluded_from_generic_recovery(self):
|
||||
assert {
|
||||
"attack-paths-scan-perform",
|
||||
"attack-paths-cleanup-stale-scans",
|
||||
} <= _SKIP_RECOVERY
|
||||
|
||||
def test_all_groups_enabled_by_default(self):
|
||||
tasks = reenqueueable_tasks()
|
||||
assert "scan-summary" in tasks
|
||||
@@ -374,33 +399,23 @@ class TestRecoveryFeatureFlags:
|
||||
class TestRecoveryMasterFlag:
|
||||
@override_settings(TASK_RECOVERY_ENABLED=False)
|
||||
def test_master_flag_disables_task_recovery(self):
|
||||
with (
|
||||
patch(
|
||||
"tasks.jobs.orphan_recovery._reconcile_task_results"
|
||||
) as mock_reconcile,
|
||||
patch(
|
||||
"tasks.jobs.attack_paths.cleanup.cleanup_stale_attack_paths_scans",
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
with patch(
|
||||
"tasks.jobs.orphan_recovery._reconcile_task_results"
|
||||
) as mock_reconcile:
|
||||
result = reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False)
|
||||
|
||||
mock_reconcile.assert_not_called()
|
||||
assert result["acquired"] is True
|
||||
assert result["enabled"] is False
|
||||
assert "attack_paths" not in result
|
||||
|
||||
@override_settings(TASK_RECOVERY_ENABLED=True)
|
||||
def test_master_flag_enabled_runs_task_recovery(self):
|
||||
with (
|
||||
patch(
|
||||
"tasks.jobs.orphan_recovery._reconcile_task_results",
|
||||
return_value={"recovered": [], "failed": [], "skipped": []},
|
||||
) as mock_reconcile,
|
||||
patch(
|
||||
"tasks.jobs.attack_paths.cleanup.cleanup_stale_attack_paths_scans",
|
||||
return_value={},
|
||||
),
|
||||
):
|
||||
reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False)
|
||||
with patch(
|
||||
"tasks.jobs.orphan_recovery._reconcile_task_results",
|
||||
return_value={"recovered": [], "failed": [], "skipped": []},
|
||||
) as mock_reconcile:
|
||||
result = reconcile_orphans(grace_minutes=2, max_attempts=3, dry_run=False)
|
||||
|
||||
mock_reconcile.assert_called_once()
|
||||
assert "attack_paths" not in result
|
||||
|
||||
Reference in New Issue
Block a user