aws-okta-wrapper.py

#!/usr/bin/env python
"""Wrapper tool for writing aws credentials into profiles."""
import argparse
import configparser
import platform
import subprocess
from os import path
from shutil import which

parser = argparse.ArgumentParser(description='Wrapper script for aws-okta tool.')
parser.add_argument(
    '-p',
    '--profile',
    required=True,
    default=False,
    type=str,
    dest="profile",
    metavar="<profile>",
    help="AWS Profile for credentials")
args = parser.parse_args()
print("Requesting credentials for profile: {0}".format(args.profile))

aws_file = path.join(path.expanduser("~"), '.aws/credentials')
config = configparser.ConfigParser()
config.read([aws_file])

if which("aws-okta") is None:
    if platform.system().lower() == "darwin":
        subprocess.check_output("brew install aws-okta", shell=True)
    else:
        raise Exception("System is not MAC and aws-okta tool is not installed. "
                        "Please, install: https://github.com/segmentio/aws-okta#installing")

output = subprocess.check_output("aws-okta env {0}".format(args.profile), shell=True)
nice_output = output.decode('UTF-8').rstrip().replace('export ', '')

data = {}
for item in nice_output.split('\n'):
    key, value = item.split('=', maxsplit=1)
    data[key] = value

config[args.profile] = data

with open(aws_file, 'w') as configfile:  # save
    config.write(configfile)

print("Updated credentials for profile: {0}".format(args.profile))