from __future__ import print_function
from __future__ import unicode_literals

import argparse
import os

from six.moves import configparser


def get_aws_credential_files_from_env():
    """Extract credential file paths from environment variables."""
    files = set()
    for env_var in {'AWS_CONFIG_FILE', 'AWS_CREDENTIAL_FILE',
                    'AWS_SHARED_CREDENTIALS_FILE', 'BOTO_CONFIG'}:
        try:
            files.add(os.environ[env_var])
        except KeyError:
            pass
    return files


def get_aws_secrets_from_env():
    """Extract AWS secrets from environment variables."""
    keys = set()
    for env_var in {'AWS_SECRET_ACCESS_KEY', 'AWS_SECURITY_TOKEN',
                    'AWS_SESSION_TOKEN'}:
        try:
            keys.add(os.environ[env_var])
        except KeyError:
            pass
    return keys


def get_aws_secrets_from_file(credentials_file):
    """Extract AWS secrets from configuration files.

    Read an ini-style configuration file and return a set with all found AWS
    secret access keys.
    """
    aws_credentials_file_path = os.path.expanduser(credentials_file)
    if not os.path.exists(aws_credentials_file_path):
        return set()

    parser = configparser.ConfigParser()
    try:
        parser.read(aws_credentials_file_path)
    except configparser.MissingSectionHeaderError:
        return set()

    keys = set()
    for section in parser.sections():
        for var in {'aws_secret_access_key', 'aws_security_token',
                    'aws_session_token'}:
            try:
                keys.add(parser.get(section, var))
            except configparser.NoOptionError:
                pass
    return keys


def check_file_for_aws_keys(filenames, keys):
    """Check if files contain AWS secrets.

    Return a list of all files containing AWS secrets and keys found, with all
    but the first four characters obfuscated to ease debugging.
    """
    bad_files = []

    for filename in filenames:
        with open(filename, 'r') as content:
            text_body = content.read()
            for key in keys:
                # naively match the entire file, low chance of incorrect
                # collision
                if key in text_body:
                    bad_files.append({'filename': filename,
                                      'key': key[:4].ljust(32, str('*'))})
    return bad_files


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument('filenames', nargs='+', help='Filenames to run')
    parser.add_argument(
        '--credentials-file',
        dest='credential_files',
        action='append',
        default=['~/.aws/config', '~/.aws/credentials', '/etc/boto.cfg',
                 '~/.boto'],
        help=(
            'Location of additional AWS credential files from which to get '
            'secret keys from'
        )
    )
    args = parser.parse_args(argv)

    credential_files = set(args.credential_files)

    # Add the credentials files configured via environment variables to the set
    # of files to to gather AWS secrets from.
    credential_files |= get_aws_credential_files_from_env()

    keys = set()
    for credential_file in credential_files:
        keys |= get_aws_secrets_from_file(credential_file)

    # Secrets might be part of environment variables, so add such secrets to
    # the set of keys.
    keys |= get_aws_secrets_from_env()

    if not keys:
        print('No AWS keys were found in the configured credential files and '
              'environment variables.\nPlease ensure you have the correct '
              'setting for --credentials-file')
        return 2

    bad_filenames = check_file_for_aws_keys(args.filenames, keys)
    if bad_filenames:
        for bad_file in bad_filenames:
            print('AWS secret found in {filename}: {key}'.format(
                **bad_file))
        return 1
    else:
        return 0


if __name__ == '__main__':
    exit(main())