feat(SecurityHub): add test_connection method (#5350)

Co-authored-by: pedrooot <pedromarting3@gmail.com>

prowler/__main__.py

@@ -621,7 +621,11 @@ def prowler():
             )
 
             security_hub_regions = (
-                global_provider.get_available_aws_service_regions("securityhub")
+                global_provider.get_available_aws_service_regions(
+                    "securityhub",
+                    global_provider.identity.partition,
+                    global_provider.identity.audited_regions,
+                )
                 if not global_provider.identity.audited_regions
                 else global_provider.identity.audited_regions
             )

prowler/providers/aws/aws_provider.py

@@ -629,7 +629,9 @@ class AwsProvider(Provider):
         """
         try:
             regional_clients = {}
-            service_regions = self.get_available_aws_service_regions(service)
+            service_regions = AwsProvider.get_available_aws_service_regions(
+                service, self._identity.partition, self._identity.audited_regions
+            )
 
             # Get the regions enabled for the account and get the intersection with the service available regions
             if self._enabled_regions:
@@ -650,14 +652,26 @@ class AwsProvider(Provider):
                 f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
             )
 
-    def get_available_aws_service_regions(self, service: str) -> set:
+    @staticmethod
+    def get_available_aws_service_regions(
+        service: str, partition: str = "aws", audited_regions: set = None
+    ) -> set:
+        """
+        get_available_aws_service_regions returns the available regions for the given service and partition.
+
+        Arguments:
+            - service: The AWS service name.
+            - partition: The AWS partition name. Default is "aws".
+            - audited_regions: A set of regions to audit. Default is None.
+
+        Returns:
+            - A set of strings representing the available regions for the given service and partition.
+        """
         data = read_aws_regions_file()
-        json_regions = set(
-            data["services"][service]["regions"][self._identity.partition]
-        )
-        if self._identity.audited_regions:
+        json_regions = set(data["services"][service]["regions"][partition])
+        if audited_regions:
             # Get common regions between input and json
-            regions = json_regions.intersection(self._identity.audited_regions)
+            regions = json_regions.intersection(audited_regions)
         else:  # Get all regions from json of the service and partition
             regions = json_regions
         return regions
@@ -793,7 +807,9 @@ class AwsProvider(Provider):
     def get_default_region(self, service: str) -> str:
         """get_default_region returns the default region based on the profile and audited service regions"""
         try:
-            service_regions = self.get_available_aws_service_regions(service)
+            service_regions = AwsProvider.get_available_aws_service_regions(
+                service, self._identity.partition, self._identity.audited_regions
+            )
             default_region = self.get_global_region()
             # global region of the partition when all regions are audited and there is no profile region
             if self._identity.profile_region in service_regions:

prowler/providers/aws/lib/security_hub/exceptions/exceptions.py

@@ -0,0 +1,44 @@
+from prowler.exceptions.exceptions import ProwlerException
+
+
+# Exceptions codes from 7000 to 7999 are reserved for Security Hub exceptions
+class SecurityHubBaseException(ProwlerException):
+    """Base class for Security Hub exceptions."""
+
+    SECURITYHUB_ERROR_CODES = {
+        (7000, "SecurityHubNoEnabledRegionsError"): {
+            "message": "No regions were found to with the Security Hub integration enabled.",
+            "remediation": "Please check the connection settings and permissions and try again.",
+        },
+        (7001, "SecurityHubInvalidRegionError"): {
+            "message": "Given region has not Security Hub enabled.",
+            "remediation": "Please provide a valid region.",
+        },
+    }
+
+    def __init__(self, code, file=None, original_exception=None, message=None):
+        module = "SecurityHub"
+        error_info = self.SECURITYHUB_ERROR_CODES.get((code, self.__class__.__name__))
+        if message:
+            error_info["message"] = message
+        super().__init__(
+            code=code,
+            source=module,
+            file=file,
+            original_exception=original_exception,
+            error_info=error_info,
+        )
+
+
+class SecurityHubNoEnabledRegionsError(SecurityHubBaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            7000, file=file, original_exception=original_exception, message=message
+        )
+
+
+class SecurityHubInvalidRegionError(SecurityHubBaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            7001, file=file, original_exception=original_exception, message=message
+        )

prowler/providers/aws/lib/security_hub/security_hub.py

@@ -1,14 +1,35 @@
+from dataclasses import dataclass
+
 from boto3 import Session
 from botocore.client import ClientError
 
 from prowler.config.config import timestamp_utc
 from prowler.lib.logger import logger
 from prowler.lib.outputs.asff.asff import AWSSecurityFindingFormat
+from prowler.providers.aws.aws_provider import AwsProvider
+from prowler.providers.aws.lib.security_hub.exceptions.exceptions import (
+    SecurityHubInvalidRegionError,
+    SecurityHubNoEnabledRegionsError,
+)
+from prowler.providers.common.models import Connection
 
 SECURITY_HUB_INTEGRATION_NAME = "prowler/prowler"
 SECURITY_HUB_MAX_BATCH = 100
 
 
+@dataclass
+class SecurityHubConnection(Connection):
+    """
+    Represents a Security Hub connection object.
+    Attributes:
+        enabled_regions (set): Set of regions where Security Hub is enabled.
+        disabled_regions (set): Set of regions where Security Hub is disabled.
+    """
+
+    enabled_regions: set = None
+    disabled_regions: set = None
+
+
 class SecurityHub:
     """
     Class representing a SecurityHub object for managing findings and interactions with AWS Security Hub.
@@ -53,7 +74,10 @@ class SecurityHub:
 
         if aws_security_hub_available_regions:
             self._enabled_regions = self.verify_enabled_per_region(
-                aws_security_hub_available_regions
+                aws_security_hub_available_regions,
+                aws_session,
+                aws_account_id,
+                aws_partition,
             )
         if findings and self._enabled_regions:
             self._findings_per_region = self.filter(findings, send_only_fails)
@@ -103,9 +127,12 @@ class SecurityHub:
             )
         return findings_per_region
 
+    @staticmethod
     def verify_enabled_per_region(
-        self,
         aws_security_hub_available_regions: list[str],
+        session: Session,
+        aws_account_id: str,
+        aws_partition: str,
     ) -> dict[str, Session]:
         """
         Filters the given list of regions where AWS Security Hub is enabled and returns a dictionary containing the region and their boto3 client if the region and the Prowler integration is enabled.
@@ -123,13 +150,11 @@ class SecurityHub:
                     f"Checking if the {SECURITY_HUB_INTEGRATION_NAME} is enabled in the {region} region."
                 )
                 # Check if security hub is enabled in current region
-                security_hub_client = self._session.client(
-                    "securityhub", region_name=region
-                )
+                security_hub_client = session.client("securityhub", region_name=region)
                 security_hub_client.describe_hub()
 
                 # Check if Prowler integration is enabled in Security Hub
-                security_hub_prowler_integration_arn = f"arn:{self._aws_partition}:securityhub:{region}:{self._aws_account_id}:product-subscription/{SECURITY_HUB_INTEGRATION_NAME}"
+                security_hub_prowler_integration_arn = f"arn:{aws_partition}:securityhub:{region}:{aws_account_id}:product-subscription/{SECURITY_HUB_INTEGRATION_NAME}"
                 if security_hub_prowler_integration_arn not in str(
                     security_hub_client.list_enabled_products_for_import()
                 ):
@@ -137,7 +162,7 @@ class SecurityHub:
                         f"Security Hub is enabled in {region} but Prowler integration does not accept findings. More info: https://docs.prowler.cloud/en/latest/tutorials/aws/securityhub/"
                     )
                 else:
-                    enabled_regions[region] = self._session.client(
+                    enabled_regions[region] = session.client(
                         "securityhub", region_name=region
                     )
 
@@ -148,7 +173,7 @@ class SecurityHub:
                 error_message = client_error.response["Error"]["Message"]
                 if (
                     error_code == "InvalidAccessException"
-                    and f"Account {self._aws_account_id} is not subscribed to AWS Security Hub"
+                    and f"Account {aws_account_id} is not subscribed to AWS Security Hub"
                     in error_message
                 ):
                     logger.warning(
@@ -284,3 +309,104 @@ class SecurityHub:
                 f"{error.__class__.__name__} -- [{error.__traceback__.tb_lineno}]:{error} in region {region}"
             )
             return success_count
+
+    @staticmethod
+    def test_connection(
+        session: Session,
+        aws_account_id: str,
+        aws_partition: str,
+        regions: set = None,
+        raise_on_exception: bool = True,
+    ) -> SecurityHubConnection:
+        """
+        Test the connection to AWS Security Hub by checking if Security Hub is enabled in the provided region
+        and if the Prowler integration is active.
+
+        Args:
+            session (Session): AWS session to use for authentication.
+            regions (set): Set of regions to check for Security Hub integration.
+            aws_account_id (str): AWS account ID to check for Prowler integration.
+            aws_partition (str): AWS partition (e.g., aws, aws-cn, aws-us-gov).
+            raise_on_exception (bool): Whether to raise an exception if an error occurs.
+
+        Returns:
+            Connection: An object that contains the result of the test connection operation.
+                - is_connected (bool): Indicates whether the connection was successful.
+                - error (Exception): An exception object if an error occurs during the connection test.
+            enabled_regions (set): Set of regions where Security Hub is enabled.
+            disabled_regions (set): Set of regions where Security Hub is disabled.
+        """
+        try:
+            disabled_regions = set()
+            enabled_regions = set()
+
+            all_regions = AwsProvider.get_available_aws_service_regions(
+                service="securityhub", partition=aws_partition
+            )
+
+            enabled_regions = SecurityHub.verify_enabled_per_region(
+                aws_security_hub_available_regions=all_regions,
+                session=session,
+                aws_account_id=aws_account_id,
+                aws_partition=aws_partition,
+            ).keys()
+            disabled_regions = all_regions - enabled_regions
+            if regions:
+                if not any(region in enabled_regions for region in regions):
+                    logger.warning(
+                        f"Prowler integration is not enabled in regions: {regions - enabled_regions}."
+                    )
+                    invalid_region_error = SecurityHubInvalidRegionError(
+                        message="Given regions have not Security Hub enabled."
+                    )
+                    if raise_on_exception:
+                        raise invalid_region_error
+                    return SecurityHubConnection(
+                        is_connected=False,
+                        error=invalid_region_error,
+                        enabled_regions=enabled_regions,
+                        disabled_regions=disabled_regions,
+                    )
+                else:
+                    logger.info(
+                        f"Prowler integration is enabled in regions: {', '.join(regions)}."
+                    )
+                    return SecurityHubConnection(
+                        is_connected=True,
+                        error=None,
+                        enabled_regions=enabled_regions,
+                        disabled_regions=disabled_regions,
+                    )
+
+            if len(enabled_regions) == 0:
+                error_str = (
+                    "No regions found with the Security Hub integration enabled."
+                )
+                logger.warning(error_str)
+                no_enabled_regions_error = SecurityHubNoEnabledRegionsError(
+                    message=error_str
+                )
+                if raise_on_exception:
+                    raise no_enabled_regions_error
+                return SecurityHubConnection(
+                    is_connected=False,
+                    error=no_enabled_regions_error,
+                    enabled_regions=enabled_regions,
+                    disabled_regions=disabled_regions,
+                )
+            else:
+                logger.info(
+                    f"Security Hub is enabled in the following regions: {', '.join(enabled_regions)}."
+                )
+                return SecurityHubConnection(
+                    is_connected=True,
+                    error=None,
+                    enabled_regions=enabled_regions,
+                    disabled_regions=disabled_regions,
+                )
+
+        except Exception as error:
+            logger.error(
+                f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}]: {error}"
+            )
+            raise error

tests/providers/aws/aws_provider_test.py

@@ -978,9 +978,9 @@ aws:
                 }
             },
         ):
-            assert aws_provider.get_available_aws_service_regions("ec2") == {
-                AWS_REGION_US_EAST_1
-            }
+            assert aws_provider.get_available_aws_service_regions(
+                "ec2", "aws", {AWS_REGION_US_EAST_1}
+            ) == {AWS_REGION_US_EAST_1}
 
     @mock_aws
     def test_get_available_aws_service_regions_with_all_regions_audited(self):
@@ -1017,7 +1017,9 @@ aws:
                 }
             },
         ):
-            assert len(aws_provider.get_available_aws_service_regions("ec2")) == 17
+            assert (
+                len(aws_provider.get_available_aws_service_regions("ec2", "aws")) == 17
+            )
 
     @mock_aws
     def test_get_tagged_resources(self):

tests/providers/aws/lib/security_hub/security_hub_test.py

@@ -7,6 +7,10 @@ from botocore.client import ClientError
 from mock import patch
 
 from prowler.lib.outputs.asff.asff import ASFF
+from prowler.providers.aws.lib.security_hub.exceptions.exceptions import (
+    SecurityHubInvalidRegionError,
+    SecurityHubNoEnabledRegionsError,
+)
 from prowler.providers.aws.lib.security_hub.security_hub import SecurityHub
 from tests.lib.outputs.fixtures.fixtures import generate_finding_output
 from tests.providers.aws.utils import (
@@ -404,3 +408,116 @@ class TestSecurityHub:
         )
 
         assert security_hub.batch_send_to_security_hub() == 2
+
+    @patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
+    def test_security_hub_test_connection_success(self):
+        session_mock = session.Session(region_name=AWS_REGION_EU_WEST_1)
+
+        # Test successful connection
+        connection = SecurityHub.test_connection(
+            session=session_mock,
+            regions={AWS_REGION_EU_WEST_1},
+            aws_account_id=AWS_ACCOUNT_NUMBER,
+            aws_partition=AWS_COMMERCIAL_PARTITION,
+            raise_on_exception=False,
+        )
+
+        assert connection.is_connected is True
+        assert connection.error is None
+
+    @patch("prowler.providers.aws.lib.security_hub.security_hub.Session.client")
+    def test_security_hub_test_connection_invalid_access_exception(
+        self, mock_security_hub_client
+    ):
+        # Mock an InvalidAccessException
+        error_message = f"Account {AWS_ACCOUNT_NUMBER} is not subscribed to AWS Security Hub in region {AWS_REGION_EU_WEST_1}"
+        error_code = "InvalidAccessException"
+        error_response = {
+            "Error": {
+                "Code": error_code,
+                "Message": error_message,
+            }
+        }
+        operation_name = "DescribeHub"
+        mock_security_hub_client.side_effect = ClientError(
+            error_response, operation_name
+        )
+
+        session_mock = session.Session(region_name=AWS_REGION_EU_WEST_1)
+
+        # Test connection failure due to invalid access
+        connection = SecurityHub.test_connection(
+            session=session_mock,
+            regions={AWS_REGION_EU_WEST_1},
+            aws_account_id=AWS_ACCOUNT_NUMBER,
+            aws_partition=AWS_COMMERCIAL_PARTITION,
+            raise_on_exception=False,
+        )
+
+        assert connection.is_connected is False
+        assert isinstance(connection.error, SecurityHubInvalidRegionError)
+
+    @patch("prowler.providers.aws.lib.security_hub.security_hub.Session.client")
+    def test_security_hub_test_connection_prowler_not_subscribed(
+        self, mock_security_hub_client
+    ):
+        # Mock successful Security Hub but no Prowler subscription
+        mock_security_hub_client.describe_hub.return_value = {}
+        mock_security_hub_client.list_enabled_products_for_import.return_value = {
+            "ProductSubscriptions": []
+        }
+
+        session_mock = session.Session(region_name=AWS_REGION_EU_WEST_1)
+
+        # Test connection failure due to missing Prowler subscription
+        connection = SecurityHub.test_connection(
+            session=session_mock,
+            regions={AWS_REGION_EU_WEST_1},
+            aws_account_id=AWS_ACCOUNT_NUMBER,
+            aws_partition=AWS_COMMERCIAL_PARTITION,
+            raise_on_exception=False,
+        )
+
+        assert connection.is_connected is False
+        assert isinstance(connection.error, SecurityHubInvalidRegionError)
+
+    @patch("prowler.providers.aws.lib.security_hub.security_hub.Session.client")
+    def test_security_hub_test_connection_unexpected_exception(
+        self, mock_security_hub_client
+    ):
+        # Mock unexpected exception
+        mock_security_hub_client.side_effect = Exception("Unexpected error")
+
+        session_mock = session.Session(region_name=AWS_REGION_EU_WEST_1)
+
+        # Test connection failure due to an unexpected exception
+        connection = SecurityHub.test_connection(
+            session=session_mock,
+            regions={AWS_REGION_EU_WEST_1},
+            aws_account_id=AWS_ACCOUNT_NUMBER,
+            aws_partition=AWS_COMMERCIAL_PARTITION,
+            raise_on_exception=False,
+        )
+
+        assert connection.is_connected is False
+        assert isinstance(connection.error, SecurityHubInvalidRegionError)
+
+    @patch("prowler.providers.aws.lib.security_hub.security_hub.Session.client")
+    def test_security_hub_test_connection_no_regions_enabled(
+        self, mock_security_hub_client
+    ):
+        # Mock unexpected exception
+        mock_security_hub_client.side_effect = Exception("Unexpected error")
+
+        session_mock = session.Session(region_name=AWS_REGION_EU_WEST_1)
+
+        # Test connection failure due to an unexpected exception
+        connection = SecurityHub.test_connection(
+            session=session_mock,
+            aws_account_id=AWS_ACCOUNT_NUMBER,
+            aws_partition=AWS_COMMERCIAL_PARTITION,
+            raise_on_exception=False,
+        )
+
+        assert connection.is_connected is False
+        assert isinstance(connection.error, SecurityHubNoEnabledRegionsError)