handling arn option

tests/unit_test.py

@@ -440,18 +440,23 @@ def test_bad_user_session_token(sample_json_response, sample_headers, mocker):
     [("duo", 123, 123), ("okta", 345, 345), ("google", 456, 456)],
 )
 def test_mfa_provider_type(
-    mfa_provider, session_token, expected, mocker, sample_headers
+    mfa_provider,
+    session_token,
+    expected,
+    mocker,
+    sample_headers,
 ):
     """Test whether function return key on specific MFA provider."""
     from tokendito.okta_helpers import mfa_provider_type
 
     payload = {"x": "y", "t": "z"}
     callback_url = "https://www.acme.org"
-    mfa_verify = {"sessionToken": session_token}
     selected_mfa_option = 1
     mfa_challenge_url = 1
     primary_auth = 1
     selected_factor = 1
+
+    mfa_verify = {"sessionToken": session_token}
     mocker.patch(
         "tokendito.duo_helpers.authenticate_duo",
         return_value=(payload, sample_headers, callback_url),
@@ -480,11 +485,12 @@ def test_bad_mfa_provider_type(mocker, sample_headers):
 
     payload = {"x": "y", "t": "z"}
     callback_url = "https://www.acme.org"
-    mfa_verify = {"sessionToken": "123"}
     selected_mfa_option = 1
     mfa_challenge_url = 1
     primary_auth = 1
     selected_factor = 1
+
+    mfa_verify = {"sessionToken": "123"}
     mfa_bad_provider = "bad_provider"
     mocker.patch(
         "tokendito.duo_helpers.authenticate_duo",
@@ -626,3 +632,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",
+    [
+        ("token", None, "arn:aws:iam::123:role/token"),
+        (
+            "dito",
+            None,
+            "arn:aws:iam::124:role/dito",
+        ),
+        (
+            None,
+            "arn:aws:iam::124:role/dito",
+            "arn:aws:iam::124:role/dito",
+        ),
+        (
+            None,
+            "arn:aws:iam::123:role/token",
+            "arn:aws:iam::123:role/token",
+        ),
+    ],
+)
+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::123:role/token",
+        "arn:aws:iam::124:role/dito",
+    ]
+    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::123:role/token",
+        "arn:aws:iam::123:role/token",
+    ]
+    monkeypatch.setattr("tokendito.settings.aws_profile", "token")
+
+    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::123:role/token",
+        "arn:aws:iam::124:role/dito",
+    ]
+    monkeypatch.setattr("tokendito.settings.aws_profile", "wrong_response")
+    monkeypatch.setattr(
+        "tokendito.settings.role_arn",
+        "arn:aws:iam::123:role/wrong_response",
+    )
+    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