Skip to content

fix: remove unrestrictive principal * from KMS policy tests. #712

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 30 commits into from
Mar 21, 2019
Merged

fix: remove unrestrictive principal * from KMS policy tests. #712

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
34bdbf7
pass kms id as parameter for uploading code with Server side encryption
mvsusp Mar 11, 2019
b40476e
Merge branch 'master' into mvs-kms-sse
mvsusp Mar 11, 2019
2099ab9
Merge branch 'master' into mvs-kms-sse
mvsusp Mar 12, 2019
40b559c
Fix tests
mvsusp Mar 12, 2019
6e817a3
Fix tests
mvsusp Mar 12, 2019
de46eaa
Merge branch 'master' into mvs-kms-sse
mvsusp Mar 13, 2019
d5cd080
Fix merge changes
mvsusp Mar 13, 2019
364a7ce
Fix tests
mvsusp Mar 13, 2019
5c9566a
Support for specific principal
mvsusp Mar 19, 2019
a9e60f0
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 19, 2019
503224c
WIP
mvsusp Mar 19, 2019
87d1ce4
Add principal to the policy
mvsusp Mar 19, 2019
7f04058
Add principal to the policy
mvsusp Mar 20, 2019
d42a34b
Add principal to the policy
mvsusp Mar 20, 2019
fa64437
Merge branch 'master' into mvs-fix-principal
mvsusp Mar 20, 2019
98bafd7
Update test_tf_script_mode.py
mvsusp Mar 20, 2019
5436096
Add principal to the policy
mvsusp Mar 20, 2019
d083637
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 20, 2019
b5d0cf3
Merge branch 'mvs-fix-principal' of github.com:mvsusp/sagemaker-pytho…
mvsusp Mar 20, 2019
7d9c229
Add principal to the policy
mvsusp Mar 20, 2019
43efe7f
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 20, 2019
e392b76
Merge branch 'master' into mvs-fix-principal
mvsusp Mar 20, 2019
3ee6952
Merge branch 'master' into mvs-fix-principal
mvsusp Mar 20, 2019
5f51727
Add principal to the policy
mvsusp Mar 20, 2019
9af4f88
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 20, 2019
53039f6
Merge branch 'mvs-fix-principal' of github.com:mvsusp/sagemaker-pytho…
mvsusp Mar 20, 2019
08f5646
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 21, 2019
50ef480
Handle PR comments
mvsusp Mar 21, 2019
79f860e
Add principal to the policy
mvsusp Mar 21, 2019
2950559
Merge remote-tracking branch 'origin/master' into mvs-fix-principal
mvsusp Mar 21, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 82 additions & 15 deletions tests/integ/kms_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,27 @@
# language governing permissions and limitations under the License.
from __future__ import absolute_import

import contextlib
import json

from botocore import exceptions

KEY_ALIAS = "SageMakerIntegTestKmsKey"
PRINCIPAL_TEMPLATE = '["{account_id}", "{role_arn}", ' \
'"arn:aws:iam::{account_id}:role/{sagemaker_role}"] '

KEY_ALIAS = 'SageMakerTestKMSKey'
KMS_S3_ALIAS = 'SageMakerTestS3KMSKey'
POLICY_NAME = 'default'
KEY_POLICY = '''
{{
"Version": "2012-10-17",
"Id": "sagemaker-kms-integ-test-policy",
"Id": "{id}",
"Statement": [
{{
"Sid": "Enable IAM User Permissions",
"Effect": "Allow",
"Principal": {{
"AWS": "*"
"AWS": {principal}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why not grant to account root instead of a role? that way any entity in the account should get these permissions.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Granting to the account was my first attempt, unfortunately KMS permissions don't propagate to the assumed role when we do that.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Btw, the account root is already one of the principals in the list.

}},
"Action": "kms:*",
"Resource": "*"
Expand All @@ -42,22 +50,75 @@ def _get_kms_key_arn(kms_client, alias):
return None


def _create_kms_key(kms_client, account_id):
def _get_kms_key_id(kms_client, alias):
try:
response = kms_client.describe_key(KeyId='alias/' + alias)
return response['KeyMetadata']['KeyId']
except kms_client.exceptions.NotFoundException:
return None


def _create_kms_key(kms_client,
account_id,
role_arn=None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why accept a role name and a role arn?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are three principals that need to be created: the account root, the role that is running the tests, and the role that will be used by SageMaker. Example:

            "Principal": {
                "AWS": [
                    "arn:aws:iam::142577830533:role/SageMakerRole",
                    "arn:aws:sts::142577830533:assumed-role/DevBuildStackPDX-PullRequestBuildRole4B40B95C-G4AE4UBWNU2K/AWSCodeBuild-fe90a41f-9006-4ccc-918c-437c0f5b26d2",
                    "arn:aws:iam::142577830533:root"
                ]
            },

sagemaker_role='SageMakerRole',
alias=KEY_ALIAS):
if role_arn:
principal = PRINCIPAL_TEMPLATE.format(account_id=account_id,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is the idea here a policy that allows the root account, some role arn, plus another role constructed from the sagemaker_role name? why not simplify that and just require test environments be set up correctly?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Talked with @jesterhazy offline. We didn't find an easier way to simplify this test without breaking running tests outside the CI.

role_arn=role_arn,
sagemaker_role=sagemaker_role)
else:
principal = "{account_id}".format(account_id=account_id)

response = kms_client.create_key(
Policy=KEY_POLICY.format(account_id=account_id),
Policy=KEY_POLICY.format(id=POLICY_NAME, principal=principal, sagemaker_role=sagemaker_role),
Description='KMS key for SageMaker Python SDK integ tests',
)
key_arn = response['KeyMetadata']['Arn']
response = kms_client.create_alias(AliasName='alias/' + KEY_ALIAS, TargetKeyId=key_arn)

if alias:
kms_client.create_alias(AliasName='alias/' + alias, TargetKeyId=key_arn)
return key_arn


def get_or_create_kms_key(kms_client, account_id):
kms_key_arn = _get_kms_key_arn(kms_client, KEY_ALIAS)
if kms_key_arn is not None:
return kms_key_arn
else:
return _create_kms_key(kms_client, account_id)
def _add_role_to_policy(kms_client,
account_id,
role_arn,
alias=KEY_ALIAS,
sagemaker_role='SageMakerRole'):
key_id = _get_kms_key_id(kms_client, alias)
policy = kms_client.get_key_policy(KeyId=key_id, PolicyName=POLICY_NAME)
policy = json.loads(policy['Policy'])
principal = policy['Statement'][0]['Principal']['AWS']

if role_arn not in principal or sagemaker_role not in principal:
principal = PRINCIPAL_TEMPLATE.format(account_id=account_id,
role_arn=role_arn,
sagemaker_role=sagemaker_role)

kms_client.put_key_policy(KeyId=key_id,
PolicyName=POLICY_NAME,
Policy=KEY_POLICY.format(id=POLICY_NAME, principal=principal))


def get_or_create_kms_key(kms_client,
account_id,
role_arn=None,
alias=KEY_ALIAS,
sagemaker_role='SageMakerRole'):
kms_key_arn = _get_kms_key_arn(kms_client, alias)

if kms_key_arn is None:
return _create_kms_key(kms_client, account_id, role_arn, sagemaker_role, alias)

if role_arn:
_add_role_to_policy(kms_client,
account_id,
role_arn,
alias,
sagemaker_role)

return kms_key_arn


KMS_BUCKET_POLICY = """{
Expand Down Expand Up @@ -92,9 +153,13 @@ def get_or_create_kms_key(kms_client, account_id):
}"""


def get_or_create_bucket_with_encryption(boto_session):
@contextlib.contextmanager
def bucket_with_encryption(boto_session, sagemaker_role):
account = boto_session.client('sts').get_caller_identity()['Account']
kms_key_arn = get_or_create_kms_key(boto_session.client('kms'), account)
role_arn = boto_session.client('sts').get_caller_identity()['Arn']

kms_client = boto_session.client('kms')
kms_key_arn = _create_kms_key(kms_client, account, role_arn, sagemaker_role, None)

region = boto_session.region_name
bucket_name = 'sagemaker-{}-{}-with-kms'.format(region, account)
Expand Down Expand Up @@ -132,4 +197,6 @@ def get_or_create_bucket_with_encryption(boto_session):
Policy=KMS_BUCKET_POLICY % (bucket_name, bucket_name)
)

return 's3://' + bucket_name, kms_key_arn
yield 's3://' + bucket_name, kms_key_arn

kms_client.schedule_key_deletion(KeyId=kms_key_arn, PendingWindowInDays=7)
46 changes: 24 additions & 22 deletions tests/integ/test_tf_script_mode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from tests.integ import kms_utils
import tests.integ.timeout as timeout

ROLE = 'SageMakerRole'

RESOURCE_PATH = os.path.join(os.path.dirname(__file__), '..', 'data', 'tensorflow_mnist')
SCRIPT = os.path.join(RESOURCE_PATH, 'mnist.py')
PARAMETER_SERVER_DISTRIBUTION = {'parameter_server': {'enabled': True}}
Expand Down Expand Up @@ -56,39 +58,39 @@ def test_mnist(sagemaker_session, instance_type):
['graph.pbtxt', 'model.ckpt-0.index', 'model.ckpt-0.meta'])


@pytest.mark.skip('this test is broken')
def test_server_side_encryption(sagemaker_session):

bucket_with_kms, kms_key = kms_utils.get_or_create_bucket_with_encryption(sagemaker_session.boto_session)
boto_session = sagemaker_session.boto_session
with kms_utils.bucket_with_encryption(boto_session, ROLE) as (bucket_with_kms, kms_key):

output_path = os.path.join(bucket_with_kms, 'test-server-side-encryption', time.strftime('%y%m%d-%H%M'))
output_path = os.path.join(bucket_with_kms, 'test-server-side-encryption', time.strftime('%y%m%d-%H%M'))

estimator = TensorFlow(entry_point=SCRIPT,
role='SageMakerRole',
train_instance_count=1,
train_instance_type='ml.c5.xlarge',
sagemaker_session=sagemaker_session,
py_version='py3',
framework_version='1.11',
base_job_name='test-server-side-encryption',
code_location=output_path,
output_path=output_path,
model_dir='/opt/ml/model',
output_kms_key=kms_key)
estimator = TensorFlow(entry_point=SCRIPT,
role=ROLE,
train_instance_count=1,
train_instance_type='ml.c5.xlarge',
sagemaker_session=sagemaker_session,
py_version='py3',
framework_version='1.11',
base_job_name='test-server-side-encryption',
code_location=output_path,
output_path=output_path,
model_dir='/opt/ml/model',
output_kms_key=kms_key)

inputs = estimator.sagemaker_session.upload_data(
path=os.path.join(RESOURCE_PATH, 'data'),
key_prefix='scriptmode/mnist')
inputs = estimator.sagemaker_session.upload_data(
path=os.path.join(RESOURCE_PATH, 'data'),
key_prefix='scriptmode/mnist')

with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
estimator.fit(inputs)
with timeout.timeout(minutes=integ.TRAINING_DEFAULT_TIMEOUT_MINUTES):
estimator.fit(inputs)


@pytest.mark.canary_quick
@pytest.mark.skipif(integ.PYTHON_VERSION != 'py3', reason="Script Mode tests are only configured to run with Python 3")
def test_mnist_distributed(sagemaker_session, instance_type):
estimator = TensorFlow(entry_point=SCRIPT,
role='SageMakerRole',
role=ROLE,
train_instance_count=2,
# TODO: change train_instance_type to instance_type once the test is passing consistently
train_instance_type='ml.c5.xlarge',
Expand All @@ -110,7 +112,7 @@ def test_mnist_distributed(sagemaker_session, instance_type):

def test_mnist_async(sagemaker_session):
estimator = TensorFlow(entry_point=SCRIPT,
role='SageMakerRole',
role=ROLE,
train_instance_count=1,
train_instance_type='ml.c5.4xlarge',
sagemaker_session=sagemaker_session,
Expand Down