|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +# !/usr/bin/python |
| 3 | + |
| 4 | +import os |
| 5 | +import sys |
| 6 | +import boto3 |
| 7 | +import requests |
| 8 | +import configparser |
| 9 | +import base64 |
| 10 | +import xml.etree.ElementTree as ET |
| 11 | +import re |
| 12 | +from bs4 import BeautifulSoup |
| 13 | +from urllib.parse import urlparse |
| 14 | +import click |
| 15 | +import yaml |
| 16 | + |
| 17 | + |
| 18 | +CFG = {} |
| 19 | +with open("config.yml", 'r') as ymlfile: |
| 20 | + CFG = yaml.load(ymlfile) |
| 21 | +regions = '\n'.join( |
| 22 | + ["%s) %s" % (i, x['name']) for i, x in enumerate(CFG.get('regions', []))] |
| 23 | +) |
| 24 | + |
| 25 | + |
| 26 | +def saml_assertion(username, password): |
| 27 | + # Initiate session handler |
| 28 | + session = requests.Session() |
| 29 | + url = CFG.get('url_sso_aws') |
| 30 | + |
| 31 | + # Programmatically get the SAML assertion |
| 32 | + # Opens the initial IdP url and follows all of the HTTP302 redirects, and |
| 33 | + # gets the resulting login page |
| 34 | + formresponse = session.get(url, verify=True) |
| 35 | + # Capture the idpauthformsubmiturl, which is the final url after all the 302s |
| 36 | + idpauthformsubmiturl = formresponse.url |
| 37 | + |
| 38 | + # Parse the response and extract all the necessary values |
| 39 | + # in order to build a dictionary of all of the form values the IdP expects |
| 40 | + formsoup = BeautifulSoup(formresponse.text, "html.parser") |
| 41 | + payload = {} |
| 42 | + |
| 43 | + for inputtag in formsoup.find_all(re.compile('(INPUT|input)')): |
| 44 | + name = inputtag.get('name', '') |
| 45 | + value = inputtag.get('value', '') |
| 46 | + if "user" in name.lower(): |
| 47 | + # Make an educated guess that this is the right field for the username |
| 48 | + payload[name] = username |
| 49 | + elif "email" in name.lower(): |
| 50 | + # Some IdPs also label the username field as 'email' |
| 51 | + payload[name] = username |
| 52 | + elif "pass" in name.lower(): |
| 53 | + # Make an educated guess that this is the right field for the password |
| 54 | + payload[name] = password |
| 55 | + else: |
| 56 | + # Simply populate the parameter with the existing value |
| 57 | + # (picks up hidden fields |
| 58 | + # in the login form) |
| 59 | + payload[name] = value |
| 60 | + # Debug the parameter payload if needed |
| 61 | + # Use with caution since this will print sensitive output to the screen |
| 62 | + |
| 63 | + # Some IdPs don't explicitly set a form action, but if one is set we should |
| 64 | + # build the idpauthformsubmiturl by combining the scheme and hostname |
| 65 | + # from the entry url with the form action target |
| 66 | + # If the action tag doesn't exist, we just stick with the |
| 67 | + # idpauthformsubmiturl above |
| 68 | + for inputtag in formsoup.find_all(re.compile('(FORM|form)')): |
| 69 | + action = inputtag.get('action') |
| 70 | + if action: |
| 71 | + parsedurl = urlparse(url) |
| 72 | + idpauthformsubmiturl = '%s://%s/%s' % ( |
| 73 | + parsedurl.scheme, |
| 74 | + parsedurl.netloc, |
| 75 | + action |
| 76 | + ) |
| 77 | + # Performs the submission of the IdP login form with the above post data |
| 78 | + response = session.post( |
| 79 | + idpauthformsubmiturl, |
| 80 | + data=payload, |
| 81 | + verify=True |
| 82 | + ) |
| 83 | + |
| 84 | + # Decode the response and extract the SAML assertion |
| 85 | + soup = BeautifulSoup(response.text, "html.parser") |
| 86 | + assertion = '' |
| 87 | + |
| 88 | + # Look for the SAMLResponse attribute of the input tag (determined by |
| 89 | + # analyzing the debug print lines above) |
| 90 | + for inputtag in soup.find_all('input'): |
| 91 | + if(inputtag.get('name') == 'SAMLResponse'): |
| 92 | + assertion = inputtag.get('value') |
| 93 | + |
| 94 | + # Better error handling is required for production use. |
| 95 | + if (assertion == ''): |
| 96 | + # TODO: Insert valid error checking/handling |
| 97 | + print('Response did not contain a valid SAML assertion') |
| 98 | + sys.exit(0) |
| 99 | + |
| 100 | + # Parse the returned assertion and extract the authorized roles |
| 101 | + awsroles = [] |
| 102 | + root = ET.fromstring(base64.b64decode(assertion)) |
| 103 | + for saml2attribute in root.iter('{urn:oasis:names:tc:SAML:2.0:assertion}Attribute'): |
| 104 | + if (saml2attribute.get('Name') == 'https://aws.amazon.com/SAML/Attributes/Role'): |
| 105 | + saml2attribute = saml2attribute.iter( |
| 106 | + '{urn:oasis:names:tc:SAML:2.0:assertion}AttributeValue' |
| 107 | + ) |
| 108 | + for saml2attributevalue in saml2attribute: |
| 109 | + awsroles.append(saml2attributevalue.text) |
| 110 | + # Note the format of the attribute value should be role_arn,principal_arn |
| 111 | + # but lots of blogs list it as principal_arn,role_arn so let's reverse |
| 112 | + # them if needed |
| 113 | + for awsrole in awsroles: |
| 114 | + chunks = awsrole.split(',') |
| 115 | + if'saml-provider' in chunks[0]: |
| 116 | + newawsrole = chunks[1] + ',' + chunks[0] |
| 117 | + index = awsroles.index(awsrole) |
| 118 | + awsroles.insert(index, newawsrole) |
| 119 | + awsroles.remove(awsrole) |
| 120 | + |
| 121 | + # TODO |
| 122 | + # If I have more than one role, ask the user which one they want, |
| 123 | + # otherwise just proceed |
| 124 | + # print ("") |
| 125 | + # if len(awsroles) > 1: |
| 126 | + # i = 0 |
| 127 | + # print ("Please choose the role you would like to assume:") |
| 128 | + # for awsrole in awsroles: |
| 129 | + # print ('[', i, ']: ', awsrole.split(',')[0]) |
| 130 | + # i += 1 |
| 131 | + # print ("Selection: ",) |
| 132 | + # selectedroleindex = raw_input() |
| 133 | + |
| 134 | + # # Basic sanity check of input |
| 135 | + # if int(selectedroleindex) > (len(awsroles) - 1): |
| 136 | + # print ('You selected an invalid role index, please try again') |
| 137 | + # sys.exit(0) |
| 138 | + |
| 139 | + # role_arn = awsroles[int(selectedroleindex)].split(',')[0] |
| 140 | + # principal_arn = awsroles[int(selectedroleindex)].split(',')[1] |
| 141 | + # else: |
| 142 | + # role_arn = awsroles[0].split(',')[0] |
| 143 | + # principal_arn = awsroles[0].split(',')[1] |
| 144 | + role_arn = awsroles[0].split(',')[0] |
| 145 | + principal_arn = awsroles[0].split(',')[1] |
| 146 | + client = boto3.client('sts') |
| 147 | + token = client.assume_role_with_saml( |
| 148 | + RoleArn=role_arn, |
| 149 | + PrincipalArn=principal_arn, |
| 150 | + SAMLAssertion=assertion |
| 151 | + ) |
| 152 | + return token |
| 153 | + |
| 154 | + |
| 155 | +@click.command() |
| 156 | +@click.option('--email', prompt=True, help="Your email") |
| 157 | +@click.option('--password', prompt=True, hide_input=True, help="Your Password") |
| 158 | +@click.option('--region', |
| 159 | + prompt=regions + "\nSelected Region", |
| 160 | + help="Selected a region:\n\n" + regions |
| 161 | +) |
| 162 | +def cli(email, password, region): |
| 163 | + |
| 164 | + token = saml_assertion(email, password) |
| 165 | + path_filename = os.path.join( |
| 166 | + os.path.expanduser(CFG.get('path_file')), |
| 167 | + CFG.get('filename') |
| 168 | + ) |
| 169 | + |
| 170 | + config = configparser.RawConfigParser() |
| 171 | + config.read(path_filename) |
| 172 | + config.set(configparser.DEFAULTSECT, 'output', CFG.get('outputformat')) |
| 173 | + config.set( |
| 174 | + configparser.DEFAULTSECT, |
| 175 | + 'region', |
| 176 | + CFG['regions'][int(region)]['id'] |
| 177 | + ) |
| 178 | + config.set( |
| 179 | + configparser.DEFAULTSECT, |
| 180 | + 'aws_access_key_id', |
| 181 | + token['Credentials']['AccessKeyId'] |
| 182 | + ) |
| 183 | + config.set( |
| 184 | + configparser.DEFAULTSECT, |
| 185 | + 'aws_secret_access_key', |
| 186 | + token['Credentials']['SecretAccessKey'] |
| 187 | + ) |
| 188 | + config.set( |
| 189 | + configparser.DEFAULTSECT, |
| 190 | + 'aws_session_token', |
| 191 | + token['Credentials']['SessionToken'] |
| 192 | + ) |
| 193 | + if not os.path.exists(os.path.expanduser(CFG.get('path_file'))): |
| 194 | + os.mkdir(os.path.expanduser(CFG.get('path_file'))) |
| 195 | + |
| 196 | + # Write the updated config file |
| 197 | + with open(path_filename, 'w+') as configfile: |
| 198 | + config.write(configfile) |
| 199 | + # Give the user some basic info as to what has just happened |
| 200 | + print ('\n----------------------------------------------------------------') |
| 201 | + print ("""Your new access key pair has been stored in the AWS configuration file {0} under the saml profile.""".format(path_filename)) |
| 202 | + print ('Note that it will expire at {0}.'.format(token['Credentials']['Expiration'])) |
| 203 | + print ('After this time, you may safely rerun this script to refresh your access key pair.') |
| 204 | + print ('To use this credential, call the AWS CLI with the --profile option (e.g. aws --profile saml ec2 describe-instances).') |
| 205 | + print ('----------------------------------------------------------------\n\n') |
0 commit comments