From ab7486c218a8e6f869fa889e04e9d2d4ff28dc72 Mon Sep 17 00:00:00 2001
From: Chris Rutter <crutter@circle.com>
Date: Fri, 20 Jan 2017 10:18:25 -0500
Subject: [PATCH] add a --role-arn argument

This will give the user the ability to pass in the ARN of the role
he/she is attempting to assume the credentials for.  This is useful if
the user has access to multiple roles he/she can assume, and saves them
from being prompted to select one from a list.
---
 README.rst                      | 9 +++++++++
 aws_role_credentials/actions.py | 7 ++++++-
 aws_role_credentials/cli.py     | 4 ++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/README.rst b/README.rst
index 64f1f00..9be3336 100644
--- a/README.rst
+++ b/README.rst
@@ -20,6 +20,12 @@ Simply pipe a SAML assertion into awssaml
     # create credentials from saml assertion
     $ oktaauth -u jobloggs | aws_role_credentials saml --profile dev
 
+Or for assuming a known role name:
+
+.. code-block:: shell
+
+    # create credentials from saml assertion using a known role ARN
+    $ oktaauth -u jobloggs | aws_role_credentials saml --profile dev --role-arn arn:aws:iam::098765432109:role/ReadOnly
 
 Or for assuming a role using an IAM user:
 
@@ -64,8 +70,11 @@ Options
 
     --profile          Use a specific profile in your credential file (e.g. Development).  Defaults to sts.
     --region           The region to use. Overrides config/env settings.  Defaults to us-east-1.
+    --role-arn         Optional `role ARN`_ to use when multiple roles are available.
     --exec             The command to execute with the AWS credentials
 
+.. _role ARN: http://docs.aws.amazon.com/IAM/latest/UserGuide/reference_identifiers.html
+
 Thanks
 ======
 
diff --git a/aws_role_credentials/actions.py b/aws_role_credentials/actions.py
index 92d7ae3..c3e9757 100644
--- a/aws_role_credentials/actions.py
+++ b/aws_role_credentials/actions.py
@@ -75,7 +75,12 @@ def exec_handler(region, exec_command, **kwargs):
     def saml_token(region, assertion, **kwargs):
         assertion = SamlAssertion(assertion)
         roles = assertion.roles()
-        if len(roles) > 1:
+        if kwargs.get('role_arn', False):
+            for i, role in enumerate(roles):
+                if role['role'] == kwargs['role_arn']:
+                    role = roles[i]
+                    break
+        elif len(roles) > 1:
             print('Please select the role you would like to assume:')
             for i, role in enumerate(roles):
                 print('[{}] - {}'.format(i, role['role']))
diff --git a/aws_role_credentials/cli.py b/aws_role_credentials/cli.py
index 64cf6a8..c4f35ff 100644
--- a/aws_role_credentials/cli.py
+++ b/aws_role_credentials/cli.py
@@ -73,6 +73,10 @@ def create_parser(prog, epilog,
         default='us-east-1',
         help='The region to use. Overrides config/env settings.')
 
+    parent_parser.add_argument(
+        '--role-arn', type=str,
+        help='Optional role ARN to use when multiple roles are available.')
+
     parent_parser.add_argument(
         '--exec', type=str,
         dest='exec_command',