handling aws profile option

tests/unit_test.py

@@ -626,3 +626,79 @@ def test_bad_with_no_mfa_methods_user_mfa_challenge(
 
     with pytest.raises(SystemExit) as error:
         assert user_mfa_challenge(sample_headers, primary_auth) == error
+
+
+@pytest.mark.parametrize(
+    "aws_profile, role_arn, selected_role",
+    [
+        ("readonly", None, "arn:aws:iam::067463091988:role/ncu/redwood/readonly"),
+        (
+            "infnonprod.okta.ReadOnly",
+            None,
+            "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+        ),
+        (
+            None,
+            "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+            "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+        ),
+        (
+            None,
+            "arn:aws:iam::067463091988:role/ncu/redwood/readonly",
+            "arn:aws:iam::067463091988:role/ncu/redwood/readonly",
+        ),
+    ],
+)
+def test_good_select_role_arn(
+    mocker, monkeypatch, aws_profile, role_arn, selected_role
+):
+    """Test which role does the user has chosen."""
+    from tokendito.helpers import select_role_arn
+
+    saml_xml = "x"
+    saml_response_string = "y"
+
+    role_arns = [
+        "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+        "arn:aws:iam::067463091988:role/ncu/redwood/readonly",
+    ]
+    monkeypatch.setattr("tokendito.settings.aws_profile", aws_profile)
+    mocker.patch("tokendito.helpers.prompt_role_choices", return_value=selected_role)
+    assert select_role_arn(role_arns, saml_xml, saml_response_string) == selected_role
+
+
+def test_repeated_line_select_role_arn(monkeypatch):
+    """Test behaviour repeated role."""
+    from tokendito.helpers import select_role_arn
+
+    saml_xml = "x"
+    saml_response_string = "y"
+
+    role_arns = [
+        "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+        "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+    ]
+    monkeypatch.setattr("tokendito.settings.aws_profile", "infnonprod.okta.ReadOnly")
+
+    with pytest.raises(SystemExit) as error:
+        assert select_role_arn(role_arns, saml_xml, saml_response_string) == error
+
+
+def test_bad_select_role_arn(monkeypatch):
+    """Test behaviour wrong aws_profile and role_arn."""
+    from tokendito.helpers import select_role_arn
+
+    saml_xml = "x"
+    saml_response_string = "y"
+
+    role_arns = [
+        "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly",
+        "arn:aws:iam::067463091988:role/ncu/redwood/readonly",
+    ]
+    monkeypatch.setattr("tokendito.settings.aws_profile", "infnonprod.okta.ReadOnly1")
+    monkeypatch.setattr(
+        "tokendito.settings.role_arn",
+        "arn:aws:iam::067463091988:role/infnonprod.okta.ReadOnly1",
+    )
+    with pytest.raises(SystemExit) as error:
+        assert select_role_arn(role_arns, saml_xml, saml_response_string) == error

tokendito/__version__.py

@@ -1,7 +1,7 @@
 # vim: set filetype=python ts=4 sw=4
 # -*- coding: utf-8 -*-
 """tokendito version."""
-__version__ = "1.2.1"
+__version__ = "1.3.0"
 __title__ = "tokendito"
 __description__ = "Get AWS STS tokens from Okta SSO"
 __long_description_content_type__ = "text/x-rst"

tokendito/helpers.py

@@ -268,7 +268,21 @@ def select_role_arn(role_arns, saml_xml, saml_response_string):
 
     """
     logging.debug("Select the role user wants to pick [{}]".format(role_arns))
-    if settings.role_arn is None:
+
+    role_names = dict((role.split("/")[-1], role) for role in role_arns)
+    roles = [role.split("/")[-1] for role in role_arns]
+
+    if roles.count(settings.aws_profile) > 1:
+        logging.error(
+            "You have multiple arn matches, you can select arn with --arn option"
+        )
+        sys.exit(2)
+    elif settings.aws_profile in role_names.keys():
+        selected_role = role_names[settings.aws_profile]
+        logging.debug(
+            "Using aws_profile env var for role: [{}]".format(settings.aws_profile)
+        )
+    elif settings.role_arn is None:
         selected_role = prompt_role_choices(role_arns, saml_xml, saml_response_string)
     elif settings.role_arn in role_arns:
         selected_role = settings.role_arn