generate-sts-role-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 optparse as OP
import sys

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 --role ROLE [options]"
    p = OP.OptionParser(description=description, usage=usage)
    p.add_option(
        "-d",
        "--debug",
        action="count",
        dest="debug",
        help="Output additional information."
    )
    p.add_option(
        "--account",
        dest="account",
        help=("Account to use in generating ARN string OR a raw ARN string "
              "minus the role name.")
    )
    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(
        "--role",
        dest="role",
        help=("ARN of the role to assume (e.g. "
              "aau-temporary-escalation)")
    )
    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, 1-3600; default %default"
    )
    p.add_option(
        "--token",
        dest="token",
        help="Security token if using STS temporary credentials"
    )
    p.add_option(
        "--access-key",
        dest="accessKey",
        help="Use this access key instead of the environment variable"
    )
    p.add_option(
        "--secret-key",
        dest="secretKey",
        help="Use this secret key instead of the environment variable"
    )
    p.add_option(
        "--shell",
        action="store_true",
        dest="shell",
        default=False,
        help="Emit copy/pastable environment variables rather than JSON"
    )

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

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

    if (opts.accessKey is None) != (opts.secretKey is None):
        sys.stderr.write("Must specify both --access-key and --secret-key or neither\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.account in accounts:
        arn = accounts[opts.account] + 'role/'
    else:
        # Raw ARN string
        arn = opts.account

    arn += opts.role

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

    if opts.accessKey is not None:
        sts_conn = sts.connect_to_region(opts.region,
                                         aws_access_key_id=opts.accessKey,
                                         aws_secret_access_key=opts.secretKey,
                                         security_token=opts.token)
    elif opts.token is not None:
        sts_conn = sts.connect_to_region(opts.region,
                                         security_token=opts.token)
    else:
        sts_conn = sts.connect_to_region(opts.region)

    assume = sts_conn.assume_role(role_arn=arn, role_session_name='foo')

    cred = assume.credentials

    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)
    else:
        print(json.dumps(result, sort_keys=True, indent=4, separators=(',', ': ')))