Skip to content

Commit

Permalink
Merge pull request #32 from nathan-v/Fix_Mutliple_Account_Role_Select…
Browse files Browse the repository at this point in the history
…ion_Bug

Fix role selection bug in multiple account case
  • Loading branch information
nathan-v authored Nov 22, 2019
2 parents 1f43d1a + 722c398 commit 373d938
Show file tree
Hide file tree
Showing 5 changed files with 54 additions and 47 deletions.
27 changes: 18 additions & 9 deletions aws_okta_keyman/aws.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,8 +142,8 @@ def __init__(self,
'SessionToken': None,
'Expiration': None}
self.session_token = None
self.roles = self.assertion.roles()
self.role = role
self.available_roles()

@property
def is_valid(self):
Expand Down Expand Up @@ -176,17 +176,26 @@ def available_roles(self):
multiple_accounts = False
first_account = ''
formatted_roles = []
for role in self.roles:
for role in self.assertion.roles():
account = role['role'].split(':')[4]
role = role['role'].split(':')[5].split('/')[1]
formatted_roles.append(
{'account': account, 'role': role})
role_name = role['role'].split(':')[5].split('/')[1]
formatted_roles.append({
'account': account,
'role_name': role_name,
'arn': role['role'],
'principle': role['principle']
})
if first_account == '':
first_account = account
elif first_account != account:
multiple_accounts = True

return formatted_roles, multiple_accounts
if multiple_accounts:
formatted_roles = self.account_ids_to_names(formatted_roles)

self.roles = sorted(formatted_roles,
key=lambda k: (k['account'], k['role_name']))
return self.roles

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

LOG.info('Assuming role: {}'.format(self.roles[self.role]['role']))
LOG.info('Assuming role: {}'.format(self.roles[self.role]['arn']))

session = self.sts.assume_role_with_saml(
RoleArn=self.roles[self.role]['role'],
RoleArn=self.roles[self.role]['arn'],
PrincipalArn=self.roles[self.role]['principle'],
SAMLAssertion=self.assertion.encode())
self.creds = session['Credentials']
Expand Down Expand Up @@ -238,7 +247,7 @@ def account_ids_to_names(self, roles):
for role in roles:
role['account'] = accounts[role['account']]
LOG.debug("AWS roles with friendly names: {}".format(accounts))
return sorted(roles, key=lambda k: (k['account'], k['role']))
return roles

def get_account_name_map(self):
""" Get the friendly to ID mappings from AWS via hacktastic HTML
Expand Down
6 changes: 2 additions & 4 deletions aws_okta_keyman/keyman.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,10 +338,8 @@ def handle_multiple_roles(self, session):
with a list to pick from
"""
self.log.warning('Multiple AWS roles found; please select one')
roles, multiple_accounts = session.available_roles()
if multiple_accounts:
roles = session.account_ids_to_names(roles)
header = [{'account': 'Account'}, {'role': 'Role'}]
roles = session.available_roles()
header = [{'account': 'Account'}, {'role_name': 'Role'}]
self.role = self.selector_menu(roles, header)
session.role = self.role

Expand Down
2 changes: 1 addition & 1 deletion aws_okta_keyman/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,5 +14,5 @@
# Copyright 2018 Nathan V
"""Package metadata."""

__version__ = '0.7.3'
__version__ = '0.7.4'
__desc__ = 'AWS Okta Keyman'
48 changes: 30 additions & 18 deletions aws_okta_keyman/test/aws_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,9 +210,9 @@ def test_write(self, mock_add_profile):
def test_assume_role(self, mock_write):
mock_write.return_value = None
assertion = mock.Mock()
assertion.roles.return_value = [{'role': '', 'principle': ''}]
assertion.roles.return_value = [{'arn': '', 'principle': ''}]
session = aws.Session('BogusAssertion')
session.roles = [{'role': '', 'principle': ''}]
session.roles = [{'arn': '', 'principle': ''}]
session.assertion = assertion
sts = {'Credentials':
{'AccessKeyId': 'AKI',
Expand All @@ -239,8 +239,8 @@ def test_assume_role(self, mock_write):
def test_assume_role_multiple(self, mock_write):
mock_write.return_value = None
assertion = mock.Mock()
roles = [{'role': '1', 'principle': ''},
{'role': '2', 'principle': ''}]
roles = [{'arn': '1', 'principle': ''},
{'arn': '2', 'principle': ''}]
assertion.roles.return_value = roles
session = aws.Session('BogusAssertion')
session.assertion = assertion
Expand All @@ -260,10 +260,10 @@ def test_assume_role_multiple(self, mock_write):
def test_assume_role_preset(self, mock_write):
mock_write.return_value = None
assertion = mock.Mock()
assertion.roles.return_value = [{'role': '', 'principle': ''}]
assertion.roles.return_value = [{'arn': '', 'principle': ''}]
session = aws.Session('BogusAssertion')
session.role = 0
session.roles = [{'role': '', 'principle': ''}]
session.roles = [{'arn': '', 'principle': ''}]
session.assertion = assertion
sts = {'Credentials':
{'AccessKeyId': 'AKI',
Expand All @@ -287,29 +287,41 @@ def test_assume_role_preset(self, mock_write):
])

def test_available_roles(self):
roles = [{'role': '::::1:role/role'},
{'role': '::::1:role/role'}]
roles = [{'role': '::::1:role/role', 'principle': ''},
{'role': '::::1:role/role', 'principle': ''}]
session = aws.Session('BogusAssertion')
session.roles = roles
session.assertion = mock.MagicMock()
session.assertion.roles.return_value = roles
expected = [
{'account': '1', 'role': 'role'},
{'account': '1', 'role': 'role'}
], False
{'account': '1', 'role_name': 'role',
'principle': '', 'arn': '::::1:role/role'},
{'account': '1', 'role_name': 'role',
'principle': '', 'arn': '::::1:role/role'}
]

result = session.available_roles()

print(result)
self.assertEqual(expected, result)

def test_available_roles_multiple_accounts(self):
roles = [{'role': '::::1:role/role'},
{'role': '::::2:role/role'}]
roles = [{'role': '::::1:role/role', 'principle': ''},
{'role': '::::2:role/role', 'principle': ''}]
roles_full = [{'account': '1', 'role_name': 'role',
'arn': '::::1:role/role', 'principle': ''},
{'account': '2', 'role_name': 'role',
'arn': '::::2:role/role', 'principle': ''}]
session = aws.Session('BogusAssertion')
session.roles = roles
session.assertion = mock.MagicMock()
session.assertion.roles.return_value = roles
session.account_ids_to_names = mock.MagicMock()
session.account_ids_to_names.return_value = roles_full
expected = [
{'account': '1', 'role': 'role'},
{'account': '2', 'role': 'role'}
], True
{'account': '1', 'role_name': 'role',
'principle': '', 'arn': '::::1:role/role'},
{'account': '2', 'role_name': 'role',
'principle': '', 'arn': '::::2:role/role'}
]

result = session.available_roles()

Expand Down
18 changes: 3 additions & 15 deletions aws_okta_keyman/test/keyman_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -572,32 +572,20 @@ def test_handle_multiple_roles(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = ([{}, {}], False)
roles = ([{}, {}])
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles

keyman.handle_multiple_roles(mock_session)

keyman.selector_menu.assert_has_calls([
mock.call([{}, {}], [{'account': 'Account'}, {'role': 'Role'}])
mock.call([{}, {}],
[{'account': 'Account'}, {'role_name': 'Role'}])
])
mock_session.assert_has_calls([
mock.call.available_roles()
])

@mock.patch('aws_okta_keyman.keyman.Config')
def test_handle_multiple_roles_and_accounts(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = ([{}, {}], True)
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles

keyman.handle_multiple_roles(mock_session)

mock_session.account_ids_to_names.assert_called_with([{}, {}])

@mock.patch('aws_okta_keyman.keyman.Config')
@mock.patch('aws_okta_keyman.keyman.aws')
def test_start_session(self, aws_mock, _config_mock):
Expand Down

0 comments on commit 373d938

Please sign in to comment.