feat(attack-surfaces): add new endpoints to retrieve overview data (#9309)

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

@@ -12,7 +12,7 @@ from celery.utils.log import get_task_logger
 from config.env import env
 from config.settings.celery import CELERY_DEADLOCK_ATTEMPTS
 from django.db import IntegrityError, OperationalError
-from django.db.models import Case, Count, IntegerField, Prefetch, Sum, When
+from django.db.models import Case, Count, IntegerField, Prefetch, Q, Sum, When
 from tasks.utils import CustomEncoder
 
 from api.compliance import PROWLER_COMPLIANCE_OVERVIEW_TEMPLATE
@@ -26,6 +26,7 @@ from api.db_utils import (
 )
 from api.exceptions import ProviderConnectionError
 from api.models import (
+    AttackSurfaceOverview,
     ComplianceOverviewSummary,
     ComplianceRequirementOverview,
     Finding,
@@ -43,6 +44,7 @@ from api.models import (
 from api.models import StatusChoices as FindingStatus
 from api.utils import initialize_prowler_provider, return_prowler_provider
 from api.v1.serializers import ScanTaskSerializer
+from prowler.lib.check.models import CheckMetadata
 from prowler.lib.outputs.finding import Finding as ProwlerFinding
 from prowler.lib.scan.scan import Scan as ProwlerScan
 
@@ -75,6 +77,44 @@ FINDINGS_MICRO_BATCH_SIZE = env.int("DJANGO_FINDINGS_MICRO_BATCH_SIZE", default=
 SCAN_DB_BATCH_SIZE = env.int("DJANGO_SCAN_DB_BATCH_SIZE", default=500)
 
 
+ATTACK_SURFACE_PROVIDER_COMPATIBILITY = {
+    "internet-exposed": None,  # Compatible with all providers
+    "secrets": None,  # Compatible with all providers
+    "privilege-escalation": ["aws", "kubernetes"],
+    "ec2-imdsv1": ["aws"],
+}
+
+_ATTACK_SURFACE_MAPPING_CACHE: dict[str, dict] = {}
+
+
+def _get_attack_surface_mapping_from_provider(provider_type: str) -> dict:
+    global _ATTACK_SURFACE_MAPPING_CACHE
+
+    if provider_type in _ATTACK_SURFACE_MAPPING_CACHE:
+        return _ATTACK_SURFACE_MAPPING_CACHE[provider_type]
+
+    attack_surface_check_mappings = {
+        "internet-exposed": None,
+        "secrets": None,
+        "privilege-escalation": {
+            "iam_policy_allows_privilege_escalation",
+            "iam_inline_policy_allows_privilege_escalation",
+        },
+        "ec2-imdsv1": {
+            "ec2_instance_imdsv2_enabled"
+        },  # AWS only - IMDSv1 enabled findings
+    }
+    for category_name, check_ids in attack_surface_check_mappings.items():
+        if check_ids is None:
+            sdk_check_ids = CheckMetadata.list(
+                provider=provider_type, category=category_name
+            )
+            attack_surface_check_mappings[category_name] = sdk_check_ids
+
+    _ATTACK_SURFACE_MAPPING_CACHE[provider_type] = attack_surface_check_mappings
+    return attack_surface_check_mappings
+
+
 def _create_finding_delta(
     last_status: FindingStatus | None | str, new_status: FindingStatus | None
 ) -> Finding.DeltaChoices:
@@ -1196,3 +1236,115 @@ def create_compliance_requirements(tenant_id: str, scan_id: str):
     except Exception as e:
         logger.error(f"Error creating compliance requirements for scan {scan_id}: {e}")
         raise e
+
+
+def aggregate_attack_surface(tenant_id: str, scan_id: str):
+    """
+    Aggregate findings into attack surface overview records.
+
+    Creates one AttackSurfaceOverview record per attack surface type
+    for the given scan, based on check_id mappings.
+
+    Args:
+        tenant_id: Tenant that owns the scan.
+        scan_id: Scan UUID whose findings should be aggregated.
+    """
+    with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS):
+        scan_instance = Scan.all_objects.select_related("provider").get(pk=scan_id)
+        provider_type = scan_instance.provider.provider
+
+    provider_attack_surface_mapping = _get_attack_surface_mapping_from_provider(
+        provider_type=provider_type
+    )
+
+    # Filter out attack surfaces that are not compatible or have no resolved check IDs
+    supported_mappings: dict[str, list[str]] = {}
+    for attack_surface_type, check_ids in provider_attack_surface_mapping.items():
+        compatible_providers = ATTACK_SURFACE_PROVIDER_COMPATIBILITY.get(
+            attack_surface_type
+        )
+        if (
+            compatible_providers is not None
+            and provider_type not in compatible_providers
+        ):
+            logger.info(
+                f"Skipping {attack_surface_type} - not supported for {provider_type}"
+            )
+            continue
+
+        if not check_ids:
+            logger.info(
+                f"Skipping {attack_surface_type} - no check IDs resolved for {provider_type}"
+            )
+            continue
+
+        supported_mappings[attack_surface_type] = list(check_ids)
+
+    if not supported_mappings:
+        logger.info(
+            f"No attack surface mappings available for scan {scan_id} and provider {provider_type}"
+        )
+        logger.info(f"No attack surface overview records created for scan {scan_id}")
+        return
+
+    # Map every check_id to its attack surface, so we can aggregate with a single query
+    check_id_to_surface: dict[str, str] = {}
+    for attack_surface_type, check_ids in supported_mappings.items():
+        for check_id in check_ids:
+            check_id_to_surface[check_id] = attack_surface_type
+
+    aggregated_counts = {
+        attack_surface_type: {"total": 0, "failed": 0, "muted": 0}
+        for attack_surface_type in supported_mappings.keys()
+    }
+
+    with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS):
+        finding_stats = (
+            Finding.all_objects.filter(
+                tenant_id=tenant_id,
+                scan_id=scan_id,
+                check_id__in=list(check_id_to_surface.keys()),
+            )
+            .values("check_id")
+            .annotate(
+                total=Count("id"),
+                failed=Count("id", filter=Q(status="FAIL", muted=False)),
+                muted=Count("id", filter=Q(status="FAIL", muted=True)),
+            )
+        )
+
+        for stats in finding_stats:
+            attack_surface_type = check_id_to_surface.get(stats["check_id"])
+            if not attack_surface_type:
+                continue
+
+            aggregated_counts[attack_surface_type]["total"] += stats["total"] or 0
+            aggregated_counts[attack_surface_type]["failed"] += stats["failed"] or 0
+            aggregated_counts[attack_surface_type]["muted"] += stats["muted"] or 0
+
+    overview_objects = []
+    for attack_surface_type, counts in aggregated_counts.items():
+        total = counts["total"]
+        if not total:
+            continue
+
+        overview_objects.append(
+            AttackSurfaceOverview(
+                tenant_id=tenant_id,
+                scan_id=scan_id,
+                attack_surface_type=attack_surface_type,
+                total_findings=total,
+                failed_findings=counts["failed"],
+                muted_failed_findings=counts["muted"],
+            )
+        )
+
+    # Bulk create overview records
+    if overview_objects:
+        with rls_transaction(tenant_id):
+            AttackSurfaceOverview.objects.bulk_create(overview_objects, batch_size=500)
+            logger.info(
+                f"Created {len(overview_objects)} attack surface overview records for scan {scan_id}"
+            )
+    else:
+        logger.info(f"No attack surface overview records created for scan {scan_id}")

api/src/backend/tasks/tasks.py

@@ -37,6 +37,7 @@ from tasks.jobs.lighthouse_providers import (
 from tasks.jobs.muting import mute_historical_findings
 from tasks.jobs.report import generate_compliance_reports_job
 from tasks.jobs.scan import (
+    aggregate_attack_surface,
     aggregate_findings,
     create_compliance_requirements,
     perform_prowler_scan,
@@ -69,6 +70,9 @@ def _perform_scan_complete_tasks(tenant_id: str, scan_id: str, provider_id: str)
     create_compliance_requirements_task.apply_async(
         kwargs={"tenant_id": tenant_id, "scan_id": scan_id}
     )
+    aggregate_attack_surface_task.apply_async(
+        kwargs={"tenant_id": tenant_id, "scan_id": scan_id}
+    )
     chain(
         perform_scan_summary_task.si(tenant_id=tenant_id, scan_id=scan_id),
         generate_outputs_task.si(
@@ -529,6 +533,21 @@ def create_compliance_requirements_task(tenant_id: str, scan_id: str):
     return create_compliance_requirements(tenant_id=tenant_id, scan_id=scan_id)
 
 
+@shared_task(name="scan-attack-surface-overviews", queue="overview")
+def aggregate_attack_surface_task(tenant_id: str, scan_id: str):
+    """
+    Creates attack surface overview records for a scan.
+
+    This task processes findings and aggregates them into attack surface categories
+    (internet-exposed, secrets, privilege-escalation, ec2-imdsv1) for quick overview queries.
+
+    Args:
+        tenant_id (str): The tenant ID for which to create records.
+        scan_id (str): The ID of the scan for which to create records.
+    """
+    return aggregate_attack_surface(tenant_id=tenant_id, scan_id=scan_id)
+
+
 @shared_task(base=RLSTask, name="lighthouse-connection-check")
 @set_tenant
 def check_lighthouse_connection_task(lighthouse_config_id: str, tenant_id: str = None):

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

@@ -9,14 +9,17 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 from tasks.jobs.scan import (
+    _ATTACK_SURFACE_MAPPING_CACHE,
     _aggregate_findings_by_region,
     _copy_compliance_requirement_rows,
     _create_compliance_summaries,
     _create_finding_delta,
+    _get_attack_surface_mapping_from_provider,
     _normalized_compliance_key,
     _persist_compliance_requirement_rows,
     _process_finding_micro_batch,
     _store_resources,
+    aggregate_attack_surface,
     aggregate_findings,
     create_compliance_requirements,
     perform_prowler_scan,
@@ -3474,3 +3477,282 @@ class TestAggregateFindingsByRegion:
 
         assert check_status_by_region == {}
         assert findings_count_by_compliance == {}
+
+
+@pytest.mark.django_db
+class TestAggregateAttackSurface:
+    """Test aggregate_attack_surface function and related caching."""
+
+    def setup_method(self):
+        """Clear cache before each test."""
+        _ATTACK_SURFACE_MAPPING_CACHE.clear()
+
+    def teardown_method(self):
+        """Clear cache after each test."""
+        _ATTACK_SURFACE_MAPPING_CACHE.clear()
+
+    @patch("tasks.jobs.scan.CheckMetadata.list")
+    def test_get_attack_surface_mapping_caches_result(self, mock_check_metadata_list):
+        """Test that _get_attack_surface_mapping_from_provider caches results."""
+        mock_check_metadata_list.return_value = {"check_internet_exposed_1"}
+
+        # First call should hit CheckMetadata.list
+        result1 = _get_attack_surface_mapping_from_provider("aws")
+        assert mock_check_metadata_list.call_count == 2  # internet-exposed, secrets
+
+        # Second call should use cache
+        result2 = _get_attack_surface_mapping_from_provider("aws")
+        assert mock_check_metadata_list.call_count == 2  # No additional calls
+
+        assert result1 is result2
+        assert "aws" in _ATTACK_SURFACE_MAPPING_CACHE
+
+    @patch("tasks.jobs.scan.CheckMetadata.list")
+    def test_get_attack_surface_mapping_different_providers(
+        self, mock_check_metadata_list
+    ):
+        """Test caching works independently for different providers."""
+        mock_check_metadata_list.return_value = {"check_1"}
+
+        _get_attack_surface_mapping_from_provider("aws")
+        aws_call_count = mock_check_metadata_list.call_count
+
+        _get_attack_surface_mapping_from_provider("gcp")
+        gcp_call_count = mock_check_metadata_list.call_count
+
+        # Both providers should have made calls
+        assert gcp_call_count > aws_call_count
+        assert "aws" in _ATTACK_SURFACE_MAPPING_CACHE
+        assert "gcp" in _ATTACK_SURFACE_MAPPING_CACHE
+
+    @patch("tasks.jobs.scan.CheckMetadata.list")
+    def test_get_attack_surface_mapping_returns_hardcoded_checks(
+        self, mock_check_metadata_list
+    ):
+        """Test that hardcoded check IDs are returned for privilege-escalation and ec2-imdsv1."""
+        mock_check_metadata_list.return_value = set()
+
+        result = _get_attack_surface_mapping_from_provider("aws")
+
+        # Hardcoded checks should be present
+        assert (
+            "iam_policy_allows_privilege_escalation" in result["privilege-escalation"]
+        )
+        assert (
+            "iam_inline_policy_allows_privilege_escalation"
+            in result["privilege-escalation"]
+        )
+        assert "ec2_instance_imdsv2_enabled" in result["ec2-imdsv1"]
+
+    @patch("tasks.jobs.scan.AttackSurfaceOverview.objects.bulk_create")
+    @patch("tasks.jobs.scan.Finding.all_objects.filter")
+    @patch("tasks.jobs.scan._get_attack_surface_mapping_from_provider")
+    @patch("tasks.jobs.scan.rls_transaction")
+    def test_aggregate_attack_surface_creates_overview_records(
+        self,
+        mock_rls_transaction,
+        mock_get_mapping,
+        mock_findings_filter,
+        mock_bulk_create,
+        tenants_fixture,
+        scans_fixture,
+    ):
+        """Test that aggregate_attack_surface creates AttackSurfaceOverview records."""
+        tenant = tenants_fixture[0]
+        scan = scans_fixture[0]
+        scan.provider.provider = "aws"
+        scan.provider.save()
+
+        mock_get_mapping.return_value = {
+            "internet-exposed": {"check_internet_1", "check_internet_2"},
+            "secrets": {"check_secrets_1"},
+            "privilege-escalation": {"check_privesc_1"},
+            "ec2-imdsv1": {"check_imdsv1_1"},
+        }
+
+        # Mock findings aggregation
+        mock_queryset = MagicMock()
+        mock_queryset.values.return_value = mock_queryset
+        mock_queryset.annotate.return_value = [
+            {"check_id": "check_internet_1", "total": 10, "failed": 3, "muted": 1},
+            {"check_id": "check_secrets_1", "total": 5, "failed": 2, "muted": 0},
+        ]
+
+        ctx = MagicMock()
+        ctx.__enter__.return_value = None
+        ctx.__exit__.return_value = False
+        mock_rls_transaction.return_value = ctx
+        mock_findings_filter.return_value = mock_queryset
+
+        aggregate_attack_surface(str(tenant.id), str(scan.id))
+
+        mock_bulk_create.assert_called_once()
+        args, kwargs = mock_bulk_create.call_args
+        objects = args[0]
+
+        # Should create records for internet-exposed and secrets (the ones with findings)
+        assert len(objects) == 2
+        assert kwargs["batch_size"] == 500
+
+    @patch("tasks.jobs.scan.AttackSurfaceOverview.objects.bulk_create")
+    @patch("tasks.jobs.scan.Finding.all_objects.filter")
+    @patch("tasks.jobs.scan._get_attack_surface_mapping_from_provider")
+    @patch("tasks.jobs.scan.rls_transaction")
+    def test_aggregate_attack_surface_skips_unsupported_provider(
+        self,
+        mock_rls_transaction,
+        mock_get_mapping,
+        mock_findings_filter,
+        mock_bulk_create,
+        tenants_fixture,
+        scans_fixture,
+    ):
+        """Test that ec2-imdsv1 is skipped for non-AWS providers."""
+        tenant = tenants_fixture[0]
+        scan = scans_fixture[0]
+        scan.provider.provider = "gcp"
+        scan.provider.uid = "gcp-test-project-id"
+        scan.provider.save()
+
+        mock_get_mapping.return_value = {
+            "internet-exposed": {"check_internet_1"},
+            "secrets": {"check_secrets_1"},
+            "privilege-escalation": set(),  # Not supported for GCP
+            "ec2-imdsv1": {"check_imdsv1_1"},  # Should be skipped for GCP
+        }
+
+        mock_queryset = MagicMock()
+        mock_queryset.values.return_value = mock_queryset
+        mock_queryset.annotate.return_value = [
+            {"check_id": "check_internet_1", "total": 5, "failed": 1, "muted": 0},
+        ]
+
+        ctx = MagicMock()
+        ctx.__enter__.return_value = None
+        ctx.__exit__.return_value = False
+        mock_rls_transaction.return_value = ctx
+        mock_findings_filter.return_value = mock_queryset
+
+        aggregate_attack_surface(str(tenant.id), str(scan.id))
+
+        # ec2-imdsv1 check_ids should not be in the filter
+        filter_call = mock_findings_filter.call_args
+        check_ids_in_filter = filter_call[1]["check_id__in"]
+        assert "check_imdsv1_1" not in check_ids_in_filter
+
+    @patch("tasks.jobs.scan.AttackSurfaceOverview.objects.bulk_create")
+    @patch("tasks.jobs.scan.Finding.all_objects.filter")
+    @patch("tasks.jobs.scan._get_attack_surface_mapping_from_provider")
+    @patch("tasks.jobs.scan.rls_transaction")
+    def test_aggregate_attack_surface_no_findings(
+        self,
+        mock_rls_transaction,
+        mock_get_mapping,
+        mock_findings_filter,
+        mock_bulk_create,
+        tenants_fixture,
+        scans_fixture,
+    ):
+        """Test that no records are created when there are no findings."""
+        tenant = tenants_fixture[0]
+        scan = scans_fixture[0]
+
+        mock_get_mapping.return_value = {
+            "internet-exposed": {"check_1"},
+            "secrets": {"check_2"},
+            "privilege-escalation": set(),
+            "ec2-imdsv1": set(),
+        }
+
+        mock_queryset = MagicMock()
+        mock_queryset.values.return_value = mock_queryset
+        mock_queryset.annotate.return_value = []  # No findings
+
+        ctx = MagicMock()
+        ctx.__enter__.return_value = None
+        ctx.__exit__.return_value = False
+        mock_rls_transaction.return_value = ctx
+        mock_findings_filter.return_value = mock_queryset
+
+        aggregate_attack_surface(str(tenant.id), str(scan.id))
+
+        mock_bulk_create.assert_not_called()
+
+    @patch("tasks.jobs.scan.AttackSurfaceOverview.objects.bulk_create")
+    @patch("tasks.jobs.scan.Finding.all_objects.filter")
+    @patch("tasks.jobs.scan._get_attack_surface_mapping_from_provider")
+    @patch("tasks.jobs.scan.rls_transaction")
+    def test_aggregate_attack_surface_aggregates_counts_correctly(
+        self,
+        mock_rls_transaction,
+        mock_get_mapping,
+        mock_findings_filter,
+        mock_bulk_create,
+        tenants_fixture,
+        scans_fixture,
+    ):
+        """Test that counts from multiple check_ids are aggregated per attack surface type."""
+        tenant = tenants_fixture[0]
+        scan = scans_fixture[0]
+        scan.provider.provider = "aws"
+        scan.provider.save()
+
+        mock_get_mapping.return_value = {
+            "internet-exposed": {"check_internet_1", "check_internet_2"},
+            "secrets": set(),
+            "privilege-escalation": set(),
+            "ec2-imdsv1": set(),
+        }
+
+        mock_queryset = MagicMock()
+        mock_queryset.values.return_value = mock_queryset
+        mock_queryset.annotate.return_value = [
+            {"check_id": "check_internet_1", "total": 10, "failed": 3, "muted": 1},
+            {"check_id": "check_internet_2", "total": 5, "failed": 2, "muted": 0},
+        ]
+
+        ctx = MagicMock()
+        ctx.__enter__.return_value = None
+        ctx.__exit__.return_value = False
+        mock_rls_transaction.return_value = ctx
+        mock_findings_filter.return_value = mock_queryset
+
+        aggregate_attack_surface(str(tenant.id), str(scan.id))
+
+        args, kwargs = mock_bulk_create.call_args
+        objects = args[0]
+
+        assert len(objects) == 1
+        overview = objects[0]
+        assert overview.attack_surface_type == "internet-exposed"
+        assert overview.total_findings == 15  # 10 + 5
+        assert overview.failed_findings == 5  # 3 + 2
+        assert overview.muted_failed_findings == 1  # 1 + 0
+
+    @patch("tasks.jobs.scan.Scan.all_objects.select_related")
+    @patch("tasks.jobs.scan.rls_transaction")
+    def test_aggregate_attack_surface_uses_select_related(
+        self, mock_rls_transaction, mock_select_related, tenants_fixture, scans_fixture
+    ):
+        """Test that select_related is used to avoid N+1 query."""
+        tenant = tenants_fixture[0]
+        scan = scans_fixture[0]
+
+        mock_scan = MagicMock()
+        mock_scan.provider.provider = "aws"
+
+        mock_select_related.return_value.get.return_value = mock_scan
+
+        ctx = MagicMock()
+        ctx.__enter__.return_value = None
+        ctx.__exit__.return_value = False
+        mock_rls_transaction.return_value = ctx
+
+        with patch(
+            "tasks.jobs.scan._get_attack_surface_mapping_from_provider"
+        ) as mock_map:
+            mock_map.return_value = {}
+
+            aggregate_attack_surface(str(tenant.id), str(scan.id))
+
+        mock_select_related.assert_called_once_with("provider")

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

@@ -529,6 +529,7 @@ class TestGenerateOutputs:
 
 
 class TestScanCompleteTasks:
+    @patch("tasks.tasks.aggregate_attack_surface_task.apply_async")
     @patch("tasks.tasks.create_compliance_requirements_task.apply_async")
     @patch("tasks.tasks.perform_scan_summary_task.si")
     @patch("tasks.tasks.generate_outputs_task.si")
@@ -541,6 +542,7 @@ class TestScanCompleteTasks:
         mock_outputs_task,
         mock_scan_summary_task,
         mock_compliance_requirements_task,
+        mock_attack_surface_task,
     ):
         """Test that scan complete tasks are properly orchestrated with optimized reports."""
         _perform_scan_complete_tasks("tenant-id", "scan-id", "provider-id")
@@ -550,6 +552,11 @@ class TestScanCompleteTasks:
             kwargs={"tenant_id": "tenant-id", "scan_id": "scan-id"},
         )
 
+        # Verify attack surface task is called
+        mock_attack_surface_task.assert_called_once_with(
+            kwargs={"tenant_id": "tenant-id", "scan_id": "scan-id"},
+        )
+
         # Verify scan summary task is called
         mock_scan_summary_task.assert_called_once_with(
             scan_id="scan-id",