feat(integrations): implement AWS Security Hub integration (#8365)

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

@@ -2,15 +2,21 @@ import os
 from glob import glob
 
 from celery.utils.log import get_task_logger
+from config.django.base import DJANGO_FINDINGS_BATCH_SIZE
+from tasks.utils import batched
 
 from api.db_utils import rls_transaction
-from api.models import Integration
+from api.models import Finding, Integration, Provider
+from api.utils import initialize_prowler_provider
 from prowler.lib.outputs.asff.asff import ASFF
 from prowler.lib.outputs.compliance.generic.generic import GenericCompliance
 from prowler.lib.outputs.csv.csv import CSV
+from prowler.lib.outputs.finding import Finding as FindingOutput
 from prowler.lib.outputs.html.html import HTML
 from prowler.lib.outputs.ocsf.ocsf import OCSF
+from prowler.providers.aws.aws_provider import AwsProvider
 from prowler.providers.aws.lib.s3.s3 import S3
+from prowler.providers.aws.lib.security_hub.security_hub import SecurityHub
 from prowler.providers.common.models import Connection
 
 logger = get_task_logger(__name__)
@@ -154,3 +160,265 @@ def upload_s3_integration(
     except Exception as e:
         logger.error(f"S3 integrations failed for provider {provider_id}: {str(e)}")
         return False
+
+
+def get_security_hub_client_from_integration(
+    integration: Integration, tenant_id: str, findings: list
+) -> tuple[bool, SecurityHub | Connection]:
+    """
+    Create and return a SecurityHub client using AWS credentials from an integration.
+
+    Args:
+        integration (Integration): The integration to get the Security Hub client from.
+        tenant_id (str): The tenant identifier.
+        findings (list): List of findings in ASFF format to send to Security Hub.
+
+    Returns:
+        tuple[bool, SecurityHub | Connection]: A tuple containing a boolean indicating
+        if the connection was successful and the SecurityHub client or connection object.
+    """
+    # Get the provider associated with this integration
+    with rls_transaction(tenant_id):
+        provider_relationship = integration.integrationproviderrelationship_set.first()
+        if not provider_relationship:
+            return Connection(
+                is_connected=False, error="No provider associated with this integration"
+            )
+        provider_uid = provider_relationship.provider.uid
+        provider_secret = provider_relationship.provider.secret.secret
+
+    credentials = (
+        integration.credentials if integration.credentials else provider_secret
+    )
+    connection = SecurityHub.test_connection(
+        aws_account_id=provider_uid,
+        raise_on_exception=False,
+        **credentials,
+    )
+
+    if connection.is_connected:
+        all_security_hub_regions = AwsProvider.get_available_aws_service_regions(
+            "securityhub", connection.partition
+        )
+
+        # Create regions status dictionary
+        regions_status = {}
+        for region in set(all_security_hub_regions):
+            regions_status[region] = region in connection.enabled_regions
+
+        # Save regions information in the integration configuration
+        with rls_transaction(tenant_id):
+            integration.configuration["regions"] = regions_status
+            integration.save()
+
+        # Create SecurityHub client with all necessary parameters
+        security_hub = SecurityHub(
+            aws_account_id=provider_uid,
+            findings=findings,
+            send_only_fails=integration.configuration.get("send_only_fails", False),
+            aws_security_hub_available_regions=list(connection.enabled_regions),
+            **credentials,
+        )
+        return True, security_hub
+
+    return False, connection
+
+
+def upload_security_hub_integration(
+    tenant_id: str, provider_id: str, scan_id: str
+) -> bool:
+    """
+    Upload findings to AWS Security Hub using configured integrations.
+
+    This function retrieves findings from the database, transforms them to ASFF format,
+    and sends them to AWS Security Hub using the configured integration credentials.
+
+    Args:
+        tenant_id (str): The tenant identifier.
+        provider_id (str): The provider identifier.
+        scan_id (str): The scan identifier for which to send findings.
+
+    Returns:
+        bool: True if all integrations executed successfully, False otherwise.
+    """
+    logger.info(f"Processing Security Hub integrations for provider {provider_id}")
+
+    try:
+        with rls_transaction(tenant_id):
+            # Get Security Hub integrations for this provider
+            integrations = list(
+                Integration.objects.filter(
+                    integrationproviderrelationship__provider_id=provider_id,
+                    integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB,
+                    enabled=True,
+                )
+            )
+
+            if not integrations:
+                logger.error(
+                    f"No Security Hub integrations found for provider {provider_id}"
+                )
+                return False
+
+            # Get the provider object
+            provider = Provider.objects.get(id=provider_id)
+
+            # Initialize prowler provider for finding transformation
+            prowler_provider = initialize_prowler_provider(provider)
+
+        # Process each Security Hub integration
+        integration_executions = 0
+        total_findings_sent = {}  # Track findings sent per integration
+
+        for integration in integrations:
+            try:
+                # Initialize Security Hub client for this integration
+                # We'll create the client once and reuse it for all batches
+                security_hub_client = None
+                send_only_fails = integration.configuration.get(
+                    "send_only_fails", False
+                )
+                total_findings_sent[integration.id] = 0
+
+                # Process findings in batches to avoid memory issues
+                has_findings = False
+                batch_number = 0
+
+                with rls_transaction(tenant_id):
+                    qs = (
+                        Finding.all_objects.filter(tenant_id=tenant_id, scan_id=scan_id)
+                        .order_by("uid")
+                        .iterator()
+                    )
+
+                    for batch, _ in batched(qs, DJANGO_FINDINGS_BATCH_SIZE):
+                        batch_number += 1
+                        has_findings = True
+
+                        # Transform findings for this batch
+                        transformed_findings = [
+                            FindingOutput.transform_api_finding(
+                                finding, prowler_provider
+                            )
+                            for finding in batch
+                        ]
+
+                        # Convert to ASFF format
+                        asff_transformer = ASFF(
+                            findings=transformed_findings,
+                            file_path="",
+                            file_extension="json",
+                        )
+                        asff_transformer.transform(transformed_findings)
+
+                        # Get the batch of ASFF findings
+                        batch_asff_findings = asff_transformer.data
+
+                        if batch_asff_findings:
+                            # Create Security Hub client for first batch or reuse existing
+                            if not security_hub_client:
+                                connected, security_hub = (
+                                    get_security_hub_client_from_integration(
+                                        integration, tenant_id, batch_asff_findings
+                                    )
+                                )
+
+                                if not connected:
+                                    logger.error(
+                                        f"Security Hub connection failed for integration {integration.id}: {security_hub.error}"
+                                    )
+                                    integration.connected = False
+                                    integration.save()
+                                    break  # Skip this integration
+
+                                security_hub_client = security_hub
+                                logger.info(
+                                    f"Sending {'fail' if send_only_fails else 'all'} findings to Security Hub via integration {integration.id}"
+                                )
+                            else:
+                                # Update findings in existing client for this batch
+                                security_hub_client._findings_per_region = (
+                                    security_hub_client.filter(
+                                        batch_asff_findings, send_only_fails
+                                    )
+                                )
+
+                            # Send this batch to Security Hub
+                            try:
+                                findings_sent = (
+                                    security_hub_client.batch_send_to_security_hub()
+                                )
+                                total_findings_sent[integration.id] += findings_sent
+
+                                if findings_sent > 0:
+                                    logger.debug(
+                                        f"Sent batch {batch_number} with {findings_sent} findings to Security Hub"
+                                    )
+                            except Exception as batch_error:
+                                logger.error(
+                                    f"Failed to send batch {batch_number} to Security Hub: {str(batch_error)}"
+                                )
+
+                        # Clear memory after processing each batch
+                        asff_transformer._data.clear()
+                        del batch_asff_findings
+                        del transformed_findings
+
+                if not has_findings:
+                    logger.info(
+                        f"No findings to send to Security Hub for scan {scan_id}"
+                    )
+                    integration_executions += 1
+                elif security_hub_client:
+                    if total_findings_sent[integration.id] > 0:
+                        logger.info(
+                            f"Successfully sent {total_findings_sent[integration.id]} total findings to Security Hub via integration {integration.id}"
+                        )
+                        integration_executions += 1
+                    else:
+                        logger.warning(
+                            f"No findings were sent to Security Hub via integration {integration.id}"
+                        )
+
+                    # Archive previous findings if configured to do so
+                    if integration.configuration.get(
+                        "archive_previous_findings", False
+                    ):
+                        logger.info(
+                            f"Archiving previous findings in Security Hub via integration {integration.id}"
+                        )
+                        try:
+                            findings_archived = (
+                                security_hub_client.archive_previous_findings()
+                            )
+                            logger.info(
+                                f"Successfully archived {findings_archived} previous findings in Security Hub"
+                            )
+                        except Exception as archive_error:
+                            logger.warning(
+                                f"Failed to archive previous findings: {str(archive_error)}"
+                            )
+
+            except Exception as e:
+                logger.error(
+                    f"Security Hub integration {integration.id} failed: {str(e)}"
+                )
+                continue
+
+        result = integration_executions == len(integrations)
+        if result:
+            logger.info(
+                f"All Security Hub integrations completed successfully for provider {provider_id}"
+            )
+        else:
+            logger.error(
+                f"Some Security Hub integrations failed for provider {provider_id}"
+            )
+
+        return result
+
+    except Exception as e:
+        logger.error(
+            f"Security Hub integrations failed for provider {provider_id}: {str(e)}"
+        )
+        return False

api/src/backend/tasks/tasks.py

@@ -21,7 +21,10 @@ from tasks.jobs.export import (
     _generate_output_directory,
     _upload_to_s3,
 )
-from tasks.jobs.integrations import upload_s3_integration
+from tasks.jobs.integrations import (
+    upload_s3_integration,
+    upload_security_hub_integration,
+)
 from tasks.jobs.scan import (
     aggregate_findings,
     create_compliance_requirements,
@@ -62,6 +65,7 @@ def _perform_scan_complete_tasks(tenant_id: str, scan_id: str, provider_id: str)
         check_integrations_task.si(
             tenant_id=tenant_id,
             provider_id=provider_id,
+            scan_id=scan_id,
         ),
     ).apply_async()
 
@@ -323,7 +327,11 @@ def generate_outputs_task(scan_id: str, provider_id: str, tenant_id: str):
         ScanSummary.objects.filter(scan_id=scan_id)
     )
 
-    qs = Finding.all_objects.filter(scan_id=scan_id).order_by("uid").iterator()
+    qs = (
+        Finding.all_objects.filter(tenant_id=tenant_id, scan_id=scan_id)
+        .order_by("uid")
+        .iterator()
+    )
     for batch, is_last in batched(qs, DJANGO_FINDINGS_BATCH_SIZE):
         fos = [FindingOutput.transform_api_finding(f, prowler_provider) for f in batch]
 
@@ -474,17 +482,19 @@ def check_lighthouse_connection_task(lighthouse_config_id: str, tenant_id: str =
 
 
 @shared_task(name="integration-check")
-def check_integrations_task(tenant_id: str, provider_id: str):
+def check_integrations_task(tenant_id: str, provider_id: str, scan_id: str = None):
     """
     Check and execute all configured integrations for a provider.
 
     Args:
         tenant_id (str): The tenant identifier
         provider_id (str): The provider identifier
+        scan_id (str, optional): The scan identifier for integrations that need scan data
     """
     logger.info(f"Checking integrations for provider {provider_id}")
 
     try:
+        integration_tasks = []
         with rls_transaction(tenant_id):
             integrations = Integration.objects.filter(
                 integrationproviderrelationship__provider_id=provider_id,
@@ -495,7 +505,16 @@ def check_integrations_task(tenant_id: str, provider_id: str):
                 logger.info(f"No integrations configured for provider {provider_id}")
                 return {"integrations_processed": 0}
 
-        integration_tasks = []
+            # Security Hub integration
+            security_hub_integrations = integrations.filter(
+                integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB
+            )
+            if security_hub_integrations.exists():
+                integration_tasks.append(
+                    security_hub_integration_task.s(
+                        tenant_id=tenant_id, provider_id=provider_id, scan_id=scan_id
+                    )
+                )
 
         # TODO: Add other integration types here
         # slack_integrations = integrations.filter(
@@ -541,3 +560,24 @@ def s3_integration_task(
         output_directory (str): Path to the directory containing output files
     """
     return upload_s3_integration(tenant_id, provider_id, output_directory)
+
+
+@shared_task(
+    base=RLSTask,
+    name="integration-security-hub",
+    queue="integrations",
+)
+def security_hub_integration_task(
+    tenant_id: str,
+    provider_id: str,
+    scan_id: str,
+):
+    """
+    Process Security Hub integrations for a provider.
+
+    Args:
+        tenant_id (str): The tenant identifier
+        provider_id (str): The provider identifier
+        scan_id (str): The scan identifier
+    """
+    return upload_security_hub_integration(tenant_id, provider_id, scan_id)

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

@@ -7,6 +7,7 @@ from tasks.tasks import (
     check_integrations_task,
     generate_outputs_task,
     s3_integration_task,
+    security_hub_integration_task,
 )
 
 from api.models import Integration
@@ -521,31 +522,68 @@ class TestCheckIntegrationsTask:
             enabled=True,
         )
 
+    @patch("tasks.tasks.security_hub_integration_task")
     @patch("tasks.tasks.group")
     @patch("tasks.tasks.rls_transaction")
     @patch("tasks.tasks.Integration.objects.filter")
-    def test_check_integrations_s3_success(
-        self, mock_integration_filter, mock_rls, mock_group
+    def test_check_integrations_security_hub_success(
+        self, mock_integration_filter, mock_rls, mock_group, mock_security_hub_task
     ):
-        # Mock that we have some integrations
-        mock_integration_filter.return_value.exists.return_value = True
+        """Test that SecurityHub integrations are processed correctly."""
+        # Mock that we have SecurityHub integrations
+        mock_integrations = MagicMock()
+        mock_integrations.exists.return_value = True
+
+        # Mock SecurityHub integrations to return existing integrations
+        mock_security_hub_integrations = MagicMock()
+        mock_security_hub_integrations.exists.return_value = True
+
+        # Set up the filter chain
+        mock_integration_filter.return_value = mock_integrations
+        mock_integrations.filter.return_value = mock_security_hub_integrations
+
+        # Mock the task signature
+        mock_task_signature = MagicMock()
+        mock_security_hub_task.s.return_value = mock_task_signature
+
+        # Mock group job
+        mock_job = MagicMock()
+        mock_group.return_value = mock_job
+
         # Ensure rls_transaction is mocked
         mock_rls.return_value.__enter__.return_value = None
 
-        # Since the current implementation doesn't actually create tasks yet (TODO comment),
-        # we test that no tasks are created but the function returns the correct count
+        # Execute the function
         result = check_integrations_task(
             tenant_id=self.tenant_id,
             provider_id=self.provider_id,
+            scan_id="test-scan-id",
         )
 
-        assert result == {"integrations_processed": 0}
+        # Should process 1 SecurityHub integration
+        assert result == {"integrations_processed": 1}
+
+        # Verify the integration filter was called
         mock_integration_filter.assert_called_once_with(
             integrationproviderrelationship__provider_id=self.provider_id,
             enabled=True,
         )
-        # group should not be called since no integration tasks are created yet
-        mock_group.assert_not_called()
+
+        # Verify SecurityHub integrations were filtered
+        mock_integrations.filter.assert_called_once_with(
+            integration_type=Integration.IntegrationChoices.AWS_SECURITY_HUB
+        )
+
+        # Verify SecurityHub task was created with correct parameters
+        mock_security_hub_task.s.assert_called_once_with(
+            tenant_id=self.tenant_id,
+            provider_id=self.provider_id,
+            scan_id="test-scan-id",
+        )
+
+        # Verify group was called and job was executed
+        mock_group.assert_called_once_with([mock_task_signature])
+        mock_job.apply_async.assert_called_once()
 
     @patch("tasks.tasks.rls_transaction")
     @patch("tasks.tasks.Integration.objects.filter")
@@ -598,3 +636,33 @@ class TestCheckIntegrationsTask:
         mock_upload.assert_called_once_with(
             self.tenant_id, self.provider_id, output_directory
         )
+
+    @patch("tasks.tasks.upload_security_hub_integration")
+    def test_security_hub_integration_task_success(self, mock_upload):
+        """Test successful SecurityHub integration task execution."""
+        mock_upload.return_value = True
+        scan_id = "test-scan-123"
+
+        result = security_hub_integration_task(
+            tenant_id=self.tenant_id,
+            provider_id=self.provider_id,
+            scan_id=scan_id,
+        )
+
+        assert result is True
+        mock_upload.assert_called_once_with(self.tenant_id, self.provider_id, scan_id)
+
+    @patch("tasks.tasks.upload_security_hub_integration")
+    def test_security_hub_integration_task_failure(self, mock_upload):
+        """Test SecurityHub integration task handling failure."""
+        mock_upload.return_value = False
+        scan_id = "test-scan-123"
+
+        result = security_hub_integration_task(
+            tenant_id=self.tenant_id,
+            provider_id=self.provider_id,
+            scan_id=scan_id,
+        )
+
+        assert result is False
+        mock_upload.assert_called_once_with(self.tenant_id, self.provider_id, scan_id)