Skip to content

Commit 373d938

Browse files
nathan-vnathan-v
authored
Merge pull request #32 from nathan-v/Fix_Mutliple_Account_Role_Selection_Bug
Fix role selection bug in multiple account case
2 parents 1f43d1a + 722c398 commit 373d938

File tree

5 files changed

+54
-47
lines changed

5 files changed

+54
-47
lines changed

‎aws_okta_keyman/aws.py

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -142,8 +142,8 @@ def __init__(self,
142142
'SessionToken': None,
143143
'Expiration': None}
144144
self.session_token = None
145-
self.roles = self.assertion.roles()
146145
self.role = role
146+
self.available_roles()
147147

148148
@property
149149
def is_valid(self):
@@ -176,17 +176,26 @@ def available_roles(self):
176176
multiple_accounts = False
177177
first_account = ''
178178
formatted_roles = []
179-
for role in self.roles:
179+
for role in self.assertion.roles():
180180
account = role['role'].split(':')[4]
181-
role = role['role'].split(':')[5].split('/')[1]
182-
formatted_roles.append(
183-
{'account': account, 'role': role})
181+
role_name = role['role'].split(':')[5].split('/')[1]
182+
formatted_roles.append({
183+
'account': account,
184+
'role_name': role_name,
185+
'arn': role['role'],
186+
'principle': role['principle']
187+
})
184188
if first_account == '':
185189
first_account = account
186190
elif first_account != account:
187191
multiple_accounts = True
188192

189-
return formatted_roles, multiple_accounts
193+
if multiple_accounts:
194+
formatted_roles = self.account_ids_to_names(formatted_roles)
195+
196+
self.roles = sorted(formatted_roles,
197+
key=lambda k: (k['account'], k['role_name']))
198+
return self.roles
190199

191200
def assume_role(self):
192201
"""Use the SAML Assertion to actually get the credentials.
@@ -200,10 +209,10 @@ def assume_role(self):
200209
raise MultipleRoles
201210
self.role = 0
202211

203-
LOG.info('Assuming role: {}'.format(self.roles[self.role]['role']))
212+
LOG.info('Assuming role: {}'.format(self.roles[self.role]['arn']))
204213

205214
session = self.sts.assume_role_with_saml(
206-
RoleArn=self.roles[self.role]['role'],
215+
RoleArn=self.roles[self.role]['arn'],
207216
PrincipalArn=self.roles[self.role]['principle'],
208217
SAMLAssertion=self.assertion.encode())
209218
self.creds = session['Credentials']
@@ -238,7 +247,7 @@ def account_ids_to_names(self, roles):
238247
for role in roles:
239248
role['account'] = accounts[role['account']]
240249
LOG.debug("AWS roles with friendly names: {}".format(accounts))
241-
return sorted(roles, key=lambda k: (k['account'], k['role']))
250+
return roles
242251

243252
def get_account_name_map(self):
244253
""" Get the friendly to ID mappings from AWS via hacktastic HTML

‎aws_okta_keyman/keyman.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -338,10 +338,8 @@ def handle_multiple_roles(self, session):
338338
with a list to pick from
339339
"""
340340
self.log.warning('Multiple AWS roles found; please select one')
341-
roles, multiple_accounts = session.available_roles()
342-
if multiple_accounts:
343-
roles = session.account_ids_to_names(roles)
344-
header = [{'account': 'Account'}, {'role': 'Role'}]
341+
roles = session.available_roles()
342+
header = [{'account': 'Account'}, {'role_name': 'Role'}]
345343
self.role = self.selector_menu(roles, header)
346344
session.role = self.role
347345

‎aws_okta_keyman/metadata.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,5 +14,5 @@
1414
# Copyright 2018 Nathan V
1515
"""Package metadata."""
1616

17-
__version__ = '0.7.3'
17+
__version__ = '0.7.4'
1818
__desc__ = 'AWS Okta Keyman'

‎aws_okta_keyman/test/aws_test.py

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -210,9 +210,9 @@ def test_write(self, mock_add_profile):
210210
def test_assume_role(self, mock_write):
211211
mock_write.return_value = None
212212
assertion = mock.Mock()
213-
assertion.roles.return_value = [{'role': '', 'principle': ''}]
213+
assertion.roles.return_value = [{'arn': '', 'principle': ''}]
214214
session = aws.Session('BogusAssertion')
215-
session.roles = [{'role': '', 'principle': ''}]
215+
session.roles = [{'arn': '', 'principle': ''}]
216216
session.assertion = assertion
217217
sts = {'Credentials':
218218
{'AccessKeyId': 'AKI',
@@ -239,8 +239,8 @@ def test_assume_role(self, mock_write):
239239
def test_assume_role_multiple(self, mock_write):
240240
mock_write.return_value = None
241241
assertion = mock.Mock()
242-
roles = [{'role': '1', 'principle': ''},
243-
{'role': '2', 'principle': ''}]
242+
roles = [{'arn': '1', 'principle': ''},
243+
{'arn': '2', 'principle': ''}]
244244
assertion.roles.return_value = roles
245245
session = aws.Session('BogusAssertion')
246246
session.assertion = assertion
@@ -260,10 +260,10 @@ def test_assume_role_multiple(self, mock_write):
260260
def test_assume_role_preset(self, mock_write):
261261
mock_write.return_value = None
262262
assertion = mock.Mock()
263-
assertion.roles.return_value = [{'role': '', 'principle': ''}]
263+
assertion.roles.return_value = [{'arn': '', 'principle': ''}]
264264
session = aws.Session('BogusAssertion')
265265
session.role = 0
266-
session.roles = [{'role': '', 'principle': ''}]
266+
session.roles = [{'arn': '', 'principle': ''}]
267267
session.assertion = assertion
268268
sts = {'Credentials':
269269
{'AccessKeyId': 'AKI',
@@ -287,29 +287,41 @@ def test_assume_role_preset(self, mock_write):
287287
])
288288

289289
def test_available_roles(self):
290-
roles = [{'role': '::::1:role/role'},
291-
{'role': '::::1:role/role'}]
290+
roles = [{'role': '::::1:role/role', 'principle': ''},
291+
{'role': '::::1:role/role', 'principle': ''}]
292292
session = aws.Session('BogusAssertion')
293-
session.roles = roles
293+
session.assertion = mock.MagicMock()
294+
session.assertion.roles.return_value = roles
294295
expected = [
295-
{'account': '1', 'role': 'role'},
296-
{'account': '1', 'role': 'role'}
297-
], False
296+
{'account': '1', 'role_name': 'role',
297+
'principle': '', 'arn': '::::1:role/role'},
298+
{'account': '1', 'role_name': 'role',
299+
'principle': '', 'arn': '::::1:role/role'}
300+
]
298301

299302
result = session.available_roles()
300303

301304
print(result)
302305
self.assertEqual(expected, result)
303306

304307
def test_available_roles_multiple_accounts(self):
305-
roles = [{'role': '::::1:role/role'},
306-
{'role': '::::2:role/role'}]
308+
roles = [{'role': '::::1:role/role', 'principle': ''},
309+
{'role': '::::2:role/role', 'principle': ''}]
310+
roles_full = [{'account': '1', 'role_name': 'role',
311+
'arn': '::::1:role/role', 'principle': ''},
312+
{'account': '2', 'role_name': 'role',
313+
'arn': '::::2:role/role', 'principle': ''}]
307314
session = aws.Session('BogusAssertion')
308-
session.roles = roles
315+
session.assertion = mock.MagicMock()
316+
session.assertion.roles.return_value = roles
317+
session.account_ids_to_names = mock.MagicMock()
318+
session.account_ids_to_names.return_value = roles_full
309319
expected = [
310-
{'account': '1', 'role': 'role'},
311-
{'account': '2', 'role': 'role'}
312-
], True
320+
{'account': '1', 'role_name': 'role',
321+
'principle': '', 'arn': '::::1:role/role'},
322+
{'account': '2', 'role_name': 'role',
323+
'principle': '', 'arn': '::::2:role/role'}
324+
]
313325

314326
result = session.available_roles()
315327

‎aws_okta_keyman/test/keyman_test.py

Lines changed: 3 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -572,32 +572,20 @@ def test_handle_multiple_roles(self, _config_mock):
572572
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
573573
keyman.selector_menu = mock.MagicMock()
574574
keyman.selector_menu.return_value = 0
575-
roles = ([{}, {}], False)
575+
roles = ([{}, {}])
576576
mock_session = mock.MagicMock()
577577
mock_session.available_roles.return_value = roles
578578

579579
keyman.handle_multiple_roles(mock_session)
580580

581581
keyman.selector_menu.assert_has_calls([
582-
mock.call([{}, {}], [{'account': 'Account'}, {'role': 'Role'}])
582+
mock.call([{}, {}],
583+
[{'account': 'Account'}, {'role_name': 'Role'}])
583584
])
584585
mock_session.assert_has_calls([
585586
mock.call.available_roles()
586587
])
587588

588-
@mock.patch('aws_okta_keyman.keyman.Config')
589-
def test_handle_multiple_roles_and_accounts(self, _config_mock):
590-
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
591-
keyman.selector_menu = mock.MagicMock()
592-
keyman.selector_menu.return_value = 0
593-
roles = ([{}, {}], True)
594-
mock_session = mock.MagicMock()
595-
mock_session.available_roles.return_value = roles
596-
597-
keyman.handle_multiple_roles(mock_session)
598-
599-
mock_session.account_ids_to_names.assert_called_with([{}, {}])
600-
601589
@mock.patch('aws_okta_keyman.keyman.Config')
602590
@mock.patch('aws_okta_keyman.keyman.aws')
603591
def test_start_session(self, aws_mock, _config_mock):

0 commit comments

Comments
 (0)