assume_role_for_adfs.py

#!/usr/bin/env python

import os
import requests
import boto3
import base64
import argparse
from getpass import getpass
from bs4 import BeautifulSoup as bs


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--adfs-fqdn", required=True, dest="ADFS_FQDN")
    parser.add_argument("--username", required=True, dest="USERNAME")
    parser.add_argument("--duration", dest="DURATION")
    args = parser.parse_args()

    PASSWORD = getpass(args.USERNAME + "'s Password: ")

    response = request(args.ADFS_FQDN, args.USERNAME, PASSWORD)
    soup = bs(response.text, 'html.parser')
    saml_response = soup.find(attrs={"name":"SAMLResponse"})["value"]
    role, principal = get_saml_params(saml_response)
    if args.DURATION:
        json = get_security_token(role, principal, saml_response, args.DURATION)
    else:
        json = get_security_token(role, principal, saml_response)
    output_aws_credentials(json)


def request(adfs_name, username, password):
    url = 'https://' + adfs_name + '/adfs/ls/idpinitiatedsignon/?loginToRp=urn:amazon:webservices'
    data = 'Username=' + username + '&Password=' + password + '&AuthMethod=FormsAuthentication'
    requests.packages.urllib3.disable_warnings()
    return requests.post(url, data, verify=False)


def get_saml_params(saml_response):
    soup = bs(base64.b64decode(saml_response), "html.parser")
    arns = soup.findAll("attribute",attrs={"name":"https://aws.amazon.com/SAML/Attributes/Role"})[0].get_text("attributevalue")
    role = ""
    principal = ""
    for arn in arns.split(","):
        if "role/" in arn:
            role = arn.strip()
        elif "saml-provider/" in arn:
            principal = arn.strip()
    return role, principal


def get_security_token(role, principal, saml_assertion, duration=3600):
    if isinstance(duration, str):
        duration = int(duration)
    client = boto3.client('sts')
    return client.assume_role_with_saml(RoleArn=role, PrincipalArn=principal, SAMLAssertion=saml_assertion, DurationSeconds=duration)


def output_aws_credentials(json):
    print("Execute following commands to set AWS credentials")
    print("-----")
    print("export AWS_ACCESS_KEY_ID=" + json["Credentials"]["AccessKeyId"])
    print("export AWS_SECRET_ACCESS_KEY=" + json["Credentials"]["SecretAccessKey"])
    print("export AWS_SESSION_TOKEN=" + json["Credentials"]["SessionToken"])


if __name__ == "__main__":
    main()