Repository: nathan-v/aws_okta_keyman

Issue:
Role Selection In Windows
I've attached the response i've received while choosing a role, the issue in this is it picks the wrong role  , in the below output i chooses role 2 which is ideally **aws-exp-exp3** but it assumed the different one please check


12:56:37 (INFO) AWS Okta Keyman 🔐 v0.8.3
Password:
12:56:45 (WARNING) Okta Verify Push being sent...
12:56:45 (INFO) Waiting for MFA success...
12:56:48 (INFO) Waiting for MFA success...
12:56:50 (INFO) Successfully authed Rishabh Ag
12:56:51 (INFO) Getting SAML Assertion from example
12:56:51 (WARNING) Application-level MFA present; re-authenticating Okta
12:56:52 (WARNING) Okta Verify Push being sent...
12:56:52 (INFO) Waiting for MFA success...
12:56:55 (INFO) Waiting for MFA success...
12:56:57 (INFO) Getting SAML Assertion from example
12:56:59 (INFO) Starting AWS session for us-east-1
12:56:59 (WARNING) Multiple AWS roles found; please select one

    Account       Role
[0] 123456789009  aws-exp-exp1
[1] 123456789009  aws-exp-exp2
[2] 123456789009  aws-exp-exp3
Selection: 2

12:57:02 (INFO) Assuming role: arn:aws:iam::123456789009:role/aws-exp-exp2
12:57:04 (INFO) Wrote profile "default" to C:\Users\Rishabh/.aws\credentials 💾
12:57:04 (INFO) Current time is 2020-07-15 07:27:04.316262
12:57:04 (INFO) Session expires at 2020-07-15 08:27:03+00:00 ⏳
12:57:04 (INFO) All done! 👍

Discussion:

Patch:
diff --git a/aws_okta_keyman/aws.py b/aws_okta_keyman/aws.py
index 80484db..91e961b 100644
--- a/aws_okta_keyman/aws.py
+++ b/aws_okta_keyman/aws.py
@@ -179,7 +179,6 @@ class Session(object):
         multiple_accounts = False
         first_account = ''
         formatted_roles = []
-        i = 0
         for role in self.assertion.roles():
             account = role['role'].split(':')[4]
             role_name = role['role'].split(':')[5].split('/')[1]
@@ -187,20 +186,27 @@ class Session(object):
                 'account': account,
                 'role_name': role_name,
                 'arn': role['role'],
-                'principle': role['principle'],
-                'roleIdx': i
+                'principle': role['principle']
             })
             if first_account == '':
                 first_account = account
             elif first_account != account:
                 multiple_accounts = True
-            i = i + 1
 
         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']))
+        formatted_roles = sorted(formatted_roles,
+                                 key=lambda k: (k['account'], k['role_name']))
+
+        # set the role role index after sorting
+        i = 0
+        for role in formatted_roles:
+            role['roleIdx'] = i
+            i = i + 1
+
+        self.roles = formatted_roles
+
         return self.roles
 
     def assume_role(self, print_only=False):


Test patch:
diff --git a/aws_okta_keyman/test/aws_test.py b/aws_okta_keyman/test/aws_test.py
index 60be405..f4f1de1 100644
--- a/aws_okta_keyman/test/aws_test.py
+++ b/aws_okta_keyman/test/aws_test.py
@@ -271,10 +271,15 @@ class TestSession(unittest.TestCase):
     def test_assume_role_preset(self, mock_write):
         mock_write.return_value = None
         assertion = mock.Mock()
-        assertion.roles.return_value = [{'arn': '', 'principle': ''}]
+
+        roles = [{'role': '::::1:role/role1', 'principle': '', 'arn': '1'},
+                 {'role': '::::1:role/role2', 'principle': '', 'arn': '2'},
+                 {'role': '::::1:role/role3', 'principle': '', 'arn': '3'}]
+
+        assertion.roles.return_value = roles
         session = aws.Session('BogusAssertion')
-        session.role = 0
-        session.roles = [{'arn': '', 'principle': ''}]
+        session.role = 1
+        session.roles = roles
         session.assertion = assertion
         sts = {'Credentials':
                {'AccessKeyId':     'AKI',
@@ -296,6 +301,13 @@ class TestSession(unittest.TestCase):
         mock_write.assert_has_calls([
             mock.call()
         ])
+        session.sts.assert_has_calls([
+            mock.call.assume_role_with_saml(
+                RoleArn='2',
+                PrincipalArn='',
+                SAMLAssertion=mock.ANY,
+                DurationSeconds=3600)
+        ])
 
     @mock.patch('aws_okta_keyman.aws.Session._print_creds')
     @mock.patch('aws_okta_keyman.aws.Session._write')
@@ -420,23 +432,29 @@ class TestSession(unittest.TestCase):
         self.assertEqual(ret, expected)
 
     def test_available_roles(self):
-        roles = [{'role': '::::1:role/role', 'principle': ''},
-                 {'role': '::::1:role/role', 'principle': ''}]
+        roles = [{'role': '::::1:role/role1', 'principle': ''},
+                 {'role': '::::1:role/role3', 'principle': ''},
+                 {'role': '::::1:role/role2', 'principle': ''}]
         session = aws.Session('BogusAssertion')
         session.assertion = mock.MagicMock()
         session.assertion.roles.return_value = roles
-        expected = [
-            {'account': '1', 'role_name': 'role',
-             'principle': '', 'arn': '::::1:role/role',
-             'roleIdx': 0},
-            {'account': '1', 'role_name': 'role',
-             'principle': '', 'arn': '::::1:role/role',
-             'roleIdx': 1}
-            ]
 
         result = session.available_roles()
 
         print(result)
+
+        expected = [
+            {'account': '1', 'role_name': 'role1',
+             'principle': '', 'arn': '::::1:role/role1',
+             'roleIdx': 0},
+            {'account': '1', 'role_name': 'role2',
+             'principle': '', 'arn': '::::1:role/role2',
+             'roleIdx': 1},
+            {'account': '1', 'role_name': 'role3',
+             'principle': '', 'arn': '::::1:role/role3',
+             'roleIdx': 2}
+            ]
+
         self.assertEqual(expected, result)
 
     def test_available_roles_multiple_accounts(self):
@@ -453,9 +471,9 @@ class TestSession(unittest.TestCase):
         session.account_ids_to_names.return_value = roles_full
         expected = [
             {'account': '1', 'role_name': 'role',
-             'principle': '', 'arn': '::::1:role/role'},
+             'principle': '', 'arn': '::::1:role/role', 'roleIdx': 0},
             {'account': '2', 'role_name': 'role',
-             'principle': '', 'arn': '::::2:role/role'}
+             'principle': '', 'arn': '::::2:role/role', 'roleIdx': 1}
             ]
 
         result = session.available_roles()