From b87e6d20d7b316f7adaac26c97e1ec40ab866e66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pedro=20Mart=C3=ADn?= Date: Wed, 30 Oct 2024 11:45:22 +0100 Subject: [PATCH] feat(s3): add test_connection method (#5332) --- .../aws/lib/s3/exceptions/exceptions.py | 77 ++++++++++++++++ prowler/providers/aws/lib/s3/s3.py | 90 +++++++++++++++++++ tests/providers/aws/lib/s3/s3_test.py | 31 +++++++ 3 files changed, 198 insertions(+) create mode 100644 prowler/providers/aws/lib/s3/exceptions/exceptions.py diff --git a/prowler/providers/aws/lib/s3/exceptions/exceptions.py b/prowler/providers/aws/lib/s3/exceptions/exceptions.py new file mode 100644 index 0000000000..b2d531f95d --- /dev/null +++ b/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 + ) diff --git a/prowler/providers/aws/lib/s3/s3.py b/prowler/providers/aws/lib/s3/s3.py index 2460a1b9c2..0bfa1aca06 100644 --- a/prowler/providers/aws/lib/s3/s3.py +++ b/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 diff --git a/tests/providers/aws/lib/s3/s3_test.py b/tests/providers/aws/lib/s3/s3_test.py index f60bc4cceb..2724e9749c 100644 --- a/tests/providers/aws/lib/s3/s3_test.py +++ b/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