#!/usr/bin/env bash

# you MUST source this file to get the exports in your shell!


# https://stackoverflow.com/questions/2683279/how-to-detect-if-a-script-is-being-sourced/28776166#28776166
([[ -n $ZSH_EVAL_CONTEXT && $ZSH_EVAL_CONTEXT =~ :file$ ]] ||
 [[ -n $KSH_VERSION && $(cd "$(dirname -- "$0")" &&
    printf '%s' "${PWD%/}/")$(basename -- "$0") != "${.sh.file}" ]] ||
 [[ -n $BASH_VERSION ]] && (return 0 2>/dev/null)) &&
    SOURCED=1 || SOURCED=0


# https://stackoverflow.com/questions/592620/how-can-i-check-if-a-program-exists-from-a-bash-script
if ! command -v python3 >/dev/null 2>&1; then
    echo "python3 could not be found!"
    [[ $SOURCED -eq 1 ]] && return 0 || exit 0
fi

if ! command -v aws >/dev/null 2>&1; then
    echo "aws could not be found!"
    [[ $SOURCED -eq 1 ]] && return 0 || exit 0
fi

if ! command -v jq >/dev/null 2>&1; then
    echo "jq could not be found!"
    [[ $SOURCED -eq 1 ]] && return 0 || exit 0
fi


PYCODE=$(cat << EOF

#!/usr/bin/env python3

# Follows the directions from AWS on using MFA with CLI:
#  https://aws.amazon.com/premiumsupport/knowledge-center/authenticate-mfa-cli/

# That means: call STS with an ARN and an MFA token,
#  with the response, populate an MFA section in aws creds file


import os
import json
import configparser
import argparse

from pathlib import Path


def parse_cli_args():
    argparser = argparse.ArgumentParser(prog='aws-cli-mfa', description='Login to AWS CLI using MFA token with STS')

    argparser.add_argument('profile_arn', metavar='profile-arn',
        help='the AWS ARN for your MFA profile')
    argparser.add_argument('mfa_token', metavar='mfa-token',
        help='the MFA token from your authenticator app for the MFA profile your ARN is for')
    argparser.add_argument('-p', '--aws-profile',
        help='AWS profile to use when contacting STS to get MFA credentials')
    argparser.add_argument('-s', '--aws-creds-mfa-section', default='mfa',
        help='section to save MFA credentials to in AWS credentials file (default: mfa)')
    argparser.add_argument('-f', '--aws-creds-file', default=str(Path.home())+'/.aws/credentials',
        help='file path to AWS credentials file (default: ~/.aws/credentials)')
    argparser.add_argument('-e', '--aws-env-vars', action='store_true',
        help='export/set AWS access/secret/session env vars instead of updating creds file')
    argparser.add_argument('-x', '--no-export-profile', action='store_true',
        help='do not export/set AWS_PROFILE to the one generated by STS')
    argparser.add_argument('-d', '--lifetime-duration', type=int,
        help='set the time, in seconds, that the access should last (default 12 hours)')

    return argparser.parse_args()


def gen_sts_cmd(cli_args):
    sts_args = []

    if cli_args.aws_profile:
        sts_args.append(f'--profile {cli_args.aws_profile}')

    sts_args.append(f'--serial-number {cli_args.profile_arn}')
    sts_args.append(f'--token-code {cli_args.mfa_token}')

    if cli_args.lifetime_duration:
        sts_args.append(f'--duration-seconds {cli_args.lifetime_duration}')

    return f"aws sts get-session-token {' '.join(sts_args)}"


def write_config(creds_file, mfa_profile_section, creds):
    config = configparser.ConfigParser()
    config.read(creds_file)

    config[mfa_profile_section] = creds

    with open(creds_file, 'w') as configfile:
       config.write(configfile)


def apply_sts_json(
        sts_json,
        use_env_vars,
        creds_file,
        mfa_profile_section,
        export_profile
    ):

    envvars = {}
    creds = {
        'AWS_ACCESS_KEY_ID': sts_json["Credentials"]["AccessKeyId"],
        'AWS_SECRET_ACCESS_KEY': sts_json["Credentials"]["SecretAccessKey"],
        'AWS_SESSION_TOKEN': sts_json["Credentials"]["SessionToken"],
    }

    if use_env_vars:
        envvars = creds.copy()
    else:
        write_config(creds_file, mfa_profile_section, creds)

        if export_profile:
            envvars['AWS_PROFILE'] = mfa_profile_section

    return envvars


def build_response():
    response = {}

    cli_args = parse_cli_args()
    sts_cmd = gen_sts_cmd(cli_args)
    response['sts_cmd'] = sts_cmd

    sts_output = os.popen(sts_cmd).read()
    try:
        sts_json = json.loads(sts_output)

        try:
            response['envvars'] = apply_sts_json(
                sts_json,
                cli_args.aws_env_vars,
                cli_args.aws_creds_file,
                cli_args.aws_creds_mfa_section,
                not cli_args.no_export_profile
            )
        except Exception as ex:
            response['output'] = str(ex)
    except:
        response['output'] = sts_output

    return response

if __name__ == '__main__':
    print(json.dumps(build_response()))
EOF
)

RESPONSE=$(COLUMNS=999 /usr/bin/env python3 -c "$PYCODE" $@)


if [[ "$RESPONSE" == usage* ]] ; then
    echo -E "$RESPONSE"
    [[ $SOURCED -eq 1 ]] && return 0 || exit 0
fi

if ! echo -E "$RESPONSE" | jq -e . >/dev/null 2>&1; then
    echo "JSON parsing failed:"
    echo -E "$RESPONSE"
    [[ $SOURCED -eq 1 ]] && return 1 || exit 1
fi

STS_CMD=$(echo -E "$RESPONSE" | jq -r .sts_cmd)
OUTPUT=$(echo -E "$RESPONSE" | jq -r .output)

[[ ! -z "$STS_CMD" ]] && echo "$STS_CMD"
[[ ! -z "$OUTPUT" && "null" != "$OUTPUT" ]] && echo "$OUTPUT"


# https://unix.stackexchange.com/questions/88850/precedence-of-the-shell-logical-operators
[[ "null" == "$(echo -E "$RESPONSE" | jq -r '.envvars')" ]] &&
    { [[ $SOURCED -eq 1 ]] && return 0 || exit 0; }

KEYS=( $(echo -E "$RESPONSE" | jq -r '.envvars | keys[]') )

[[ ${#KEYS[@]} -ge 0 ]] && [[ $SOURCED -ne 1 ]] &&
    echo "You must source this file to get the exports in your shell" &&
    exit 1

for key in "${KEYS[@]}" ; do
    value=$(echo -E "$RESPONSE" | jq -r ".envvars.$key")
    export $key=$value
    echo "Set env var: $key"
done
