forked from aws/aws-cli
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_assume_role.py
More file actions
249 lines (218 loc) · 9.22 KB
/
test_assume_role.py
File metadata and controls
249 lines (218 loc) · 9.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
# Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
import os
import shutil
import tempfile
import json
import time
from botocore.session import Session
from botocore.exceptions import ClientError
from awscli.testutils import unittest, aws, random_chars
S3_READ_POLICY_ARN = 'arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess'
class TestAssumeRoleCredentials(unittest.TestCase):
def setUp(self):
super(TestAssumeRoleCredentials, self).setUp()
self.environ = os.environ.copy()
self.parent_session = Session()
self.iam = self.parent_session.create_client('iam')
self.sts = self.parent_session.create_client('sts')
self.tempdir = tempfile.mkdtemp()
self.config_file = os.path.join(self.tempdir, 'config')
# A role trust policy that allows the current account to call assume
# role on itself.
account_id = self.sts.get_caller_identity()['Account']
self.role_policy = {
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": {
"AWS": "arn:aws:iam::%s:root" % account_id
},
"Action": "sts:AssumeRole"
}
]
}
def tearDown(self):
super(TestAssumeRoleCredentials, self).tearDown()
shutil.rmtree(self.tempdir)
def random_name(self):
return 'clitest-' + random_chars(10)
def create_role(self, policy_document, policy_arn=None):
name = self.random_name()
response = self.iam.create_role(
RoleName=name,
AssumeRolePolicyDocument=json.dumps(policy_document)
)
self.addCleanup(self.iam.delete_role, RoleName=name)
if policy_arn:
self.iam.attach_role_policy(RoleName=name, PolicyArn=policy_arn)
self.addCleanup(
self.iam.detach_role_policy, RoleName=name,
PolicyArn=policy_arn
)
return response['Role']
def create_user(self, policy_arns):
name = self.random_name()
user = self.iam.create_user(UserName=name)['User']
self.addCleanup(self.iam.delete_user, UserName=name)
for arn in policy_arns:
self.iam.attach_user_policy(
UserName=name,
PolicyArn=arn
)
self.addCleanup(
self.iam.detach_user_policy,
UserName=name, PolicyArn=arn
)
return user
def create_creds(self, user_name):
creds = self.iam.create_access_key(UserName=user_name)['AccessKey']
self.addCleanup(
self.iam.delete_access_key,
UserName=user_name, AccessKeyId=creds['AccessKeyId']
)
return creds
def wait_for_assume_role(self, role_arn, access_key, secret_key,
token=None, attempts=30, delay=10,
num_success_needed=3):
# "Why not use the policy simulator?" you might ask. The answer is
# that the policy simulator will return success far before you can
# actually make the calls.
client = self.parent_session.create_client(
'sts', aws_access_key_id=access_key,
aws_secret_access_key=secret_key, aws_session_token=token
)
attempts_remaining = attempts
role_session_name = random_chars(10)
num_success = 0
while attempts_remaining > 0:
try:
result = client.assume_role(
RoleArn=role_arn, RoleSessionName=role_session_name)
num_success += 1
if num_success == num_success_needed:
return result['Credentials']
except ClientError as e:
code = e.response.get('Error', {}).get('Code')
if code not in ["InvalidClientTokenId", "AccessDenied"]:
raise
attempts_remaining -= 1
time.sleep(delay)
raise Exception("Unable to assume role %s" % role_arn)
def create_assume_policy(self, role_arn):
policy_document = {
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Resource": role_arn,
"Action": "sts:AssumeRole"
}
]
}
name = self.random_name()
response = self.iam.create_policy(
PolicyName=name,
PolicyDocument=json.dumps(policy_document)
)
self.addCleanup(
self.iam.delete_policy, PolicyArn=response['Policy']['Arn']
)
return response['Policy']['Arn']
def assert_s3_read_only_profile(self, profile_name):
# Calls to S3 should succeed
command = 's3api list-buckets --profile %s' % profile_name
result = aws(command, env_vars=self.environ)
self.assertEqual(result.rc, 0, result.stderr)
# Calls to other services should not
command = 'iam list-groups --profile %s' % profile_name
result = aws(command, env_vars=self.environ)
self.assertNotEqual(result.rc, 0, result.stdout)
self.assertIn('AccessDenied', result.stderr)
def test_recursive_assume_role(self):
# Create the final role, the one that will actually have access to s3
final_role = self.create_role(self.role_policy, S3_READ_POLICY_ARN)
# Create the role that can assume the final role
middle_policy_arn = self.create_assume_policy(final_role['Arn'])
middle_role = self.create_role(self.role_policy, middle_policy_arn)
# Create a user that can only assume the middle-man role, and then get
# static credentials for it.
user_policy_arn = self.create_assume_policy(middle_role['Arn'])
user = self.create_user([user_policy_arn])
user_creds = self.create_creds(user['UserName'])
# Setup the config file with the profiles we'll be using. For
# convenience static credentials are placed here instead of putting
# them in the credentials file.
config = (
'[default]\n'
'aws_access_key_id = %s\n'
'aws_secret_access_key = %s\n'
'[profile middle]\n'
'source_profile = default\n'
'role_arn = %s\n'
'[profile final]\n'
'source_profile = middle\n'
'role_arn = %s\n'
)
config = config % (
user_creds['AccessKeyId'], user_creds['SecretAccessKey'],
middle_role['Arn'], final_role['Arn']
)
with open(self.config_file, 'w') as f:
f.write(config)
# Wait for IAM permissions to propagate
middle_creds = self.wait_for_assume_role(
role_arn=middle_role['Arn'],
access_key=user_creds['AccessKeyId'],
secret_key=user_creds['SecretAccessKey'],
)
self.wait_for_assume_role(
role_arn=final_role['Arn'],
access_key=middle_creds['AccessKeyId'],
secret_key=middle_creds['SecretAccessKey'],
token=middle_creds['SessionToken'],
)
# Configure our credentials file to be THE credentials file
self.environ['AWS_CONFIG_FILE'] = self.config_file
self.assert_s3_read_only_profile(profile_name='final')
def test_assume_role_with_credential_source(self):
# Create a role with read access to S3
role = self.create_role(self.role_policy, S3_READ_POLICY_ARN)
# Create a user that can assume the role and get static credentials
# for it.
user_policy_arn = self.create_assume_policy(role['Arn'])
user = self.create_user([user_policy_arn])
user_creds = self.create_creds(user['UserName'])
# Setup the config file with the profile we'll be using.
config = (
'[profile assume]\n'
'role_arn = %s\n'
'credential_source = Environment\n'
)
config = config % role['Arn']
with open(self.config_file, 'w') as f:
f.write(config)
# Wait for IAM permissions to propagate
self.wait_for_assume_role(
role_arn=role['Arn'],
access_key=user_creds['AccessKeyId'],
secret_key=user_creds['SecretAccessKey'],
)
# Setup the environment so that our new config file is THE config
# file and add the expected credentials since we're using the
# environment as our credential source.
self.environ['AWS_CONFIG_FILE'] = self.config_file
self.environ['AWS_SECRET_ACCESS_KEY'] = user_creds['SecretAccessKey']
self.environ['AWS_ACCESS_KEY_ID'] = user_creds['AccessKeyId']
self.assert_s3_read_only_profile(profile_name='assume')