Merge branch 'attack-paths-demo' into attack-paths-demo-extras

api/src/backend/api/attack_paths/database.py

@@ -6,15 +6,18 @@ from typing import Iterator
 from uuid import UUID
 
 import neo4j
+import neo4j.exceptions
 
 from django.conf import settings
 
-import neo4j.exceptions
+from api.attack_paths.retryable_session import RetryableSession
 
 # Without this Celery goes crazy with Neo4j logging
 logging.getLogger("neo4j").setLevel(logging.ERROR)
 logging.getLogger("neo4j").propagate = False
 
+SERVICE_UNAVAILABLE_MAX_RETRIES = 3
+
 # Module-level process-wide driver singleton
 _driver: neo4j.Driver | None = None
 _lock = threading.Lock()
@@ -62,14 +65,24 @@ def close_driver() -> None:  # TODO: Use it
 
 
 @contextmanager
-def get_session(database: str | None = None) -> Iterator[neo4j.Session]:
+def get_session(database: str | None = None) -> Iterator[RetryableSession]:
+    session_wrapper: RetryableSession | None = None
+
     try:
-        with get_driver().session(database=database) as session:
-            yield session
+        session_wrapper = RetryableSession(
+            session_factory=lambda: get_driver().session(database=database),
+            close_driver=close_driver,  # Just to avoid circular imports
+            max_retries=SERVICE_UNAVAILABLE_MAX_RETRIES,
+        )
+        yield session_wrapper
 
     except neo4j.exceptions.Neo4jError as exc:
         raise GraphDatabaseQueryException(message=exc.message, code=exc.code)
 
+    finally:
+        if session_wrapper is not None:
+            session_wrapper.close()
+
 
 def create_database(database: str) -> None:
     query = "CREATE DATABASE $database IF NOT EXISTS"

api/src/backend/api/attack_paths/retryable_session.py

@@ -0,0 +1,87 @@
+import logging
+
+from collections.abc import Callable
+from typing import Any
+
+import neo4j
+import neo4j.exceptions
+
+logger = logging.getLogger(__name__)
+
+
+class RetryableSession:
+    """
+    Wrapper around `neo4j.Session` that retries `neo4j.exceptions.ServiceUnavailable` errors.
+    """
+
+    def __init__(
+        self,
+        session_factory: Callable[[], neo4j.Session],
+        close_driver: Callable[[], None],  # Just to avoid circular imports
+        max_retries: int,
+    ) -> None:
+        self._session_factory = session_factory
+        self._close_driver = close_driver
+        self._max_retries = max(0, max_retries)
+        self._session = self._session_factory()
+
+    def close(self) -> None:
+        if self._session is not None:
+            self._session.close()
+            self._session = None
+
+    def __enter__(self) -> "RetryableSession":
+        return self
+
+    def __exit__(self, exc_type: Any, exc: Any, exc_tb: Any) -> None:
+        self.close()
+
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        return self._call_with_retry("run", *args, **kwargs)
+
+    def write_transaction(self, *args: Any, **kwargs: Any) -> Any:
+        return self._call_with_retry("write_transaction", *args, **kwargs)
+
+    def read_transaction(self, *args: Any, **kwargs: Any) -> Any:
+        return self._call_with_retry("read_transaction", *args, **kwargs)
+
+    def execute_write(self, *args: Any, **kwargs: Any) -> Any:
+        return self._call_with_retry("execute_write", *args, **kwargs)
+
+    def execute_read(self, *args: Any, **kwargs: Any) -> Any:
+        return self._call_with_retry("execute_read", *args, **kwargs)
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._session, item)
+
+    def _call_with_retry(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
+        attempt = 0
+        last_exc: neo4j.exceptions.ServiceUnavailable | None = None
+
+        while attempt <= self._max_retries:
+            try:
+                method = getattr(self._session, method_name)
+                return method(*args, **kwargs)
+
+            except (
+                neo4j.exceptions.ServiceUnavailable
+            ) as exc:  # pragma: no cover - depends on infra
+                last_exc = exc
+                attempt += 1
+
+                if attempt > self._max_retries:
+                    raise
+
+                logger.warning(
+                    f"Neo4j session {method_name} failed with ServiceUnavailable ({attempt}/{self._max_retries} attempts). Retrying..."
+                )
+                self._refresh_session()
+
+        raise last_exc if last_exc else RuntimeError("Unexpected retry loop exit")
+
+    def _refresh_session(self) -> None:
+        if self._session is not None:
+            self._session.close()
+
+        self._close_driver()
+        self._session = self._session_factory()

api/src/backend/tasks/jobs/attack_paths/aws.py

@@ -3,6 +3,7 @@
 
 from typing import Any
 
+import aioboto3
 import boto3
 import neo4j
 
@@ -28,7 +29,7 @@ def start_aws_ingestion(
     attack_paths_scan: ProwlerAPIAttackPathsScan,
 ) -> dict[str, dict[str, str]]:
     """
-    Code based on Cartography version 0.117.0, specifically on `cartography.intel.aws.__init__.py`.
+    Code based on Cartography version 0.122.0, specifically on `cartography.intel.aws.__init__.py`.
 
     For the scan progress updates:
         - The caller of this function (`tasks.jobs.attack_paths.scan.run`) has set it to 2.
@@ -41,6 +42,7 @@ def start_aws_ingestion(
         "permission_relationships_file": cartography_config.permission_relationships_file,
         "aws_guardduty_severity_threshold": cartography_config.aws_guardduty_severity_threshold,
         "aws_cloudtrail_management_events_lookback_hours": cartography_config.aws_cloudtrail_management_events_lookback_hours,
+        "experimental_aws_inspector_batch": cartography_config.experimental_aws_inspector_batch,
     }
 
     boto3_session = get_boto3_session(prowler_api_provider, prowler_sdk_provider)
@@ -157,6 +159,10 @@ def get_boto3_session(
     return boto3_session
 
 
+def get_aioboto3_session(boto3_session: boto3.Session) -> aioboto3.Session:
+    return aioboto3.Session(botocore_session=boto3_session._session)
+
+
 def sync_aws_account(
     prowler_api_provider: ProwlerAPIProvider,
     requested_syncs: list[str],
@@ -167,7 +173,9 @@ def sync_aws_account(
     max_progress = (
         87  # `cartography_aws.RESOURCE_FUNCTIONS["permission_relationships"]` - 1
     )
-    n_steps = len(requested_syncs)
+    n_steps = (
+        len(requested_syncs) - 2
+    )  # Excluding `permission_relationships` and `resourcegroupstaggingapi`
     progress_step = (max_progress - current_progress) / n_steps
 
     failed_syncs = {}
@@ -185,15 +193,28 @@ def sync_aws_account(
             )
 
             try:
+                # `ecr:image_layers` uses `aioboto3_session` instead of `boto3_session`
+                if func_name == "ecr:image_layers":
+                    cartography_aws.RESOURCE_FUNCTIONS[func_name](
+                        neo4j_session=sync_args.get("neo4j_session"),
+                        aioboto3_session=get_aioboto3_session(
+                            sync_args.get("boto3_session")
+                        ),
+                        regions=sync_args.get("regions"),
+                        current_aws_account_id=sync_args.get("current_aws_account_id"),
+                        update_tag=sync_args.get("update_tag"),
+                        common_job_parameters=sync_args.get("common_job_parameters"),
+                    )
+
                 # Skip permission relationships and tags for now because they rely on data already being in the graph
-                if func_name not in [
+                elif func_name in [
                     "permission_relationships",
                     "resourcegroupstaggingapi",
                 ]:
-                    cartography_aws.RESOURCE_FUNCTIONS[func_name](**sync_args)
+                    continue
 
                 else:
-                    continue
+                    cartography_aws.RESOURCE_FUNCTIONS[func_name](**sync_args)
 
             except Exception as e:
                 exception_message = utils.stringify_exception(

api/src/backend/tasks/jobs/attack_paths/prowler.py

@@ -80,7 +80,7 @@ CLEANUP_STATEMENT = """
 
 def create_indexes(neo4j_session: neo4j.Session) -> None:
     """
-    Code based on Cartography version 0.117.0, specifically on `cartography.intel.create_indexes.run`.
+    Code based on Cartography version 0.122.0, specifically on `cartography.intel.create_indexes.run`.
     """
 
     logger.info("Creating indexes for Prowler node types.")

api/src/backend/tasks/jobs/attack_paths/scan.py

@@ -7,6 +7,7 @@ from typing import Any, Callable
 from cartography.config import Config as CartographyConfig
 from cartography.intel import analysis as cartography_analysis
 from cartography.intel import create_indexes as cartography_create_indexes
+from cartography.intel import ontology as cartography_ontology
 from celery.utils.log import get_task_logger
 
 from api.attack_paths import database as graph_database
@@ -35,7 +36,7 @@ def get_cartography_ingestion_function(provider_type: str) -> Callable | None:
 
 def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
     """
-    Code based on Cartography version 0.117.0, specifically on `cartography.cli.main`, `cartography.cli.CLI.main`,
+    Code based on Cartography version 0.122.0, specifically on `cartography.cli.main`, `cartography.cli.CLI.main`,
     `cartography.sync.run_with_config` and `cartography.sync.Sync.run`.
     """
     ingestion_exceptions = {}  # This will hold any exceptions raised during ingestion
@@ -45,19 +46,35 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
         prowler_api_provider = ProwlerAPIProvider.objects.get(scan__pk=scan_id)
         prowler_sdk_provider = initialize_prowler_provider(prowler_api_provider)
 
-    # If the provider is still not supported, just return the current `ingestion_exceptions`, that is empty
-    if not get_cartography_ingestion_function(prowler_api_provider.provider):
+    # Attack Paths Scan necessary objects
+    cartography_ingestion_function = get_cartography_ingestion_function(
+        prowler_api_provider.provider
+    )
+    attack_paths_scan = db_utils.retrieve_attack_paths_scan(tenant_id, scan_id)
+
+    # Checks before starting the scan
+    if not cartography_ingestion_function:
+        ingestion_exceptions = {
+            "global_error": f"Provider {prowler_api_provider.provider} is not supported for Attack Paths scans"
+        }
+        if attack_paths_scan:
+            db_utils.finish_attack_paths_scan(
+                attack_paths_scan, StateChoices.COMPLETED, ingestion_exceptions
+            )
+
+        logger.warning(
+            f"Provider {prowler_api_provider.provider} is not supported for Attack Paths scans"
+        )
         return ingestion_exceptions
 
-    # Getting the Attack Paths Scan object and starting it
-    attack_paths_scan = db_utils.retrieve_attack_paths_scan(tenant_id, scan_id)
-    if not attack_paths_scan:
-        logger.warning(
-            f"No Attack Paths Scan found for scan {scan_id} and tenant {tenant_id}, let's create it then"
-        )
-        attack_paths_scan = db_utils.create_attack_paths_scan(
-            tenant_id, scan_id, prowler_api_provider.id
-        )
+    else:
+        if not attack_paths_scan:
+            logger.warning(
+                f"No Attack Paths Scan found for scan {scan_id} and tenant {tenant_id}, let's create it then"
+            )
+            attack_paths_scan = db_utils.create_attack_paths_scan(
+                tenant_id, scan_id, prowler_api_provider.id
+            )
 
     # While creating the Cartography configuration, attributes `neo4j_user` and `neo4j_password` are not really needed in this config object
     cartography_config = CartographyConfig(
@@ -91,7 +108,7 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
 
             # The real scan, where iterates over cloud services
             ingestion_exceptions = _call_within_event_loop(
-                get_cartography_ingestion_function(prowler_api_provider.provider),
+                cartography_ingestion_function,
                 neo4j_session,
                 cartography_config,
                 prowler_api_provider,
@@ -100,9 +117,12 @@ def run(tenant_id: str, scan_id: str, task_id: str) -> dict[str, Any]:
             )
 
             # Post-processing: Just keeping it to be more Cartography compliant
-            cartography_analysis.run(neo4j_session, cartography_config)
+            cartography_ontology.run(neo4j_session, cartography_config)
             db_utils.update_attack_paths_scan_progress(attack_paths_scan, 95)
 
+            cartography_analysis.run(neo4j_session, cartography_config)
+            db_utils.update_attack_paths_scan_progress(attack_paths_scan, 96)
+
             # Adding Prowler nodes and relationships
             prowler.analysis(
                 neo4j_session, prowler_api_provider, scan_id, cartography_config

api/src/backend/tasks/tests/test_attack_paths_scan.py

@@ -74,6 +74,9 @@ class TestAttackPathsRun:
             patch(
                 "tasks.jobs.attack_paths.scan.cartography_analysis.run"
             ) as mock_cartography_analysis,
+            patch(
+                "tasks.jobs.attack_paths.scan.cartography_ontology.run"
+            ) as mock_cartography_ontology,
             patch(
                 "tasks.jobs.attack_paths.scan.prowler.create_indexes"
             ) as mock_prowler_indexes,
@@ -115,6 +118,7 @@ class TestAttackPathsRun:
         mock_cartography_indexes.assert_called_once_with(mock_session, config)
         mock_prowler_indexes.assert_called_once_with(mock_session)
         mock_cartography_analysis.assert_called_once_with(mock_session, config)
+        mock_cartography_ontology.assert_called_once_with(mock_session, config)
         mock_prowler_analysis.assert_called_once_with(
             mock_session,
             provider,