ben.carrasco/prowler
1
0
Fork 0
mirror of https://github.com/prowler-cloud/prowler.git synced 2026-07-23 20:42:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(api): resolve attack paths scan reliability issues (#12019)

Browse Source
This commit is contained in:
Josema Camacho
2026-07-16 18:15:42 +02:00
committed by GitHub
parent bc3c922177
commit bfc6b9e577
22 changed files with 726 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/src/backend/api/attack_paths/database.py
+5
View File
@@ -27,6 +27,7 @@ from django.conf import (
MAX_CUSTOM_QUERY_NODES = env.int("ATTACK_PATHS_MAX_CUSTOM_QUERY_NODES", default=250)
TEMP_DB_PREFIX = "db-tmp-scan-"
DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound"
# Exceptions
@@ -44,6 +45,10 @@ class GraphDatabaseQueryException(Exception):
return self.message
class NeptuneWriteRetryExhaustedException(GraphDatabaseQueryException):
pass
class WriteQueryNotAllowedException(GraphDatabaseQueryException):
pass
api/src/backend/api/attack_paths/retryable_session.py
+54 -8
View File
@@ -10,6 +10,28 @@ import neo4j.exceptions
logger = logging.getLogger(__name__)
class RetryExhaustedError(Exception):
def __init__(
self,
*,
retry_context: str,
method_name: str,
attempts: int,
elapsed_seconds: float,
last_error: Exception,
) -> None:
self.retry_context = retry_context
self.method_name = method_name
self.attempts = attempts
self.elapsed_seconds = elapsed_seconds
self.last_error = last_error
last_message = getattr(last_error, "message", None) or str(last_error)
super().__init__(
f"{retry_context} {method_name} failed after {attempts} attempts over "
f"{elapsed_seconds:.3f}s. Last error: {last_message}"
)
class RetryableSession:
"""Wrapper around ``neo4j.Session`` with a refreshable retry policy."""
@@ -19,11 +41,13 @@ class RetryableSession:
max_retries: int,
retry_if: Callable[[Exception], bool] | None = None,
initial_retry_delay_seconds: float = 0,
retry_context: str | None = None,
) -> None:
self._session_factory = session_factory
self._max_retries = max(0, max_retries)
self._retry_if = retry_if
self._initial_retry_delay_seconds = max(0.0, initial_retry_delay_seconds)
self._retry_context = retry_context
self._session = self._session_factory()
def close(self) -> None:
@@ -54,6 +78,7 @@ class RetryableSession:
def _call_with_retry(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
attempt = 0
last_exc: Exception | None = None
started_at = time.monotonic()
while attempt <= self._max_retries:
try:
@@ -68,17 +93,38 @@ class RetryableSession:
attempt += 1
if attempt > self._max_retries:
if self._retry_context is not None:
raise RetryExhaustedError(
retry_context=self._retry_context,
method_name=method_name,
attempts=attempt,
elapsed_seconds=time.monotonic() - started_at,
last_error=exc,
) from exc
raise
delay = self._retry_delay(attempt)
logger.warning(
"Graph session %s failed with %s; retry %s/%s in %.3fs",
method_name,
type(exc).__name__,
attempt,
self._max_retries,
delay,
)
if self._retry_context is not None:
error_message = getattr(exc, "message", None) or str(exc)
logger.warning(
"%s %s failed with %s: %s; retry %s/%s in %.3fs",
self._retry_context,
method_name,
type(exc).__name__,
error_message,
attempt,
self._max_retries,
delay,
)
else:
logger.warning(
"Graph session %s failed with %s; retry %s/%s in %.3fs",
method_name,
type(exc).__name__,
attempt,
self._max_retries,
delay,
)
self._refresh_session()
if delay:
time.sleep(delay)
api/src/backend/api/attack_paths/sink/base.py
+2
View File
@@ -15,6 +15,8 @@ class SinkDatabase(Protocol):
has a single graph, and isolation is label-based).
"""
sync_batch_size: int
def init(self) -> None: ...
def close(self) -> None: ...
api/src/backend/api/attack_paths/sink/neo4j.py
+5 -3
View File
@@ -54,6 +54,8 @@ DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound"
class Neo4jSink(SinkDatabase):
"""Neo4j-backed sink. Multi-database cluster; tenant isolation is physical."""
sync_batch_size = env.int("ATTACK_PATHS_NEO4J_SYNC_BATCH_SIZE", default=1000)
def __init__(self) -> None:
self._driver: neo4j.Driver | None = None
self._lock = threading.Lock()
@@ -203,7 +205,7 @@ class Neo4jSink(SinkDatabase):
"""
from api.attack_paths.database import GraphDatabaseQueryException
from tasks.jobs.attack_paths.config import (
BATCH_SIZE,
GRAPH_MUTATION_BATCH_SIZE,
PROVIDER_RESOURCE_LABEL,
get_provider_label,
)
@@ -251,7 +253,7 @@ class Neo4jSink(SinkDatabase):
total_key="rels",
deleted_key="deleted_rels",
initial_total=deleted_relationships,
batch_size=BATCH_SIZE,
batch_size=GRAPH_MUTATION_BATCH_SIZE,
drop_t0=drop_t0,
)
relationship_batches += phase_batches
@@ -270,7 +272,7 @@ class Neo4jSink(SinkDatabase):
total_key="nodes",
deleted_key="deleted_nodes",
initial_total=0,
batch_size=BATCH_SIZE,
batch_size=GRAPH_MUTATION_BATCH_SIZE,
drop_t0=drop_t0,
)
api/src/backend/api/attack_paths/sink/neptune.py
+15 -4
View File
@@ -25,7 +25,7 @@ from urllib.parse import urlsplit
import neo4j
import neo4j.exceptions
from api.attack_paths.retryable_session import RetryableSession
from api.attack_paths.retryable_session import RetryableSession, RetryExhaustedError
from api.attack_paths.sink.base import SinkDatabase
from api.attack_paths.sink.drop import (
NODE_DELETE_QUERY_TEMPLATE,
@@ -85,6 +85,8 @@ def _is_retryable_write_error(exc: Exception) -> bool:
class NeptuneSink(SinkDatabase):
"""Neptune-backed sink. Single database; isolation is label-based."""
sync_batch_size = env.int("ATTACK_PATHS_NEPTUNE_SYNC_BATCH_SIZE", default=500)
def __init__(self) -> None:
self._writer: neo4j.Driver | None = None
self._reader: neo4j.Driver | None = None
@@ -206,6 +208,7 @@ class NeptuneSink(SinkDatabase):
from api.attack_paths.database import (
ClientStatementException,
GraphDatabaseQueryException,
NeptuneWriteRetryExhaustedException,
WriteQueryNotAllowedException,
)
@@ -227,9 +230,17 @@ class NeptuneSink(SinkDatabase):
initial_retry_delay_seconds=(
NEPTUNE_WRITE_RETRY_DELAY_SECONDS if is_write_session else 0
),
retry_context="Neptune write" if is_write_session else None,
)
yield session_wrapper
except RetryExhaustedError as exc:
last_error = exc.last_error
raise NeptuneWriteRetryExhaustedException(
message=str(exc),
code=getattr(last_error, "code", None),
) from last_error
except neo4j.exceptions.Neo4jError as exc:
if (
default_access_mode == neo4j.READ_ACCESS
@@ -291,7 +302,7 @@ class NeptuneSink(SinkDatabase):
graph's branching factor.
"""
from tasks.jobs.attack_paths.config import (
BATCH_SIZE,
GRAPH_MUTATION_BATCH_SIZE,
PROVIDER_RESOURCE_LABEL,
get_provider_label,
)
@@ -330,7 +341,7 @@ class NeptuneSink(SinkDatabase):
total_key="rels",
deleted_key="deleted_rels",
initial_total=deleted_relationships,
batch_size=BATCH_SIZE,
batch_size=GRAPH_MUTATION_BATCH_SIZE,
drop_t0=drop_t0,
)
relationship_batches += phase_batches
@@ -349,7 +360,7 @@ class NeptuneSink(SinkDatabase):
total_key="nodes",
deleted_key="deleted_nodes",
initial_total=0,
batch_size=BATCH_SIZE,
batch_size=GRAPH_MUTATION_BATCH_SIZE,
drop_t0=drop_t0,
)
api/src/backend/api/decorators.py
+22 -7
View File
@@ -1,12 +1,13 @@
import uuid
from functools import wraps
from api.attack_paths.database import GraphDatabaseQueryException
from api.db_router import READ_REPLICA_ALIAS
from api.db_utils import POSTGRES_TENANT_VAR, SET_CONFIG_QUERY, rls_transaction
from api.exceptions import ProviderDeletedException
from api.models import Provider, Scan
from api.models import Membership, Provider, Scan, Tenant
from django.core.exceptions import ObjectDoesNotExist
from django.db import DatabaseError, connection, transaction
from django.db import DEFAULT_DB_ALIAS, DatabaseError, connection, transaction
from rest_framework_json_api.serializers import ValidationError
@@ -75,9 +76,11 @@ def handle_provider_deletion(func):
"""
Decorator that raises `ProviderDeletedException` if provider was deleted during execution.
Catches `ObjectDoesNotExist` and `DatabaseError` (including `IntegrityError`), checks if
provider still exists, and raises `ProviderDeletedException` if not. Otherwise,
re-raises original exception.
Catches `ObjectDoesNotExist`, `DatabaseError` (including `IntegrityError`), and
`GraphDatabaseQueryException`, checks if provider still exists, and raises
`ProviderDeletedException` if not. Graph database errors also check whether the
tenant still exists and has memberships. Otherwise, re-raises the original
exception.
Requires `tenant_id` and `provider_id` in kwargs.
@@ -92,11 +95,16 @@ def handle_provider_deletion(func):
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except (ObjectDoesNotExist, DatabaseError):
except (ObjectDoesNotExist, DatabaseError, GraphDatabaseQueryException) as exc:
tenant_id = kwargs.get("tenant_id")
provider_id = kwargs.get("provider_id")
database_alias = (
DEFAULT_DB_ALIAS
if isinstance(exc, GraphDatabaseQueryException)
else READ_REPLICA_ALIAS
)
with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS):
with rls_transaction(tenant_id, using=database_alias):
if provider_id is None:
scan_id = kwargs.get("scan_id")
if scan_id is None:
@@ -113,6 +121,13 @@ def handle_provider_deletion(func):
raise ProviderDeletedException(
f"Provider '{provider_id}' was deleted during the scan"
) from None
if isinstance(exc, GraphDatabaseQueryException) and (
not Tenant.objects.filter(pk=tenant_id).exists()
or not Membership.objects.filter(tenant_id=tenant_id).exists()
):
raise ProviderDeletedException(
f"Tenant '{tenant_id}' was deleted during the scan"
) from None
raise
return wrapper
api/src/backend/api/tests/test_decorators.py
+102 -1
View File
@@ -2,11 +2,12 @@ import uuid
from unittest.mock import call, patch
import pytest
from api.attack_paths.database import GraphDatabaseQueryException
from api.db_utils import POSTGRES_TENANT_VAR, SET_CONFIG_QUERY
from api.decorators import handle_provider_deletion, set_tenant
from api.exceptions import ProviderDeletedException
from django.core.exceptions import ObjectDoesNotExist
from django.db import DatabaseError, IntegrityError
from django.db import DEFAULT_DB_ALIAS, DatabaseError, IntegrityError
@pytest.mark.django_db
@@ -204,6 +205,106 @@ class TestHandleProviderDeletionDecorator:
with pytest.raises(DatabaseError):
task_func(tenant_id=str(tenant.id), provider_id=str(provider.id))
@patch("api.decorators.rls_transaction")
@patch("api.decorators.Provider.objects.filter")
def test_graph_database_error_provider_missing_or_soft_deleted(
self, mock_provider_filter, mock_rls, tenants_fixture
):
tenant = tenants_fixture[0]
provider_id = str(uuid.uuid4())
mock_rls.return_value.__enter__ = lambda s: None
mock_rls.return_value.__exit__ = lambda s, *args: None
mock_provider_filter.return_value.exists.return_value = False
@handle_provider_deletion
def task_func(**kwargs):
raise GraphDatabaseQueryException("Temporary database not found")
with pytest.raises(ProviderDeletedException):
task_func(tenant_id=str(tenant.id), provider_id=provider_id)
@patch("api.decorators.rls_transaction")
@patch("api.decorators.Tenant.objects.filter")
@patch("api.decorators.Provider.objects.filter")
def test_graph_database_error_tenant_missing(
self, mock_provider_filter, mock_tenant_filter, mock_rls, tenants_fixture
):
tenant = tenants_fixture[0]
provider_id = str(uuid.uuid4())
mock_rls.return_value.__enter__ = lambda s: None
mock_rls.return_value.__exit__ = lambda s, *args: None
mock_provider_filter.return_value.exists.return_value = True
mock_tenant_filter.return_value.exists.return_value = False
@handle_provider_deletion
def task_func(**kwargs):
raise GraphDatabaseQueryException("Temporary database not found")
with pytest.raises(ProviderDeletedException):
task_func(tenant_id=str(tenant.id), provider_id=provider_id)
@patch("api.decorators.rls_transaction")
@patch("api.decorators.Membership.objects.filter")
@patch("api.decorators.Tenant.objects.filter")
@patch("api.decorators.Provider.objects.filter")
def test_graph_database_error_tenant_without_memberships(
self,
mock_provider_filter,
mock_tenant_filter,
mock_membership_filter,
mock_rls,
tenants_fixture,
):
tenant = tenants_fixture[0]
provider_id = str(uuid.uuid4())
mock_rls.return_value.__enter__ = lambda s: None
mock_rls.return_value.__exit__ = lambda s, *args: None
mock_provider_filter.return_value.exists.return_value = True
mock_tenant_filter.return_value.exists.return_value = True
mock_membership_filter.return_value.exists.return_value = False
@handle_provider_deletion
def task_func(**kwargs):
raise GraphDatabaseQueryException("Temporary database not found")
with pytest.raises(ProviderDeletedException):
task_func(tenant_id=str(tenant.id), provider_id=provider_id)
@patch("api.decorators.rls_transaction")
@patch("api.decorators.Membership.objects.filter")
@patch("api.decorators.Tenant.objects.filter")
@patch("api.decorators.Provider.objects.filter")
def test_graph_database_error_active_provider_and_tenant_reraises(
self,
mock_provider_filter,
mock_tenant_filter,
mock_membership_filter,
mock_rls,
tenants_fixture,
):
tenant = tenants_fixture[0]
provider_id = str(uuid.uuid4())
graph_error = GraphDatabaseQueryException("Temporary database not found")
mock_rls.return_value.__enter__ = lambda s: None
mock_rls.return_value.__exit__ = lambda s, *args: None
mock_provider_filter.return_value.exists.return_value = True
mock_tenant_filter.return_value.exists.return_value = True
mock_membership_filter.return_value.exists.return_value = True
@handle_provider_deletion
def task_func(**kwargs):
raise graph_error
with pytest.raises(GraphDatabaseQueryException) as exc_info:
task_func(tenant_id=str(tenant.id), provider_id=provider_id)
assert exc_info.value is graph_error
mock_rls.assert_called_once_with(str(tenant.id), using=DEFAULT_DB_ALIAS)
def test_missing_provider_and_scan_raises_assertion(self, tenants_fixture):
"""Raises AssertionError when neither provider_id nor scan_id in kwargs."""
api/src/backend/api/tests/test_retryable_session.py
+81 -1
View File
@@ -1,7 +1,7 @@
from unittest.mock import MagicMock, patch
import pytest
from api.attack_paths.retryable_session import RetryableSession
from api.attack_paths.retryable_session import RetryableSession, RetryExhaustedError
from neo4j.exceptions import ServiceUnavailable
@@ -24,6 +24,7 @@ class TestRetryableSession:
max_retries=3,
retry_if=lambda exc: exc is retryable_error,
initial_retry_delay_seconds=2,
retry_context="Neptune write",
)
assert session.execute_write(work) == "success"
@@ -54,6 +55,7 @@ class TestRetryableSession:
max_retries=3,
retry_if=lambda _: False,
initial_retry_delay_seconds=2,
retry_context="Neptune write",
)
with pytest.raises(RuntimeError) as exc_info:
@@ -83,3 +85,81 @@ class TestRetryableSession:
driver_sessions[0].close.assert_called_once_with()
driver_sessions[1].close.assert_called_once_with()
driver_sessions[2].close.assert_not_called()
def test_retry_exhaustion_with_context_reports_attempts_and_elapsed_time(self):
error = RuntimeError("still retryable")
driver_sessions = [MagicMock() for _ in range(3)]
for driver_session in driver_sessions:
driver_session.execute_write.side_effect = error
session = RetryableSession(
session_factory=MagicMock(side_effect=driver_sessions),
max_retries=2,
retry_if=lambda _: True,
retry_context="Neptune write",
)
with (
patch(
"api.attack_paths.retryable_session.time.monotonic",
side_effect=[100.0, 127.1234],
),
pytest.raises(RetryExhaustedError) as exc_info,
):
session.execute_write(MagicMock())
assert exc_info.value.method_name == "execute_write"
assert exc_info.value.attempts == 3
assert exc_info.value.elapsed_seconds == pytest.approx(27.1234)
assert exc_info.value.last_error is error
assert exc_info.value.__cause__ is error
assert str(exc_info.value) == (
"Neptune write execute_write failed after 3 attempts over 27.123s. "
"Last error: still retryable"
)
def test_retry_exhaustion_with_zero_retries_reports_one_attempt(self):
error = ServiceUnavailable("still unavailable")
driver_session = MagicMock()
driver_session.execute_write.side_effect = error
session = RetryableSession(
session_factory=MagicMock(return_value=driver_session),
max_retries=0,
retry_context="Neptune write",
)
with pytest.raises(RetryExhaustedError) as exc_info:
session.execute_write(MagicMock())
assert exc_info.value.attempts == 1
@patch("api.attack_paths.retryable_session.time.sleep")
@patch("api.attack_paths.retryable_session.random.uniform", return_value=3.0)
def test_contextual_retry_warning_includes_original_error(
self, _mock_uniform, _mock_sleep
):
error = RuntimeError("retryable detail")
first_session = MagicMock()
first_session.execute_write.side_effect = error
second_session = MagicMock()
second_session.execute_write.return_value = "success"
session = RetryableSession(
session_factory=MagicMock(side_effect=[first_session, second_session]),
max_retries=1,
retry_if=lambda _: True,
initial_retry_delay_seconds=2,
retry_context="Neptune write",
)
with patch("api.attack_paths.retryable_session.logger.warning") as mock_warning:
assert session.execute_write(MagicMock()) == "success"
mock_warning.assert_called_once_with(
"%s %s failed with %s: %s; retry %s/%s in %.3fs",
"Neptune write",
"execute_write",
"RuntimeError",
"retryable detail",
1,
1,
3.0,
)
api/src/backend/api/tests/test_sentry.py
+40
View File
@@ -1,6 +1,7 @@
import logging
from unittest.mock import MagicMock, patch
import pytest
from config.settings import sentry as sentry_settings
from config.settings.sentry import before_send
@@ -82,6 +83,45 @@ def test_before_send_passes_through_non_ignored_log():
assert result == event
def test_before_send_ignores_cartography_missing_temporary_database_log():
log_record = _make_log_record(
msg="Cartography job failed with %s for database %s",
name="cartography.graph.job",
args=(
"Neo.ClientError.Database.DatabaseNotFound",
"db-tmp-scan-12345678",
),
)
event = MagicMock()
assert before_send(event, {"log_record": log_record}) is None
@pytest.mark.parametrize(
("logger_name", "message"),
[
(
"cartography.graph.job.worker",
"Neo.ClientError.Database.DatabaseNotFound for db-tmp-scan-12345678",
),
(
"cartography.graph.job",
"DatabaseNotFound for db-tmp-scan-12345678",
),
(
"cartography.graph.job",
"Neo.ClientError.Database.DatabaseNotFound for db-tenant-12345678",
),
],
)
def test_before_send_passes_through_similar_cartography_logs(logger_name, message):
log_record = _make_log_record(msg=message, name=logger_name)
event = MagicMock()
assert before_send(event, {"log_record": log_record}) is event
def test_before_send_passes_through_non_ignored_exception():
"""Test that before_send passes through exceptions that don't contain ignored exceptions."""
exc_info = (Exception, Exception("Some other error message"), None)
api/src/backend/api/tests/test_sink.py
+56 -1
View File
@@ -11,7 +11,11 @@ from unittest.mock import MagicMock, patch
import neo4j
import pytest
from api.attack_paths import sink as sink_module
from api.attack_paths.database import GraphDatabaseQueryException
from api.attack_paths.database import (
GraphDatabaseQueryException,
NeptuneWriteRetryExhaustedException,
)
from api.attack_paths.retryable_session import RetryExhaustedError
from api.attack_paths.sink import factory
from api.attack_paths.sink.neo4j import DATABASE_NOT_FOUND_CODE, Neo4jSink
from api.attack_paths.sink.neptune import (
@@ -123,6 +127,14 @@ class TestSinkFactory:
assert mock_driver.call_count == 1
def test_neo4j_sync_batch_size_defaults_to_1000():
assert Neo4jSink.sync_batch_size == 1000
def test_neptune_sync_batch_size_defaults_to_500():
assert NeptuneSink.sync_batch_size == 500
class TestGetBackendForScan:
"""``get_backend_for_scan`` routes by the row's recorded sink backend."""
@@ -372,6 +384,7 @@ class TestNeptuneRetryPolicy:
assert (
kwargs["initial_retry_delay_seconds"] == NEPTUNE_WRITE_RETRY_DELAY_SECONDS
)
assert kwargs["retry_context"] == "Neptune write"
@patch("api.attack_paths.sink.neptune.RetryableSession")
def test_reader_session_does_not_enable_write_retry_policy(self, retryable_session):
@@ -384,6 +397,48 @@ class TestNeptuneRetryPolicy:
kwargs = retryable_session.call_args.kwargs
assert kwargs["retry_if"] is None
assert kwargs["initial_retry_delay_seconds"] == 0
assert kwargs["retry_context"] is None
def test_writer_retry_exhaustion_preserves_neptune_error_details(self):
message = (
"Unexpected server exception 'Operation failed due to conflicting "
"concurrent operations (please retry), 0 transactions are currently "
"rolling back.'"
)
error = neo4j.exceptions.Neo4jError._hydrate_neo4j(
code="BoltProtocol.unexpectedException",
message=message,
)
retry_error = RetryExhaustedError(
retry_context="Neptune write",
method_name="execute_write",
attempts=4,
elapsed_seconds=27.1234,
last_error=error,
)
sink = NeptuneSink()
driver = MagicMock()
retryable_session = MagicMock()
retryable_session.execute_write.side_effect = retry_error
with (
patch.object(sink, "_get_writer", return_value=driver),
patch(
"api.attack_paths.sink.neptune.RetryableSession",
return_value=retryable_session,
),
pytest.raises(NeptuneWriteRetryExhaustedException) as exc_info,
):
with sink.get_session() as session:
session.execute_write(MagicMock())
assert exc_info.value.code == "BoltProtocol.unexpectedException"
assert str(exc_info.value) == (
"BoltProtocol.unexpectedException: Neptune write execute_write failed "
"after 4 attempts over 27.123s. Last error: "
f"{message}"
)
assert exc_info.value.__cause__ is error
class TestNeptuneSinkDropSubgraph:
api/src/backend/config/django/base.py
+2 -2
View File
@@ -312,8 +312,8 @@ 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
"ATTACK_PATHS_SCAN_STALE_THRESHOLD_MINUTES", 960
) # 16h
# Selects where the persistent attack-paths graph is stored. The scan
# temporary database is always Neo4j; only the sink is configurable.
api/src/backend/config/settings/sentry.py
+7
View File
@@ -91,6 +91,13 @@ def before_send(event, hint):
log_msg = log_record.getMessage()
log_lvl = log_record.levelno
if (
getattr(log_record, "name", "") == "cartography.graph.job"
and "Neo.ClientError.Database.DatabaseNotFound" in log_msg
and "db-tmp-scan-" in log_msg
):
return None
# The Neo4j driver logs transient connection errors (defunct
# connections, resets) at ERROR level via the `neo4j.io` logger.
# `RetryableSession` handles these with retries. If all retries
api/src/backend/tasks/jobs/attack_paths/aws.py
+8
View File
@@ -8,6 +8,8 @@ import aioboto3
import boto3
import botocore
import neo4j
import neo4j.exceptions
from api.attack_paths.database import DATABASE_NOT_FOUND_CODE
from api.models import (
AttackPathsScan as ProwlerAPIAttackPathsScan,
)
@@ -347,6 +349,12 @@ def sync_aws_account(
)
except Exception as e:
if (
isinstance(e, neo4j.exceptions.Neo4jError)
and e.code == DATABASE_NOT_FOUND_CODE
):
raise
logger.info(
f"Synced function {func_name} for AWS account {prowler_api_provider.uid} in {time.perf_counter() - func_t0:.3f}s (FAILED)"
)
api/src/backend/tasks/jobs/attack_paths/config.py
+2 -5
View File
@@ -10,13 +10,10 @@ NormalizedList = _provider_config.NormalizedList
PROVIDER_CONFIGS = _provider_config.PROVIDER_CONFIGS
ProviderConfig = _provider_config.ProviderConfig
# Batch size for Neo4j write operations (resource labeling, cleanup)
BATCH_SIZE = env.int("ATTACK_PATHS_BATCH_SIZE", 1000)
# Batch size for graph mutation operations (resource labeling and subgraph deletion)
GRAPH_MUTATION_BATCH_SIZE = env.int("ATTACK_PATHS_GRAPH_MUTATION_BATCH_SIZE", 1000)
# Batch size for Postgres findings fetch (keyset pagination page size)
FINDINGS_BATCH_SIZE = env.int("ATTACK_PATHS_FINDINGS_BATCH_SIZE", 1000)
# Batch size for temp-to-tenant graph sync (nodes and relationships per cursor page)
SYNC_BATCH_SIZE = env.int("ATTACK_PATHS_SYNC_BATCH_SIZE", 1000)
# Neo4j internal labels (Prowler-specific, not provider-specific)
# - `Internet`: Singleton node representing external internet access for exposed-resource queries
# - `ProwlerFinding`: Label for finding nodes created by Prowler and linked to cloud resources
api/src/backend/tasks/jobs/attack_paths/findings.py
+2 -2
View File
@@ -21,8 +21,8 @@ from cartography.config import Config as CartographyConfig
from celery.utils.log import get_task_logger
from prowler.config import config as ProwlerConfig
from tasks.jobs.attack_paths.config import (
BATCH_SIZE,
FINDINGS_BATCH_SIZE,
GRAPH_MUTATION_BATCH_SIZE,
get_node_uid_field,
get_provider_resource_label,
get_root_node_label,
@@ -135,7 +135,7 @@ def add_resource_label(
while labeled_count > 0:
result = neo4j_session.run(
query,
{"provider_uid": provider_uid, "batch_size": BATCH_SIZE},
{"provider_uid": provider_uid, "batch_size": GRAPH_MUTATION_BATCH_SIZE},
)
labeled_count = result.single().get("labeled_count", 0)
total_labeled += labeled_count
api/src/backend/tasks/jobs/attack_paths/scan.py
+31 -12
View File
@@ -372,7 +372,19 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
except Exception as e:
exception_message = utils.stringify_exception(e, "Attack Paths scan failed")
logger.exception(exception_message)
temporary_database_missing = (
isinstance(e, graph_database.GraphDatabaseQueryException)
and e.code == graph_database.DATABASE_NOT_FOUND_CODE
and tmp_database_name in str(e)
)
if temporary_database_missing:
logger.warning(exception_message)
else:
logger.exception(exception_message)
cleanup_log_level = (
logging.WARNING if temporary_database_missing else logging.ERROR
)
cleanup_exc_info = not temporary_database_missing
ingestion_exceptions["global_error"] = exception_message
# Recover `graph_data_ready` based on how far the swap got
@@ -387,19 +399,24 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
)
except Exception:
logger.error(
f"Failed to recover `graph_data_ready` for provider {attack_paths_scan.provider_id}",
exc_info=True,
logger.log(
cleanup_log_level,
"Failed to recover `graph_data_ready` for provider "
f"{attack_paths_scan.provider_id}",
exc_info=cleanup_exc_info,
)
# Dropping the temporary database if it still exists
try:
graph_database.drop_database(tmp_cartography_config.neo4j_database)
except Exception as e:
logger.error(
f"Failed to drop temporary Neo4j database `{tmp_cartography_config.neo4j_database}` during cleanup: {e}",
exc_info=True,
except Exception as cleanup_error:
logger.log(
cleanup_log_level,
"Failed to drop temporary Neo4j database "
f"`{tmp_cartography_config.neo4j_database}` during cleanup: "
f"{cleanup_error}",
exc_info=cleanup_exc_info,
)
# Set Attack Paths scan state to FAILED
@@ -407,10 +424,12 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
db_utils.finish_attack_paths_scan(
attack_paths_scan, StateChoices.FAILED, ingestion_exceptions
)
except Exception as e:
logger.error(
f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` (row may have been deleted): {e}",
exc_info=True,
except Exception as cleanup_error:
logger.log(
cleanup_log_level,
f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` "
f"(row may have been deleted): {cleanup_error}",
exc_info=cleanup_exc_info,
)
raise
api/src/backend/tasks/jobs/attack_paths/sync.py
+9 -9
View File
@@ -30,7 +30,6 @@ from tasks.jobs.attack_paths.config import (
PROVIDER_CONFIGS,
PROVIDER_ISOLATION_PROPERTIES,
PROVIDER_RESOURCE_LABEL,
SYNC_BATCH_SIZE,
NormalizedList,
get_provider_label,
get_tenant_label,
@@ -116,6 +115,7 @@ def sync_nodes(
Source and target sessions are opened sequentially per batch to avoid
holding two Bolt connections simultaneously for the entire sync duration.
"""
batch_size = sink.sync_batch_size
t0 = time.perf_counter()
last_id = -1
parents_synced = 0
@@ -137,7 +137,7 @@ def sync_nodes(
with graph_database.get_session(source_database) as source_session:
result = source_session.run(
NODE_FETCH_QUERY,
{"last_id": last_id, "batch_size": SYNC_BATCH_SIZE},
{"last_id": last_id, "batch_size": batch_size},
)
for record in result:
batch_count += 1
@@ -156,17 +156,17 @@ def sync_nodes(
for labels, batch in parent_groups.items():
rendered_labels = _render_labels(labels, extra_labels)
for sink_batch in _iter_sink_batches(batch):
for sink_batch in _iter_sink_batches(batch, batch_size):
sink.write_nodes(target_database, rendered_labels, sink_batch)
for child_label, batch in child_groups.items():
rendered_labels = _render_labels((child_label,), extra_labels)
for sink_batch in _iter_sink_batches(batch):
for sink_batch in _iter_sink_batches(batch, batch_size):
sink.write_nodes(target_database, rendered_labels, sink_batch)
children_synced += len(batch)
for rel_type, batch in rel_groups.items():
for sink_batch in _iter_sink_batches(batch):
for sink_batch in _iter_sink_batches(batch, batch_size):
sink.write_relationships(
target_database, rel_type, provider_id, sink_batch
)
@@ -205,6 +205,7 @@ def sync_relationships(
Source and target sessions are opened sequentially per batch to avoid
holding two Bolt connections simultaneously for the entire sync duration.
"""
batch_size = sink.sync_batch_size
t0 = time.perf_counter()
last_id = -1
total_synced = 0
@@ -217,7 +218,7 @@ def sync_relationships(
with graph_database.get_session(source_database) as source_session:
result = source_session.run(
RELATIONSHIPS_FETCH_QUERY,
{"last_id": last_id, "batch_size": SYNC_BATCH_SIZE},
{"last_id": last_id, "batch_size": batch_size},
)
for record in result:
batch_count += 1
@@ -229,7 +230,7 @@ def sync_relationships(
break
for rel_type, batch in grouped.items():
for sink_batch in _iter_sink_batches(batch):
for sink_batch in _iter_sink_batches(batch, batch_size):
sink.write_relationships(
target_database, rel_type, provider_id, sink_batch
)
@@ -247,10 +248,9 @@ def sync_relationships(
def _iter_sink_batches(
rows: list[dict[str, Any]],
batch_size: int | None = None,
batch_size: int,
) -> Iterator[list[dict[str, Any]]]:
"""Yield final sink write batches after source rows have been transformed."""
batch_size = SYNC_BATCH_SIZE if batch_size is None else batch_size
if batch_size <= 0:
raise ValueError("Sink batch size must be greater than zero")
api/src/backend/tasks/tasks.py
+8 -1
View File
@@ -11,6 +11,7 @@ from api.compliance import (
from api.db_router import READ_REPLICA_ALIAS
from api.db_utils import delete_related_daily_task, rls_transaction
from api.decorators import handle_provider_deletion, set_tenant
from api.exceptions import ProviderDeletedException
from api.models import (
Finding,
Integration,
@@ -666,7 +667,13 @@ class AttackPathsScanRLSTask(RLSTask):
scan_id = kwargs.get("scan_id")
if tenant_id and scan_id:
logger.error(f"Attack paths scan task {task_id} failed: {exc}")
if isinstance(exc, ProviderDeletedException):
logger.warning(
f"Attack paths scan task {task_id} stopped because its provider "
f"or tenant was deleted: {exc}"
)
else:
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))
api/src/backend/tasks/tests/test_attack_paths_aws.py
+102
View File
@@ -0,0 +1,102 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import neo4j.exceptions
import pytest
from tasks.jobs.attack_paths import aws
DATABASE_NOT_FOUND_CODE = "Neo.ClientError.Database.DatabaseNotFound"
def _make_neo4j_error(code: str) -> neo4j.exceptions.Neo4jError:
return neo4j.exceptions.Neo4jError._hydrate_neo4j(
code=code,
message="graph query failed",
)
def _resource_functions(failing_sync, following_sync):
return {
"failing_sync": failing_sync,
"following_sync": following_sync,
"permission_relationships": MagicMock(),
"resourcegroupstaggingapi": MagicMock(),
}
def test_sync_aws_account_reraises_database_not_found_immediately():
error = _make_neo4j_error(DATABASE_NOT_FOUND_CODE)
failing_sync = MagicMock(side_effect=error)
following_sync = MagicMock()
with (
patch.object(
aws.cartography_aws,
"RESOURCE_FUNCTIONS",
_resource_functions(failing_sync, following_sync),
),
patch.object(aws.db_utils, "update_attack_paths_scan_progress"),
patch.object(aws.utils, "stringify_exception") as stringify_exception,
patch.object(aws.logger, "warning") as warning,
pytest.raises(neo4j.exceptions.Neo4jError) as exc_info,
):
aws.sync_aws_account(
SimpleNamespace(uid="123456789012"),
[
"failing_sync",
"following_sync",
"permission_relationships",
"resourcegroupstaggingapi",
],
{},
MagicMock(),
)
assert exc_info.value is error
following_sync.assert_not_called()
stringify_exception.assert_not_called()
warning.assert_not_called()
@pytest.mark.parametrize(
"error",
[
_make_neo4j_error("Neo.ClientError.Statement.SyntaxError"),
RuntimeError("resource sync failed"),
],
ids=["different-neo4j-error", "non-neo4j-error"],
)
def test_sync_aws_account_warns_and_continues_for_other_exceptions(error):
failing_sync = MagicMock(side_effect=error)
following_sync = MagicMock()
with (
patch.object(
aws.cartography_aws,
"RESOURCE_FUNCTIONS",
_resource_functions(failing_sync, following_sync),
),
patch.object(aws.db_utils, "update_attack_paths_scan_progress"),
patch.object(
aws.utils,
"stringify_exception",
return_value="formatted failure",
),
patch.object(aws.logger, "warning") as warning,
):
failed_syncs = aws.sync_aws_account(
SimpleNamespace(uid="123456789012"),
[
"failing_sync",
"following_sync",
"permission_relationships",
"resourcegroupstaggingapi",
],
{},
MagicMock(),
)
assert failed_syncs == {"failing_sync": "formatted failure"}
following_sync.assert_called_once_with()
warning.assert_called_once()
assert "Continuing to the next AWS sync function" in warning.call_args.args[0]
api/src/backend/tasks/tests/test_attack_paths_scan.py
+170 -49
View File
@@ -1,3 +1,4 @@
import logging
from contextlib import nullcontext
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
@@ -5,7 +6,9 @@ from unittest.mock import MagicMock, call, patch
from uuid import uuid4
import pytest
from api.attack_paths.database import GraphDatabaseQueryException
from api.db_utils import rls_transaction
from api.exceptions import ProviderDeletedException
from api.models import (
AttackPathsScan,
Finding,
@@ -250,6 +253,32 @@ class TestAttackPathsRun:
mock_starting.assert_not_called()
mock_create_db.assert_not_called()
@pytest.mark.parametrize(
("ingestion_error", "temporary_database_missing"),
[
(RuntimeError("ingestion boom"), False),
(
GraphDatabaseQueryException(
message="Graph not found: db-scan-id",
code="Neo.ClientError.Database.DatabaseNotFound",
),
True,
),
(
GraphDatabaseQueryException(
message="Graph not found: db-tenant-id",
code="Neo.ClientError.Database.DatabaseNotFound",
),
False,
),
],
ids=[
"regular-error",
"temporary-database-missing",
"sink-database-missing",
],
)
@patch("tasks.jobs.attack_paths.scan.logger")
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
return_value="Cartography failed: ingestion boom",
@@ -302,6 +331,9 @@ class TestAttackPathsRun:
mock_drop_db,
mock_event_loop,
mock_stringify,
mock_logger,
ingestion_error,
temporary_database_missing,
tenants_fixture,
aws_provider,
scans_fixture,
@@ -321,7 +353,11 @@ class TestAttackPathsRun:
session_ctx = MagicMock()
session_ctx.__enter__.return_value = mock_session
session_ctx.__exit__.return_value = False
ingestion_fn = MagicMock(side_effect=RuntimeError("ingestion boom"))
ingestion_fn = MagicMock(side_effect=ingestion_error)
if temporary_database_missing:
mock_finish.side_effect = DatabaseError(
"Save with update_fields did not affect any rows"
)
with (
patch(
@@ -337,13 +373,28 @@ class TestAttackPathsRun:
return_value=ingestion_fn,
),
):
with pytest.raises(RuntimeError, match="ingestion boom"):
with pytest.raises(type(ingestion_error)):
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"}
mock_drop_db.assert_called_once_with("db-scan-id")
if temporary_database_missing:
mock_logger.warning.assert_any_call("Cartography failed: ingestion boom")
mock_logger.exception.assert_not_called()
mock_logger.log.assert_called_once_with(
logging.WARNING,
f"Could not mark Attack Paths scan {attack_paths_scan.id} as `FAILED` "
"(row may have been deleted): Save with update_fields did not affect "
"any rows",
exc_info=False,
)
else:
mock_logger.exception.assert_called_once_with(
"Cartography failed: ingestion boom"
)
@patch(
"tasks.jobs.attack_paths.scan.utils.stringify_exception",
@@ -1265,6 +1316,33 @@ class TestAttackPathsScanRLSTaskOnFailure:
mock_fail.assert_called_once_with("t-1", "s-1", "boom")
def test_on_failure_logs_provider_deletion_as_warning(self):
from tasks.tasks import AttackPathsScanRLSTask
task = AttackPathsScanRLSTask()
error = ProviderDeletedException("provider deleted")
with (
patch("tasks.tasks.logger") as mock_logger,
patch(
"tasks.tasks.attack_paths_db_utils.fail_attack_paths_scan"
) as mock_fail,
):
task.on_failure(
exc=error,
task_id="task-abc",
args=(),
kwargs={"tenant_id": "t-1", "scan_id": "s-1"},
_einfo=None,
)
mock_logger.warning.assert_called_once_with(
"Attack paths scan task task-abc stopped because its provider or tenant "
"was deleted: provider deleted"
)
mock_logger.error.assert_not_called()
mock_fail.assert_called_once_with("t-1", "s-1", "provider deleted")
def test_on_failure_skips_when_missing_kwargs(self):
from tasks.tasks import AttackPathsScanRLSTask
@@ -1896,7 +1974,7 @@ class TestSyncNodes:
mock_source_1.run.return_value = [row]
mock_source_2 = MagicMock()
mock_source_2.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1000)
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
@@ -1933,7 +2011,7 @@ class TestSyncNodes:
src_1.run.return_value = [row]
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1000)
sink.write_nodes.side_effect = lambda *_a, **_kw: call_order.append(
"sink:write"
)
@@ -1969,18 +2047,15 @@ class TestSyncNodes:
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1)
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
_make_session_ctx(src_3),
],
):
result = sync_module.sync_nodes("src", "tgt", "t-1", "p-1", sink, [])
@@ -2009,17 +2084,14 @@ class TestSyncNodes:
src_1.run.return_value = [row]
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=2)
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 2),
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
],
):
result = sync_module.sync_nodes(
"src", "tgt", "t-1", "p-1", sink, normalized_lists
@@ -2037,7 +2109,7 @@ class TestSyncNodes:
def test_sync_nodes_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1000)
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
@@ -2066,7 +2138,7 @@ class TestSyncRelationships:
src_1.run.return_value = [row]
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1000)
sink.write_relationships.side_effect = lambda *_a, **_kw: call_order.append(
"sink:write"
)
@@ -2104,18 +2176,15 @@ class TestSyncRelationships:
src_2.run.return_value = [row_b]
src_3 = MagicMock()
src_3.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1)
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
_make_session_ctx(src_3),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 1),
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
_make_session_ctx(src_3),
],
):
total = sync_module.sync_relationships("src", "tgt", "p-1", sink)
@@ -2140,17 +2209,14 @@ class TestSyncRelationships:
src_1.run.return_value = rows
src_2 = MagicMock()
src_2.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=2)
with (
patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
],
),
patch("tasks.jobs.attack_paths.sync.SYNC_BATCH_SIZE", 2),
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
side_effect=[
_make_session_ctx(src_1),
_make_session_ctx(src_2),
],
):
total = sync_module.sync_relationships("src", "tgt", "p-1", sink)
@@ -2163,7 +2229,7 @@ class TestSyncRelationships:
def test_sync_relationships_empty_source_returns_zero(self):
src = MagicMock()
src.run.return_value = []
sink = MagicMock()
sink = MagicMock(sync_batch_size=1000)
with patch(
"tasks.jobs.attack_paths.sync.graph_database.get_session",
@@ -3056,6 +3122,61 @@ class TestCleanupStaleAttackPathsScans:
ap_scan.refresh_from_db()
assert ap_scan.state == StateChoices.FAILED
@pytest.mark.parametrize(
("age_seconds", "should_clean"),
[
(960 * 60 - 1, False),
(960 * 60, False),
(960 * 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_stale_threshold_boundary_is_strict(
self,
mock_ping,
mock_revoke,
mock_drop_db,
mock_recover,
age_seconds,
should_clean,
tenants_fixture,
aws_provider,
):
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,
started_at=now - timedelta(seconds=age_seconds),
worker="live-worker@host",
)
mock_ping.return_value = ({"live-worker@host"}, set())
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(