feat(s3): add test_connection method (#5332)

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

@@ -0,0 +1,77 @@
+from prowler.exceptions.exceptions import ProwlerException
+
+
+# Exceptions codes from 6000 to 6999 are reserved for S3 exceptions
+class S3BaseException(ProwlerException):
+    """Base class for S3 exceptions."""
+
+    S3_ERROR_CODES = {
+        (6000, "S3TestConnectionError"): {
+            "message": "Failed to test connection to S3 bucket.",
+            "remediation": "Check the S3 bucket name and permissions.",
+        },
+        (6001, "S3InvalidBucketNameError"): {
+            "message": "The specified S3 bucket name is invalid.",
+            "remediation": "Check the S3 bucket name.",
+        },
+        (6002, "S3BucketAccessDeniedError"): {
+            "message": "Access to the specified S3 bucket is denied.",
+            "remediation": "Check the S3 bucket permissions.",
+        },
+        (6003, "S3ClientError"): {
+            "message": "An error occurred while interacting with the S3 client.",
+            "remediation": "Check the S3 client configuration.",
+        },
+        (6004, "S3IllegalLocationConstraintError"): {
+            "message": "The specified location constraint is not valid.",
+            "remediation": "Check the location constraint.",
+        },
+    }
+
+    def __init__(self, code, file=None, original_exception=None, message=None):
+        module = "S3"
+        error_info = self.S3_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 S3TestConnectionError(S3BaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            6000, file=file, original_exception=original_exception, message=message
+        )
+
+
+class S3InvalidBucketNameError(S3BaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            6001, file=file, original_exception=original_exception, message=message
+        )
+
+
+class S3BucketAccessDeniedError(S3BaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            6002, file=file, original_exception=original_exception, message=message
+        )
+
+
+class S3ClientError(S3BaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            6003, file=file, original_exception=original_exception, message=message
+        )
+
+
+class S3IllegalLocationConstraintError(S3BaseException):
+    def __init__(self, file=None, original_exception=None, message=None):
+        super().__init__(
+            6004, file=file, original_exception=original_exception, message=message
+        )

prowler/providers/aws/lib/s3/s3.py

@@ -1,10 +1,20 @@
+import tempfile
 from os import path
 from tempfile import NamedTemporaryFile
 
 from boto3 import Session
+from botocore import exceptions
 
 from prowler.lib.logger import logger
 from prowler.lib.outputs.output import Output
+from prowler.providers.aws.lib.s3.exceptions.exceptions import (
+    S3BucketAccessDeniedError,
+    S3ClientError,
+    S3IllegalLocationConstraintError,
+    S3InvalidBucketNameError,
+    S3TestConnectionError,
+)
+from prowler.providers.common.models import Connection
 
 
 class S3:
@@ -147,3 +157,83 @@ class S3:
                 f"{error.__class__.__name__}[{error.__traceback__.tb_lineno}] -- {error}"
             )
         return uploaded_objects
+
+    @staticmethod
+    def test_connection(
+        session, bucket_name: str, raise_on_exception: bool = True
+    ) -> Connection:
+        """
+        Test the connection to the S3 bucket.
+
+        Parameters:
+        - session: An instance of the `Session` class representing the AWS session.
+        - bucket_name: A string representing the name of the S3 bucket.
+        - raise_on_exception: A boolean indicating whether to raise an exception if the connection test fails.
+
+        Returns:
+        - A Connection object indicating the status of the connection test.
+
+        Raises:
+        - Exception: An exception indicating that the connection test failed.
+        """
+        try:
+            s3_client = session.client(__class__.__name__.lower())
+            if "s3://" in bucket_name:
+                bucket_name = bucket_name.removeprefix("s3://")
+            # Check for the bucket location
+            bucket_location = s3_client.get_bucket_location(Bucket=bucket_name)
+            if bucket_location["LocationConstraint"] == "EU":
+                bucket_location["LocationConstraint"] = "eu-west-1"
+            if (
+                bucket_location["LocationConstraint"] == ""
+                or bucket_location["LocationConstraint"] is None
+            ):
+                bucket_location["LocationConstraint"] = "us-east-1"
+
+            # If the bucket location is not the same as the session region, change the session region
+            if (
+                session.region_name != bucket_location["LocationConstraint"]
+                and bucket_location["LocationConstraint"] is not None
+            ):
+                s3_client = session.client(
+                    __class__.__name__.lower(),
+                    region_name=bucket_location["LocationConstraint"],
+                )
+            # Set a Temp file to upload
+            with tempfile.TemporaryFile() as temp_file:
+                temp_file.write(b"Test Prowler Connection")
+                temp_file.seek(0)
+                s3_client.upload_fileobj(
+                    temp_file, bucket_name, "test-prowler-connection.txt"
+                )
+
+            # Try to delete the file
+            s3_client.delete_object(
+                Bucket=bucket_name, Key="test-prowler-connection.txt"
+            )
+            return Connection(is_connected=True)
+
+        except exceptions.ClientError as client_error:
+            if raise_on_exception:
+                if (
+                    "specified bucket does not exist"
+                    in client_error.response["Error"]["Message"]
+                ):
+                    raise S3InvalidBucketNameError(original_exception=client_error)
+                elif (
+                    "IllegalLocationConstraintException"
+                    in client_error.response["Error"]["Message"]
+                ):
+                    raise S3IllegalLocationConstraintError(
+                        original_exception=client_error
+                    )
+                elif "AccessDenied" in client_error.response["Error"]["Code"]:
+                    raise S3BucketAccessDeniedError(original_exception=client_error)
+                else:
+                    raise S3ClientError(original_exception=client_error)
+            return Connection(is_connected=False, error=client_error)
+
+        except Exception as error:
+            if raise_on_exception:
+                raise S3TestConnectionError(original_exception=error)
+            return False

tests/providers/aws/lib/s3/s3_test.py

@@ -2,12 +2,14 @@ from os import path, remove
 from pathlib import Path
 
 import boto3
+import pytest
 from moto import mock_aws
 
 from prowler.lib.outputs.compliance.iso27001.iso27001_aws import AWSISO27001
 from prowler.lib.outputs.csv.csv import CSV
 from prowler.lib.outputs.html.html import HTML
 from prowler.lib.outputs.ocsf.ocsf import OCSF
+from prowler.providers.aws.lib.s3.exceptions.exceptions import S3InvalidBucketNameError
 from prowler.providers.aws.lib.s3.s3 import S3
 from tests.lib.outputs.compliance.fixtures import ISO27001_2013_AWS
 from tests.lib.outputs.fixtures.fixtures import generate_finding_output
@@ -314,3 +316,32 @@ class TestS3:
 
     def test_generate_subfolder_name_by_extension_json_ocsf(self):
         assert S3.generate_subfolder_name_by_extension(".ocsf.json") == "json-ocsf"
+
+    @mock_aws
+    def test_test_connection_S3(self):
+        current_session = boto3.session.Session(region_name=AWS_REGION_US_EAST_1)
+        s3_client = current_session.client("s3")
+        s3_client.create_bucket(Bucket=S3_BUCKET_NAME)
+        s3 = S3.test_connection(
+            session=current_session,
+            bucket_name=S3_BUCKET_NAME,
+        )
+        assert s3 is not None
+        assert s3.is_connected is True
+        assert s3.error is None
+
+    @mock_aws
+    def test_test_connection_S3_bucket_invalid_name(self):
+        current_session = boto3.session.Session(region_name=AWS_REGION_US_EAST_1)
+        s3_client = current_session.client("s3")
+
+        s3_client.create_bucket(Bucket=S3_BUCKET_NAME)
+        with pytest.raises(S3InvalidBucketNameError):
+            s3 = S3.test_connection(
+                session=current_session,
+                bucket_name="invalid_bucket",
+            )
+
+            assert s3 is not None
+            assert s3.is_connected is False
+            assert s3.error is not None