generate-sts-mfa-token

#!/usr/bin/env python

# Copyright 2013 42Lines, Inc.
# Original Author: Jim Browne
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from boto import sts
import boto
import json
import logging
import optparse as OP
import sys
import os

if __name__ == '__main__':

    # Boto 2.7.0 or later is needed for assume_role in STS
    desired = '2.7.0'
    try:
        from pkg_resources import parse_version
        if parse_version(boto.__version__) < parse_version(desired):
            print('Boto version %s or later is required' % desired)
            print('Try: sudo easy_install boto')
            sys.exit(-1)
    except (AttributeError, NameError):
        print('Boto version %s or later is required' % desired)
        print('Try: sudo easy_install boto')
        sys.exit(-1)

    description = '''Assume the given role and return a JSON blob with
    information about the STS token
'''
    usage = "usage: %prog --account ACCOUNT --user USER --token TOKEN [options]"
    p = OP.OptionParser(description=description, usage=usage)
    p.add_option(
        "-d",
        "--debug",
        action="count",
        dest="debug",
        help="Output additional information."
    )
    p.add_option(
        "--mfa",
        dest="mfa",
        help=("MFA serial number (e.g."
              "arn:aws:iam::987654321012:mfa/jbrowne)")
    )
    p.add_option(
        "--token",
        dest="token",
        help="Current MFA token"
    )
    p.add_option(
        "--account",
        dest="account",
        help=("Account to use in generating ARN string OR a raw ARN string "
              "minus the username.")
    )
    p.add_option(
        "--config",
        dest="config",
        default="/etc/default/aws-sts-helpers/accounts.json",
        help="JSON format file containing ARN strings for accounts"
    )
    p.add_option(
        "--user",
        dest="user",
        help="The username to append to the ARN string"
    )
    p.add_option(
        "--region",
        dest="region",
        default='us-east-1',
        help="AWS region (default %default).  Possible only one region exists"
    )
    p.add_option(
        "--duration",
        dest="duration",
        default=3600,
        type='int',
        help="lifetime of token in seconds, 900-129600; default %default"
    )
    p.add_option(
        "--shell",
        action="store_true",
        dest="shell",
        default=False,
        help="Emit copy/pastable environment variables rather than JSON"
    )
    p.add_option(
        "--cachefiles",
        action="store_true",
        dest="cachefiles",
        default=False,
        help=("Update cache files of keys in /home/USER/aws-creds rather than",
              " emit JSON")
    )

    # Parse user input and sanity check
    (opts, args) = p.parse_args()

    if opts.debug:
        logging.basicConfig(level=logging.DEBUG)

    if opts.token is None:
        sys.stderr.write("Must specify --token\n")
        sys.exit(1)

    if (opts.account is None or opts.user is None) and opts.mfa is None:
        sys.stderr.write("Must specify (--account --user) or --mfa\n")
        sys.exit(1)

    try:
        with open(opts.config) as config_file:
            accounts = json.load(config_file)
    except IOError:
        sys.stderr.write("Unable to read file %s\n" % opts.config)
        sys.exit(1)

    if opts.mfa is not None:
        # Raw ARN string
        arn = opts.mfa
    elif opts.account in accounts:
        arn = accounts[opts.account] + 'mfa/'
        arn += opts.user
    else:
        sys.stderr.write("Unknown account %s\n" % opts.account)
        sys.exit(1)

    if opts.duration < 900 or opts.duration > 129600:
        sys.stderr.write("--duration must be between 900 and 129600\n")
        sys.exit(1)

    sts_conn = sts.connect_to_region(opts.region)

    cred = sts_conn.get_session_token(mfa_serial_number=arn,
                                      mfa_token=opts.token,
                                      duration=opts.duration)

    result = {}
    for datum in ['access_key', 'secret_key', 'session_token', 'expiration']:
        result[datum] = getattr(cred, datum)

    if opts.shell:
        print('export AWS_TOKEN_EXPIRATION=' + result['expiration'])
        # ec2 cli environment variables
        print('export AWS_ACCESS_KEY=' + result['access_key'])
        print('export AWS_SECRET_KEY=' + result['secret_key'])
        print('export AWS_DELEGATION_TOKEN=' + result['session_token'])
        # boto and aws-cli environment variables
        print('export AWS_ACCESS_KEY_ID=' + result['access_key'])
        print('export AWS_SECRET_ACCESS_KEY=' + result['secret_key'])
        # aws-cli specific environment variables
        print('export AWS_SECURITY_TOKEN=' + result['session_token'])
        print('export AWS_DEFAULT_REGION=' + opts.region)
    elif opts.cachefiles:
        base = '%s/aws-creds/.cache.%s.' % (os.environ['HOME'], opts.account)
        files = [f for f in result if f != 'expiration']
        for f in files:
            with open(base + f, 'w') as fh:
                fh.write(result[f])
        with open(base + 'expiration', 'w') as fh:
            fh.write("%s" % (int(opts.duration / 60)))
    else:
        print(json.dumps(result, sort_keys=True, indent=4,
                         separators=(',', ': ')))