ghndrx/prowler
1
0
Fork 0
mirror of https://github.com/ghndrx/prowler.git synced 2026-02-10 06:45:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(audit_info): refactor sagemaker (#3135)

Browse Source
This commit is contained in:
Nacho Rivera
2023-12-05 11:40:19 +01:00
committed by GitHub
parent 6a901bb7de
commit 50ef2729e6
11 changed files with 95 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/providers/aws/services/sagemaker/sagemaker_models_network_isolation_enabled/sagemaker_models_network_isolation_enabled_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import Model
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
subnet_id = "subnet-" + str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_models_network_isolation_enabled:
Model(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
network_isolation=True,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_models_network_isolation_enabled:
Model(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
network_isolation=False,
)
)

13
tests/providers/aws/services/sagemaker/sagemaker_models_vpc_settings_configured/sagemaker_models_vpc_settings_configured_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import Model
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
subnet_id = "subnet-" + str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_models_vpc_settings_configured:
Model(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
vpc_config_subnets=[subnet_id],
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_models_vpc_settings_configured:
Model(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(

13
tests/providers/aws/services/sagemaker/sagemaker_notebook_instance_encryption_enabled/sagemaker_notebook_instance_encryption_enabled_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import NotebookInstance
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
kms_key = str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_notebook_instance_encryption_enabled:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
kms_key_id=kms_key,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_notebook_instance_encryption_enabled:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(

13
tests/providers/aws/services/sagemaker/sagemaker_notebook_instance_root_access_disabled/sagemaker_notebook_instance_root_access_disabled_test.py
View File

@@ -2,12 +2,13 @@ from re import search
from unittest import mock
from prowler.providers.aws.services.sagemaker.sagemaker_service import NotebookInstance
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
class Test_sagemaker_notebook_instance_root_access_disabled:
@@ -33,7 +34,7 @@ class Test_sagemaker_notebook_instance_root_access_disabled:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
root_access=False,
)
)
@@ -60,7 +61,7 @@ class Test_sagemaker_notebook_instance_root_access_disabled:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
root_access=True,
)
)

13
tests/providers/aws/services/sagemaker/sagemaker_notebook_instance_vpc_settings_configured/sagemaker_notebook_instance_vpc_settings_configured_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import NotebookInstance
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
subnet_id = "subnet-" + str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_notebook_instance_vpc_settings_configured:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
subnet_id=subnet_id,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_notebook_instance_vpc_settings_configured:
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
root_access=True,
)
)

13
tests/providers/aws/services/sagemaker/sagemaker_notebook_instance_without_direct_internet_access_configured/sagemaker_notebook_instance_without_direct_internet_access_configured_test.py
View File

@@ -2,12 +2,13 @@ from re import search
from unittest import mock
from prowler.providers.aws.services.sagemaker.sagemaker_service import NotebookInstance
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}"
class Test_sagemaker_notebook_instance_without_direct_internet_access_configured:
@@ -35,7 +36,7 @@ class Test_sagemaker_notebook_instance_without_direct_internet_access_configured
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
direct_internet_access=False,
)
)
@@ -66,7 +67,7 @@ class Test_sagemaker_notebook_instance_without_direct_internet_access_configured
NotebookInstance(
name=test_notebook_instance,
arn=notebook_instance_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
direct_internet_access=True,
)
)

81
tests/providers/aws/services/sagemaker/sagemaker_service_test.py
View File

@@ -2,25 +2,22 @@ from unittest.mock import patch
from uuid import uuid4
import botocore
from boto3 import session
from prowler.providers.aws.lib.audit_info.models import AWS_Audit_Info
from prowler.providers.aws.services.sagemaker.sagemaker_service import SageMaker
from prowler.providers.common.models import Audit_Metadata
AWS_ACCOUNT_NUMBER = "123456789012"
AWS_REGION = "eu-west-1"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
set_mocked_aws_audit_info,
)
test_notebook_instance = "test-notebook-instance"
notebook_instance_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}: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}:{AWS_ACCOUNT_NUMBER}:model/{test_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}:{AWS_ACCOUNT_NUMBER}:training-job/{test_model}"
)
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())
@@ -94,9 +91,11 @@ def mock_make_api_call(self, operation_name, kwarg):
def mock_generate_regional_clients(service, audit_info, _):
regional_client = audit_info.audit_session.client(service, region_name=AWS_REGION)
regional_client.region = AWS_REGION
return {AWS_REGION: regional_client}
regional_client = audit_info.audit_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)
@@ -105,95 +104,65 @@ def mock_generate_regional_clients(service, audit_info, _):
new=mock_generate_regional_clients,
)
class Test_SageMaker_Service:
# Mocked Audit Info
def set_mocked_audit_info(self):
audit_info = AWS_Audit_Info(
session_config=None,
original_session=None,
audit_session=session.Session(
profile_name=None,
botocore_session=None,
),
audited_account=AWS_ACCOUNT_NUMBER,
audited_account_arn=f"arn:aws:iam::{AWS_ACCOUNT_NUMBER}:root",
audited_user_id=None,
audited_partition="aws",
audited_identity_arn=None,
profile=None,
profile_region=None,
credentials=None,
assumed_role_info=None,
audited_regions=None,
organizations_metadata=None,
audit_resources=None,
mfa_enabled=False,
audit_metadata=Audit_Metadata(
services_scanned=0,
expected_checks=[],
completed_checks=0,
audit_progress=0,
),
)
return audit_info
# Test SageMaker Service
def test_service(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
assert sagemaker.service == "sagemaker"
# Test SageMaker client
def test_client(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
for reg_client in sagemaker.regional_clients.values():
assert reg_client.__class__.__name__ == "SageMaker"
# Test SageMaker session
def test__get_session__(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
assert sagemaker.session.__class__.__name__ == "Session"
# Test SageMaker list notebook instances
def test_list_notebook_instances(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
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
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):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
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
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):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
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
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):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
assert len(sagemaker.sagemaker_notebook_instances) == 1
assert sagemaker.sagemaker_notebook_instances[0].root_access
@@ -203,7 +172,7 @@ class Test_SageMaker_Service:
# Test SageMaker describe model
def test_describe_model(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
assert len(sagemaker.sagemaker_models) == 1
assert sagemaker.sagemaker_models[0].network_isolation
@@ -211,7 +180,7 @@ class Test_SageMaker_Service:
# Test SageMaker describe training jobs
def test_describe_training_jobs(self):
audit_info = self.set_mocked_audit_info()
audit_info = set_mocked_aws_audit_info([AWS_REGION_EU_WEST_1])
sagemaker = SageMaker(audit_info)
assert len(sagemaker.sagemaker_training_jobs) == 1
assert sagemaker.sagemaker_training_jobs[0].container_traffic_encryption

13
tests/providers/aws/services/sagemaker/sagemaker_training_jobs_intercontainer_encryption_enabled/sagemaker_training_jobs_intercontainer_encryption_enabled_test.py
View File

@@ -2,12 +2,13 @@ from re import search
from unittest import mock
from prowler.providers.aws.services.sagemaker.sagemaker_service import TrainingJob
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_training_job = "test-training-job"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
class Test_sagemaker_training_jobs_intercontainer_encryption_enabled:
@@ -33,7 +34,7 @@ class Test_sagemaker_training_jobs_intercontainer_encryption_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
container_traffic_encryption=True,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_training_jobs_intercontainer_encryption_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(

13
tests/providers/aws/services/sagemaker/sagemaker_training_jobs_network_isolation_enabled/sagemaker_training_jobs_network_isolation_enabled_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import TrainingJob
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_training_job = "test-training-job"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
kms_key_id = str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_training_jobs_network_isolation_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
network_isolation=True,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_training_jobs_network_isolation_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(

13
tests/providers/aws/services/sagemaker/sagemaker_training_jobs_volume_and_output_encryption_enabled/sagemaker_training_jobs_volume_and_output_encryption_enabled_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import TrainingJob
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_training_job = "test-training-job"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
kms_key_id = str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_training_jobs_volume_and_output_encryption_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
volume_kms_key_id=kms_key_id,
)
)
@@ -62,7 +63,7 @@ class Test_sagemaker_training_jobs_volume_and_output_encryption_enabled:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(

13
tests/providers/aws/services/sagemaker/sagemaker_training_jobs_vpc_settings_configured/sagemaker_training_jobs_vpc_settings_configured_test.py
View File

@@ -3,12 +3,13 @@ from unittest import mock
from uuid import uuid4
from prowler.providers.aws.services.sagemaker.sagemaker_service import TrainingJob
AWS_REGION = "eu-west-1"
AWS_ACCOUNT_NUMBER = "123456789012"
from tests.providers.aws.audit_info_utils import (
AWS_ACCOUNT_NUMBER,
AWS_REGION_EU_WEST_1,
)
test_training_job = "test-training-job"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
training_job_arn = f"arn:aws:sagemaker:{AWS_REGION_EU_WEST_1}:{AWS_ACCOUNT_NUMBER}:training-job/{test_training_job}"
subnet_id = "subnet-" + str(uuid4())
@@ -35,7 +36,7 @@ class Test_sagemaker_training_jobs_vpc_settings_configured:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
vpc_config_subnets=[subnet_id],
)
)
@@ -65,7 +66,7 @@ class Test_sagemaker_training_jobs_vpc_settings_configured:
TrainingJob(
name=test_training_job,
arn=training_job_arn,
region=AWS_REGION,
region=AWS_REGION_EU_WEST_1,
)
)
with mock.patch(