#!/usr/bin/env fish

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

if contains (status current-command) $IS_SOURCED_COMMANDS; or test -n "$IS_SOURCED"
    set SOURCED 1
else
    set SOURCED 0
end

set -l REQUIREMENTS python3 aws jq

for requirement in $REQUIREMENTS
    if not command -v $requirement &> /dev/null
        echo "$requirement could not be found!"
        exit 0
    end
end

set -l PYCODE "

#!/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()))
"

set -l RESPONSE (/usr/bin/env python3 -c "$PYCODE" $argv)

if string match -r "^usage" "$RESPONSE" &> /dev/null
    echo "$RESPONSE"
    exit 0
end

if not echo "$RESPONSE" | jq -e . &> /dev/null
    echo "JSON parsing failed:"
    echo "$RESPONSE"
    exit 1
end

set -l STS_CMD (echo "$RESPONSE" | jq -r ".sts_cmd") 
set -l OUTPUT (echo "$RESPONSE" | jq -r ".output" | string collect)

if not test null = "$STS_CMD"
    echo "$STS_CMD"
end

if not test null = "$OUTPUT"
    echo "$OUTPUT"
end

set -l ENVVARS (echo "$RESPONSE" | jq -r ".envvars")

if test null = "$ENVVARS"
    exit 0
end

set -l KEYS (echo "$RESPONSE" | jq -r ".envvars | keys[]")
if test (count $KEYS) -ge 0 -a $SOURCED -ne 1
    echo "You must source this file to get the exports in your shell"
    exit 1
end

for key in $KEYS
    set value (echo "$RESPONSE" | jq -r ".envvars.$key")
    set -x $key $value
    echo "Set env var: $key"
end

