prowler/tests/providers/aws/services/batch/batch_service_test.py

from unittest.mock import patch

import botocore
from botocore.exceptions import ClientError

from prowler.providers.aws.services.batch.batch_service import Batch
from tests.providers.aws.utils import (
    AWS_REGION_EU_WEST_1,
    AWS_REGION_US_EAST_1,
    set_mocked_aws_provider,
)

make_api_call = botocore.client.BaseClient._make_api_call

COMPUTE_ENVIRONMENT_NAME = "test-compute-environment"
COMPUTE_ENVIRONMENT_ARN = f"arn:aws:batch:eu-west-1:123456789012:compute-environment/{COMPUTE_ENVIRONMENT_NAME}"


def default_api_call(self, operation_name, kwarg):
    """Fall back to the real call, stubbing the APIs a test does not exercise."""
    if operation_name == "DescribeComputeEnvironments":
        return {"computeEnvironments": []}
    return make_api_call(self, operation_name, kwarg)


def mock_make_api_call(self, operation_name, kwarg):
    if operation_name == "DescribeJobDefinitions":
        return {
            "jobDefinitions": [
                {
                    "jobDefinitionName": "test-batch-job",
                    "jobDefinitionArn": "arn:aws:batch:eu-west-1:123456789012:job-definition/test-batch-job:1",
                    "revision": 1,
                    "containerProperties": {
                        "image": "test-image:latest",
                        "command": ["python", "app.py"],
                        "environment": [
                            {"name": "DB_PASSWORD", "value": "pass-12343"},
                            {"name": "APP_NAME", "value": "myapp"},
                        ],
                    },
                }
            ]
        }
    return default_api_call(self, operation_name, kwarg)


def mock_generate_regional_clients(provider, service):
    regional_client = provider._session.current_session.client(
        service, region_name=AWS_REGION_EU_WEST_1
    )
    regional_client.region = AWS_REGION_EU_WEST_1
    return {AWS_REGION_EU_WEST_1: regional_client}


def mock_generate_multi_region_clients(provider, service):
    eu_west_1_client = provider._session.current_session.client(
        service, region_name=AWS_REGION_EU_WEST_1
    )
    eu_west_1_client.region = AWS_REGION_EU_WEST_1

    us_east_1_client = provider._session.current_session.client(
        service, region_name=AWS_REGION_US_EAST_1
    )
    us_east_1_client.region = AWS_REGION_US_EAST_1

    return {
        AWS_REGION_EU_WEST_1: eu_west_1_client,
        AWS_REGION_US_EAST_1: us_east_1_client,
    }


@patch(
    "prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients",
    new=mock_generate_regional_clients,
)
class Test_Batch_Service:
    def test_service(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)
        assert batch.service == "batch"

    def test_client(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)
        for reg_client in batch.regional_clients.values():
            assert reg_client.__class__.__name__ == "Batch"

    def test__get_session__(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)
        assert batch.session.__class__.__name__ == "Session"

    @patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
    def test_list_job_definitions(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)

        assert len(batch.job_definitions) == 1
        jd_arn = "arn:aws:batch:eu-west-1:123456789012:job-definition/test-batch-job:1"
        jd = batch.job_definitions[jd_arn]
        assert jd.name == "test-batch-job"
        assert jd.arn == jd_arn
        assert jd.revision == 1
        assert jd.region == AWS_REGION_EU_WEST_1

    @patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
    def test_describe_job_definitions(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)

        assert len(batch.job_definitions) == 1
        jd = list(batch.job_definitions.values())[0]
        assert jd.name == "test-batch-job"
        assert jd.container_properties.image == "test-image:latest"
        assert jd.container_properties.command == ["python", "app.py"]
        assert len(jd.container_properties.environment) == 2
        assert jd.container_properties.environment[0].name == "DB_PASSWORD"
        assert jd.container_properties.environment[0].value == "pass-12343"
        assert jd.container_properties.environment[1].name == "APP_NAME"
        assert jd.container_properties.environment[1].value == "myapp"

    @patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
    def test_no_job_definitions(self):
        def mock_make_api_call_empty(self, operation_name, kwarg):
            if operation_name == "DescribeJobDefinitions":
                return {"jobDefinitions": []}
            return default_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call",
            new=mock_make_api_call_empty,
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            batch = Batch(aws_provider)
            assert len(batch.job_definitions) == 0

    def test_job_definitions_are_loaded_for_analysis(self):
        describe_calls = []

        def counting_make_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeJobDefinitions":
                describe_calls.append(kwarg)
                return {
                    "jobDefinitions": [
                        {
                            "jobDefinitionName": f"job-{i}",
                            "jobDefinitionArn": f"arn:aws:batch:eu-west-1:123456789012:job-definition/job-{i}:{i}",
                            "revision": i,
                            "containerProperties": {
                                "image": "test-image:latest",
                                "environment": [],
                            },
                        }
                        for i in (3, 2, 1)
                    ]
                }
            return default_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call", new=counting_make_api_call
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            batch = Batch(aws_provider)

            assert [jd.revision for jd in batch.job_definitions.values()] == [3, 2, 1]
            assert len(describe_calls) == 1
            assert describe_calls[0].get("status") == "ACTIVE"

    def test_job_definition_limit_exposes_only_selected_resources(self):
        describe_calls = []

        def counting_make_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeJobDefinitions":
                describe_calls.append(kwarg)
                return {
                    "jobDefinitions": [
                        {
                            "jobDefinitionName": f"job-{i}",
                            "jobDefinitionArn": f"arn:aws:batch:eu-west-1:123456789012:job-definition/job-{i}:{i}",
                            "revision": i,
                            "containerProperties": {
                                "image": "test-image:latest",
                                "environment": [],
                            },
                        }
                        for i in (3, 2, 1)
                    ]
                }
            return default_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call", new=counting_make_api_call
        ):
            aws_provider = set_mocked_aws_provider(
                [AWS_REGION_EU_WEST_1],
                audit_config={"max_batch_job_definitions": 2},
            )
            batch = Batch(aws_provider)

            assert [jd.revision for jd in batch.job_definitions.values()] == [3, 2]
            assert len(describe_calls) == 1

    def test_describe_compute_environments(self):
        def compute_environments_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeComputeEnvironments":
                return {
                    "computeEnvironments": [
                        {
                            "computeEnvironmentName": COMPUTE_ENVIRONMENT_NAME,
                            "computeEnvironmentArn": COMPUTE_ENVIRONMENT_ARN,
                            "computeResources": {
                                "securityGroupIds": ["sg-1", "sg-2"],
                                "subnets": ["subnet-1"],
                            },
                        }
                    ]
                }
            return mock_make_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call",
            new=compute_environments_api_call,
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            batch = Batch(aws_provider)

            assert len(batch.compute_environments) == 1
            compute_environment = batch.compute_environments[COMPUTE_ENVIRONMENT_ARN]
            assert compute_environment.name == COMPUTE_ENVIRONMENT_NAME
            assert compute_environment.arn == COMPUTE_ENVIRONMENT_ARN
            assert compute_environment.region == AWS_REGION_EU_WEST_1
            assert compute_environment.security_groups == ["sg-1", "sg-2"]
            assert compute_environment.subnets == ["subnet-1"]
            assert batch.security_groups_in_use == {"sg-1", "sg-2"}

    @patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
    def test_no_compute_environments(self):
        aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
        batch = Batch(aws_provider)

        assert len(batch.compute_environments) == 0
        assert batch.security_groups_in_use == set()
        assert batch.compute_environment_lookup_failed_regions == set()

    def test_compute_environments_access_denied(self):
        def access_denied_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeComputeEnvironments":
                raise ClientError(
                    {
                        "Error": {
                            "Code": "AccessDeniedException",
                            "Message": "User is not authorized to perform: batch:DescribeComputeEnvironments",
                        }
                    },
                    operation_name,
                )
            return mock_make_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call", new=access_denied_api_call
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            batch = Batch(aws_provider)

            # The associations are unknown, not absent
            assert batch.compute_environments == {}
            assert batch.security_groups_in_use == set()
            assert batch.compute_environment_lookup_failed_regions == {
                AWS_REGION_EU_WEST_1
            }
            # Unrelated discovery is unaffected
            assert len(batch.job_definitions) == 1

    def test_compute_environment_without_compute_resources(self):
        def unmanaged_api_call(self, operation_name, kwarg):
            # UNMANAGED compute environments carry no `computeResources` block
            if operation_name == "DescribeComputeEnvironments":
                return {
                    "computeEnvironments": [
                        {
                            "computeEnvironmentName": COMPUTE_ENVIRONMENT_NAME,
                            "computeEnvironmentArn": COMPUTE_ENVIRONMENT_ARN,
                        }
                    ]
                }
            return mock_make_api_call(self, operation_name, kwarg)

        with patch("botocore.client.BaseClient._make_api_call", new=unmanaged_api_call):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            batch = Batch(aws_provider)

            assert len(batch.compute_environments) == 1
            compute_environment = batch.compute_environments[COMPUTE_ENVIRONMENT_ARN]
            assert compute_environment.security_groups == []
            assert compute_environment.subnets == []
            assert batch.security_groups_in_use == set()

    def test_audit_resources_filters_compute_environments(self):
        def compute_environments_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeComputeEnvironments":
                return {
                    "computeEnvironments": [
                        {
                            "computeEnvironmentName": f"compute-environment-{i}",
                            "computeEnvironmentArn": f"arn:aws:batch:eu-west-1:123456789012:compute-environment/compute-environment-{i}",
                            "computeResources": {
                                "securityGroupIds": [f"sg-{i}"],
                                "subnets": ["subnet-1"],
                            },
                        }
                        for i in (1, 2)
                    ]
                }
            return mock_make_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call",
            new=compute_environments_api_call,
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            aws_provider._audit_resources = [
                "arn:aws:batch:eu-west-1:123456789012:compute-environment/compute-environment-1"
            ]
            batch = Batch(aws_provider)

            assert list(batch.compute_environments) == [
                "arn:aws:batch:eu-west-1:123456789012:compute-environment/compute-environment-1"
            ]
            assert batch.security_groups_in_use == {"sg-1"}

    def test_audit_resources_filters_job_definitions(self):
        def counting_make_api_call(self, operation_name, kwarg):
            if operation_name == "DescribeJobDefinitions":
                return {
                    "jobDefinitions": [
                        {
                            "jobDefinitionName": f"job-{i}",
                            "jobDefinitionArn": f"arn:aws:batch:eu-west-1:123456789012:job-definition/job-{i}:{i}",
                            "revision": i,
                            "containerProperties": {
                                "image": "test-image:latest",
                                "environment": [],
                            },
                        }
                        for i in (1, 2)
                    ]
                }
            return default_api_call(self, operation_name, kwarg)

        with patch(
            "botocore.client.BaseClient._make_api_call", new=counting_make_api_call
        ):
            aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
            aws_provider._audit_resources = [
                "arn:aws:batch:eu-west-1:123456789012:job-definition/job-2:2"
            ]
            batch = Batch(aws_provider)

            assert list(batch.job_definitions.keys()) == [
                "arn:aws:batch:eu-west-1:123456789012:job-definition/job-2:2"
            ]