aws-saml-adfs.py

#!/usr/bin/python

import sys
import boto.sts
import boto.s3
import boto.iam
import requests
import getpass
import ConfigParser
import base64
import json
import logging
import xml.etree.ElementTree as ET
import re
from bs4 import BeautifulSoup
from os.path import expanduser
from urlparse import urlparse, urlunparse
import argparse

##########################################################################
# Parse arguments from command line

parser = argparse.ArgumentParser(description='Ask for user specific information')
parser.add_argument('-i', '--IdP',
                    action="store", dest="IdP",
                    help="Identity Provider, default is ADFS", default="ADFS")

parser.add_argument('-d', '--domain',
                    action="store", dest="domain",
                    help="Domain name of IdP, i.e. https://adfs.yourdomain.com \nYou may specify port https://adfs.yourdomain.com:4443", required=True)

parser.add_argument('-r', '--region',
                    action="store", dest="region",
                    help="default AWS region, default is us-east-1", default="us-east-1")

parser.add_argument('-o', '--outputformat',
                    action="store", dest="outputformat",
                    help="default AWS output format, default is json", default="json")


parser.add_argument('-c', '--configfile',
                    action="store", dest="awsconfigfile",
                    help="default AWS configfile, default is /.aws/credentials", default="/.aws/credentials")

parser.add_argument('-s', '--sslverification',
                    action="store_true", dest="sslverification",
                    help="IdP SSL certificate verification, default is true", default=False)

parser.add_argument('-R', '--Roles',
                    action="store", dest="Roles",
                    help="Index number(s) of AWS Role(s) to assume, enter '0' to 'N', \nor you may enter multiple ones seperated by space, i.e. -R '0 3 6 9'. \nYou can also enter 'all' to assume all roles")

args = parser.parse_args()

##########################################################################
# Variables

# region: The default AWS region that this script will connect
# to for all API calls
region = args.region

# output format: The AWS CLI output format that will be configured in the
# saml profile (affects subsequent CLI calls)
outputformat = args.outputformat

# awsconfigfile: The file where this script will store the temp
# credentials under the saml profile
awsconfigfile = args.awsconfigfile

# SSL certificate verification: Whether or not strict certificate
# verification is done, False should only be used for dev/test
sslverification = args.sslverification

# idpentryurl: The initial url that starts the authentication process.
# idpentryurl = 'https://<fqdn>:<port>/idp/profile/SAML2/Unsolicited/SSO?providerId=urn:amazon:webservices'
# ADFS URL looks like this
# idpentryurl = 'https://adfs.YourDomainName.com/adfs/ls/IdpInitiatedSignOn.aspx?loginToRp=urn:amazon:webservices'

if args.IdP.lower() == "adfs":
    idpentryurl = args.domain + '/adfs/ls/IdpInitiatedSignOn.aspx?loginToRp=urn:amazon:webservices'
else:
    idpentryurl = args.domain + '/idp/profile/SAML2/Unsolicited/SSO?providerId=urn:amazon:webservices'

# Uncomment to enable low level debugging
#logging.basicConfig(level=logging.DEBUG)

##########################################################################

# Get the federated credentials from the user
print "Username:",
username = raw_input()
password = getpass.getpass()
print ''

# Initiate session handler
session = requests.Session()

# Programmatically get the SAML assertion
# Opens the initial IdP url and follows all of the HTTP302 redirects, and
# gets the resulting login page
formresponse = session.get(idpentryurl, verify=sslverification, timeout=10)
# Capture the idpauthformsubmiturl, which is the final url after all the 302s
idpauthformsubmiturl = formresponse.url

# Parse the response and extract all the necessary values
# in order to build a dictionary of all of the form values the IdP expects
formsoup = BeautifulSoup(formresponse.text.decode('utf8'), "html.parser")
payload = {}

for inputtag in formsoup.find_all(re.compile('(INPUT|input)')):
    name = inputtag.get('name','')
    value = inputtag.get('value','')
    if "user" in name.lower():
        #Make an educated guess that this is the right field for the username
        payload[name] = username
    elif "email" in name.lower():
        #Some IdPs also label the username field as 'email'
        payload[name] = username
    elif "pass" in name.lower():
        #Make an educated guess that this is the right field for the password
        payload[name] = password
    else:
        #Simply populate the parameter with the existing value (picks up hidden fields in the login form)
        payload[name] = value

# Debug the parameter payload if needed
# Use with caution since this will print sensitive output to the screen
#print payload

# Some IdPs don't explicitly set a form action, but if one is set we should
# build the idpauthformsubmiturl by combining the scheme and hostname 
# from the entry url with the form action target
# If the action tag doesn't exist, we just stick with the 
# idpauthformsubmiturl above
for inputtag in formsoup.find_all(re.compile('(FORM|form)')):
    action = inputtag.get('action')
    actionURL = urlparse(action)
    if (actionURL.scheme == '' and actionURL.netloc == '') and actionURL.path:
        parsedurl = urlparse(idpentryurl)
        idpauthformsubmiturl = parsedurl.scheme + "://" + parsedurl.netloc + action
    else:
        idpauthformsubmiturl = action
# Performs the submission of the IdP login form with the above post data
response = session.post(idpauthformsubmiturl, data=payload, verify=sslverification)

# Multi-Factor Authentication (MFA) Handle
# Depending upon the MFA provider, you may replace the string found in response.text to identify the MFA
# if 'SAMLResponse' not in response.text:
#     print('For security reasons, we require additional information to verify your account.\n')
#     print('Please enter your One Time Password (OTP):')
#     otp = getpass.getpass()
#     formsoup = BeautifulSoup(response.text.decode('utf8'), "html.parser")
#     payload = {}
#     for inputtag in formsoup.find_all(re.compile('(INPUT|input)')):
#         name = inputtag.get('name','')
#         value = inputtag.get('value','')
#         if "password" in name.lower():
#             # OTP
#             payload[name] = otp
#         else:
#             #Simply populate the parameter with the existing value (picks up hidden fields in the login form)
#             payload[name] = value
# response = session.post(idpauthformsubmiturl, data=payload, verify=sslverification)


# Debug the response if needed
#print (response.text)

# Overwrite and delete the credential variables, just for safety
username = '##############################################'
password = '##############################################'
del username
del password

# Decode the response and extract the SAML assertion
soup = BeautifulSoup(response.text.decode('utf8'), "html.parser")
assertion = ''

# Look for the SAMLResponse attribute of the input tag (determined by
# analyzing the debug print lines above)
for inputtag in soup.find_all('input'):
    if(inputtag.get('name') == 'SAMLResponse'):
        #print(inputtag.get('value'))
        assertion = inputtag.get('value')

# Better error handling is required for production use.
if (assertion == ''):
    #TODO: Insert valid error checking/handling
    print 'Response did not contain a valid SAML assertion'
    sys.exit(0)

# Debug only
# print(base64.b64decode(assertion))

# Parse the returned assertion and extract the authorized roles
awsroles = []
root = ET.fromstring(base64.b64decode(assertion))
for saml2attribute in root.iter('{urn:oasis:names:tc:SAML:2.0:assertion}Attribute'):
    if (saml2attribute.get('Name') == 'https://aws.amazon.com/SAML/Attributes/Role'):
        for saml2attributevalue in saml2attribute.iter('{urn:oasis:names:tc:SAML:2.0:assertion}AttributeValue'):
            awsroles.append(saml2attributevalue.text)

# Note the format of the attribute value should be role_arn,principal_arn
# but lots of blogs list it as principal_arn,role_arn so let's reverse
# them if needed
for awsrole in awsroles:
    chunks = awsrole.split(',')
    if'saml-provider' in chunks[0]:
        newawsrole = chunks[1] + ',' + chunks[0]
        index = awsroles.index(awsrole)
        awsroles.insert(index, newawsrole)
        awsroles.remove(awsrole)

# If I have more than one role, ask the user which one they want,
# otherwise just proceed
print ""
selectedroleindexes = [int]
role_arns = []
principal_arns = []
if len(awsroles) > 1:
    i = 0
    if args.Roles:
        selectedroleindexes = args.Roles.split(' ')
    else:
        print "Please choose the role(s) you would like to assume, multiple roles should be delimited by single space. You may also enter ALL to assume all roles:"
        for awsrole in awsroles:
            print '[', i, ']: ', awsrole.split(',')[0]
            i += 1
    
        print "Selection(s): ",
        selectedroleindexes = raw_input().split(' ')

    for selectedroleindex in selectedroleindexes:
        # If the selection is ALL, save all roles as profiles
        if selectedroleindex.upper() == 'ALL':
            for i in range(0,len(awsroles)):
                role_arns.append(awsroles[i].split(',')[0])
                principal_arns.append(awsroles[i].split(',')[1])
            break
        else:
            try:
                # Basic sanity check of input
                if int(selectedroleindex) > (len(awsroles) - 1):
                    print 'You selected an invalid role index, please try again'
                    sys.exit(0)
                else:
                    role_arns.append(awsroles[int(selectedroleindex)].split(',')[0])
                    principal_arns.append(awsroles[int(selectedroleindex)].split(',')[1])
            except Exception:
                print ''

else:
    role_arns.append(awsroles[0].split(',')[0])
    principal_arns.append(awsroles[0].split(',')[1])
samlprofiles = []
for i in range(0,len(role_arns)):
    # Use the assertion to get an AWS STS token using Assume Role with SAML
    conn = boto.sts.connect_to_region(region)
    token = conn.assume_role_with_saml(role_arns[i], principal_arns[i], assertion)
    account_alias = ''

    try:
        iamconn = boto.iam.connect_to_region(region,
                                         aws_access_key_id=token.credentials.access_key,
                                         aws_secret_access_key=token.credentials.secret_key,
                                         security_token=token.credentials.session_token)
        account_alias = iamconn.get_account_alias()['list_account_aliases_response']['list_account_aliases_result']['account_aliases'][0]
    except:
        print('')
    # Write the AWS STS token into the AWS credential file
    home = expanduser("~")
    filename = home + awsconfigfile
                                         
    # Read in the existing config file
    config = ConfigParser.RawConfigParser()
    config.read(filename)
                                         
    # Put the credentials into a specific profile instead of clobbering
    # the default credentials
                                         
    if account_alias != '':
        samlprofile = account_alias + '_' + role_arns[i].split('/')[1]
    else:
        samlprofile = role_arns[i].split('/')[0].split(':')[4] + '_' + role_arns[i].split('/')[1]

    samlprofiles.append(samlprofile)

    if not config.has_section(samlprofile):
        config.add_section(samlprofile)

    config.set(samlprofile, 'output', outputformat)
    config.set(samlprofile, 'region', region)
    config.set(samlprofile, 'aws_access_key_id', token.credentials.access_key)
    config.set(samlprofile, 'aws_secret_access_key', token.credentials.secret_key)
    config.set(samlprofile, 'aws_session_token', token.credentials.session_token)
    
    # Write the updated config file
    with open(filename, 'w+') as configfile:
        config.write(configfile)

# Give the user some basic info as to what has just happened
print '\n\n----------------------------------------------------------------'
print 'Your new access key pair(s) have been stored in the AWS configuration file {0} under the saml profile.'.format(filename)
print 'Note that they will expire at {0}.'.format(token.credentials.expiration)
print 'After this time you may safely rerun this script to refresh your access key pair(s).'
print 'To use this credential call the AWS CLI with the --profile option (e.g. aws --profile saml ec2 describe-instances).'
print '----------------------------------------------------------------\n\n'

print '\nYour SAML Profiles:\n'
for samlprofile in samlprofiles:
    print samlprofile