mirror of
https://github.com/prowler-cloud/prowler.git
synced 2026-07-24 21:11:53 +00:00
Oleksandr_Sanin
bba594a1db
Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com> Signed-off-by: Oleksandr Yizchak Sanin <alexaaander.sanin@gmail.com> Co-authored-by: Daniel Barranquero <danielbo2001@gmail.com>
436 lines
18 KiB
Python
436 lines
18 KiB
Python
from unittest.mock import MagicMock, patch
|
|
from uuid import uuid4
|
|
|
|
import botocore
|
|
|
|
from prowler.providers.aws.services.sagemaker.sagemaker_service import (
|
|
Model,
|
|
SageMaker,
|
|
)
|
|
from tests.providers.aws.utils import (
|
|
AWS_ACCOUNT_NUMBER,
|
|
AWS_REGION_EU_WEST_1,
|
|
set_mocked_aws_provider,
|
|
)
|
|
|
|
test_model_package_group_name = "test-model-package-group"
|
|
test_model_package_group_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:model-package-group/{test_model_package_group_name}"
|
|
test_model_package_name = "test-model-package"
|
|
test_model_package_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:model-package/{test_model_package_name}/1"
|
|
|
|
test_notebook_instance = "test-notebook-instance"
|
|
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:notebook-instance/{test_notebook_instance}"
|
|
test_model = "test-model"
|
|
test_arn_model = (
|
|
f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:model/{test_model}"
|
|
)
|
|
test_training_job = "test-training-job"
|
|
test_arn_training_job = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:training-job/{test_model}"
|
|
subnet_id = "subnet-" + str(uuid4())
|
|
kms_key_id = str(uuid4())
|
|
endpoint_config_name = "endpoint-config-test"
|
|
endpoint_config_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:endpoint-config/{endpoint_config_name}"
|
|
prod_variant_name = "Variant1"
|
|
test_domain_name = "test-domain"
|
|
test_domain_id = "d-testdomain123"
|
|
test_domain_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:domain/{test_domain_id}"
|
|
test_sso_instance_id = "app-test-instance-id"
|
|
test_sso_application_arn = (
|
|
f"arn:aws:sso::{AWS_ACCOUNT_NUMBER}:application/sagemaker/apl-test"
|
|
)
|
|
|
|
make_api_call = botocore.client.BaseClient._make_api_call
|
|
|
|
|
|
def mock_make_api_call(self, operation_name, kwarg):
|
|
if operation_name == "ListNotebookInstances":
|
|
return {
|
|
"NotebookInstances": [
|
|
{
|
|
"NotebookInstanceName": test_notebook_instance,
|
|
"NotebookInstanceArn": notebook_instance_arn,
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "ListModels":
|
|
return {
|
|
"Models": [
|
|
{
|
|
"ModelName": test_model,
|
|
"ModelArn": test_arn_model,
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "ListTrainingJobs":
|
|
return {
|
|
"TrainingJobSummaries": [
|
|
{
|
|
"TrainingJobName": test_training_job,
|
|
"TrainingJobArn": test_arn_training_job,
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "DescribeNotebookInstance":
|
|
return {
|
|
"SubnetId": subnet_id,
|
|
"KmsKeyId": kms_key_id,
|
|
"DirectInternetAccess": "Enabled",
|
|
"RootAccess": "Enabled",
|
|
}
|
|
if operation_name == "DescribeModel":
|
|
return {
|
|
"VpcConfig": {
|
|
"Subnets": [
|
|
subnet_id,
|
|
]
|
|
},
|
|
"EnableNetworkIsolation": True,
|
|
}
|
|
if operation_name == "DescribeTrainingJob":
|
|
return {
|
|
"ResourceConfig": {
|
|
"VolumeKmsKeyId": kms_key_id,
|
|
},
|
|
"VpcConfig": {
|
|
"Subnets": [
|
|
subnet_id,
|
|
]
|
|
},
|
|
"EnableNetworkIsolation": True,
|
|
"EnableInterContainerTrafficEncryption": True,
|
|
}
|
|
if operation_name == "ListModelPackageGroups":
|
|
return {
|
|
"ModelPackageGroupSummaryList": [
|
|
{
|
|
"ModelPackageGroupName": test_model_package_group_name,
|
|
"ModelPackageGroupArn": test_model_package_group_arn,
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "ListModelPackages":
|
|
return {
|
|
"ModelPackageSummaryList": [
|
|
{
|
|
"ModelPackageName": test_model_package_name,
|
|
"ModelPackageArn": test_model_package_arn,
|
|
"ModelApprovalStatus": "Approved",
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "ListTags":
|
|
return {
|
|
"Tags": [
|
|
{"Key": "test", "Value": "test"},
|
|
],
|
|
}
|
|
if operation_name == "ListEndpointConfigs":
|
|
return {
|
|
"EndpointConfigs": [
|
|
{
|
|
"EndpointConfigName": endpoint_config_name,
|
|
"EndpointConfigArn": endpoint_config_arn,
|
|
},
|
|
],
|
|
}
|
|
if operation_name == "DescribeEndpointConfig":
|
|
return {
|
|
"ProductionVariants": [
|
|
{
|
|
"VariantName": prod_variant_name,
|
|
"InitialInstanceCount": 5,
|
|
},
|
|
{
|
|
"VariantName": "Variant2",
|
|
"InitialInstanceCount": 2,
|
|
},
|
|
]
|
|
}
|
|
if operation_name == "ListDomains":
|
|
return {
|
|
"Domains": [
|
|
{
|
|
"DomainId": test_domain_id,
|
|
"DomainName": test_domain_name,
|
|
"DomainArn": test_domain_arn,
|
|
},
|
|
],
|
|
}
|
|
if operation_name == "DescribeDomain":
|
|
return {
|
|
"DomainId": test_domain_id,
|
|
"DomainName": test_domain_name,
|
|
"DomainArn": test_domain_arn,
|
|
"AuthMode": "SSO",
|
|
"SingleSignOnManagedApplicationInstanceId": test_sso_instance_id,
|
|
"SingleSignOnApplicationArn": test_sso_application_arn,
|
|
}
|
|
|
|
return make_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}
|
|
|
|
|
|
@patch("botocore.client.BaseClient._make_api_call", new=mock_make_api_call)
|
|
@patch(
|
|
"prowler.providers.aws.aws_provider.AwsProvider.generate_regional_clients",
|
|
new=mock_generate_regional_clients,
|
|
)
|
|
class Test_SageMaker_Service:
|
|
# Test SageMaker Service
|
|
def test_service(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert sagemaker.service == "sagemaker"
|
|
|
|
# Test SageMaker client
|
|
def test_client(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
for reg_client in sagemaker.regional_clients.values():
|
|
assert reg_client.__class__.__name__ == "SageMaker"
|
|
|
|
# Test SageMaker session
|
|
def test__get_session__(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert sagemaker.session.__class__.__name__ == "Session"
|
|
|
|
# Test SageMaker list notebook instances
|
|
def test_list_notebook_instances(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_notebook_instances) == 1
|
|
assert sagemaker.sagemaker_notebook_instances[0].name == test_notebook_instance
|
|
assert sagemaker.sagemaker_notebook_instances[0].arn == notebook_instance_arn
|
|
assert sagemaker.sagemaker_notebook_instances[0].region == AWS_REGION_EU_WEST_1
|
|
assert sagemaker.sagemaker_notebook_instances[0].tags == [
|
|
{"Key": "test", "Value": "test"},
|
|
]
|
|
|
|
# Test SageMaker list models
|
|
def test_list_models(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_models) == 1
|
|
assert sagemaker.sagemaker_models[0].name == test_model
|
|
assert sagemaker.sagemaker_models[0].arn == test_arn_model
|
|
assert sagemaker.sagemaker_models[0].region == AWS_REGION_EU_WEST_1
|
|
assert sagemaker.sagemaker_models[0].tags == [
|
|
{"Key": "test", "Value": "test"},
|
|
]
|
|
|
|
# Test SageMaker list training jobs
|
|
def test_list_training_jobs(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_training_jobs) == 1
|
|
assert sagemaker.sagemaker_training_jobs[0].name == test_training_job
|
|
assert sagemaker.sagemaker_training_jobs[0].arn == test_arn_training_job
|
|
assert sagemaker.sagemaker_training_jobs[0].region == AWS_REGION_EU_WEST_1
|
|
assert sagemaker.sagemaker_training_jobs[0].tags == [
|
|
{"Key": "test", "Value": "test"},
|
|
]
|
|
|
|
# Test SageMaker describe notebook instance
|
|
def test_describe_notebook_instance(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_notebook_instances) == 1
|
|
assert sagemaker.sagemaker_notebook_instances[0].root_access
|
|
assert sagemaker.sagemaker_notebook_instances[0].subnet_id == subnet_id
|
|
assert sagemaker.sagemaker_notebook_instances[0].direct_internet_access
|
|
assert sagemaker.sagemaker_notebook_instances[0].kms_key_id == kms_key_id
|
|
|
|
# Test SageMaker describe model
|
|
def test_describe_model(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_models) == 1
|
|
assert sagemaker.sagemaker_models[0].network_isolation
|
|
assert sagemaker.sagemaker_models[0].vpc_config_subnets == [subnet_id]
|
|
|
|
# Test SageMaker describe training jobs
|
|
def test_describe_training_jobs(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_training_jobs) == 1
|
|
assert sagemaker.sagemaker_training_jobs[0].container_traffic_encryption
|
|
assert sagemaker.sagemaker_training_jobs[0].network_isolation
|
|
assert sagemaker.sagemaker_training_jobs[0].volume_kms_key_id == kms_key_id
|
|
assert sagemaker.sagemaker_training_jobs[0].vpc_config_subnets == [subnet_id]
|
|
|
|
# Test SageMaker list endpoint configs
|
|
def test_list_endpoint_configs(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.endpoint_configs) == 1
|
|
assert (
|
|
sagemaker.endpoint_configs[endpoint_config_arn].name == endpoint_config_name
|
|
)
|
|
assert (
|
|
sagemaker.endpoint_configs[endpoint_config_arn].arn == endpoint_config_arn
|
|
)
|
|
assert (
|
|
sagemaker.endpoint_configs[endpoint_config_arn].region
|
|
== AWS_REGION_EU_WEST_1
|
|
)
|
|
assert sagemaker.sagemaker_notebook_instances[0].tags == [
|
|
{"Key": "test", "Value": "test"},
|
|
]
|
|
|
|
# Test SageMaker describe training jobs
|
|
def test_describe_endpoint_configs(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.endpoint_configs) == 1
|
|
assert sagemaker.endpoint_configs[endpoint_config_arn].production_variants
|
|
for prod_variant in sagemaker.endpoint_configs[
|
|
endpoint_config_arn
|
|
].production_variants:
|
|
if prod_variant.name == prod_variant_name:
|
|
assert prod_variant.initial_instance_count == 5
|
|
else:
|
|
assert prod_variant.initial_instance_count == 2
|
|
|
|
# Test SageMaker list domains
|
|
def test_list_domains(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_domains) == 1
|
|
assert sagemaker.sagemaker_domains[0].domain_id == test_domain_id
|
|
assert sagemaker.sagemaker_domains[0].name == test_domain_name
|
|
assert sagemaker.sagemaker_domains[0].arn == test_domain_arn
|
|
assert sagemaker.sagemaker_domains[0].region == AWS_REGION_EU_WEST_1
|
|
|
|
# Test SageMaker describe domain
|
|
def test_describe_domain(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_domains) == 1
|
|
assert sagemaker.sagemaker_domains[0].auth_mode == "SSO"
|
|
assert (
|
|
sagemaker.sagemaker_domains[
|
|
0
|
|
].single_sign_on_managed_application_instance_id
|
|
== test_sso_instance_id
|
|
)
|
|
assert (
|
|
sagemaker.sagemaker_domains[0].single_sign_on_application_arn
|
|
== test_sso_application_arn
|
|
)
|
|
|
|
# Test SageMaker _list_tags_for_resource
|
|
def test_list_tags_for_resource_calls_client(self):
|
|
"""Test that _list_tags_for_resource calls the correct AWS client and updates the resource."""
|
|
# Mock audit info
|
|
audit_info = MagicMock()
|
|
audit_info.audited_partition = "aws"
|
|
audit_info.audited_account = AWS_ACCOUNT_NUMBER
|
|
audit_info.audit_resources = None
|
|
|
|
# Mock regional client
|
|
regional_client = MagicMock()
|
|
regional_client.region = AWS_REGION_EU_WEST_1
|
|
regional_client.list_tags.return_value = {
|
|
"Tags": [{"Key": "foo", "Value": "bar"}]
|
|
}
|
|
|
|
# Create service instance (mocking init to avoid full setup)
|
|
with patch.object(SageMaker, "__init__", return_value=None):
|
|
sagemaker_service = SageMaker(audit_info)
|
|
sagemaker_service.regional_clients = {AWS_REGION_EU_WEST_1: regional_client}
|
|
sagemaker_service.audit_info = audit_info
|
|
|
|
# Create a mock resource
|
|
resource = Model(
|
|
name="test-model",
|
|
region=AWS_REGION_EU_WEST_1,
|
|
arn=f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:model/test-model",
|
|
)
|
|
|
|
# Execute method under test
|
|
sagemaker_service._list_tags_for_resource(resource)
|
|
|
|
# Verification
|
|
regional_client.list_tags.assert_called_once_with(ResourceArn=resource.arn)
|
|
assert len(resource.tags) == 1
|
|
assert resource.tags[0]["Key"] == "foo"
|
|
assert resource.tags[0]["Value"] == "bar"
|
|
|
|
# Test SageMaker parallel tag listing
|
|
def test_init_calls_threading_for_tags(self):
|
|
"""Test that __init__ calls __threading_call__ for tag listing for each resource type."""
|
|
audit_info = MagicMock()
|
|
audit_info.audited_partition = "aws"
|
|
audit_info.audited_account = AWS_ACCOUNT_NUMBER
|
|
|
|
# We mock __threading_call__ to verify it is called with the right arguments
|
|
with patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker.__threading_call__"
|
|
) as mock_threading_call:
|
|
# We also need to mock the other methods called in init to avoid errors
|
|
with (
|
|
patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker._list_notebook_instances"
|
|
),
|
|
patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker._list_models"
|
|
),
|
|
patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker._list_training_jobs"
|
|
),
|
|
patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker._list_endpoint_configs"
|
|
),
|
|
patch(
|
|
"prowler.providers.aws.services.sagemaker.sagemaker_service.SageMaker._list_domains"
|
|
),
|
|
):
|
|
sagemaker_service = SageMaker(audit_info)
|
|
|
|
# Check that __threading_call__ was called for _list_tags_for_resource
|
|
# (one for each resource type: models, notebooks, training jobs, processing jobs, endpoint configs, domains)
|
|
tag_calls = [
|
|
c
|
|
for c in mock_threading_call.call_args_list
|
|
if c[0][0] == sagemaker_service._list_tags_for_resource
|
|
]
|
|
assert len(tag_calls) == 6
|
|
|
|
# Test SageMaker list model package groups
|
|
def test_list_model_package_groups(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert len(sagemaker.sagemaker_model_registries) == 1
|
|
registry = sagemaker.sagemaker_model_registries[0]
|
|
assert registry.region == AWS_REGION_EU_WEST_1
|
|
assert registry.has_groups is True
|
|
assert registry.has_approved_packages is True
|
|
|
|
def test_list_model_package_groups_access_denied(self):
|
|
aws_provider = set_mocked_aws_provider([AWS_REGION_EU_WEST_1])
|
|
|
|
def mock_access_denied(self, operation_name, kwarg):
|
|
if operation_name == "ListModelPackageGroups":
|
|
raise botocore.exceptions.ClientError(
|
|
{
|
|
"Error": {
|
|
"Code": "AccessDeniedException",
|
|
"Message": "User is not authorized to perform sagemaker:ListModelPackageGroups",
|
|
}
|
|
},
|
|
"ListModelPackageGroups",
|
|
)
|
|
return make_api_call(self, operation_name, kwarg)
|
|
|
|
with patch("botocore.client.BaseClient._make_api_call", new=mock_access_denied):
|
|
sagemaker = SageMaker(aws_provider)
|
|
assert sagemaker.sagemaker_model_registries == []
|