#!/usr/local/bin/python3

import sys
import subprocess
import json
import os
from os.path import expanduser
from datetime import datetime

mfa_config = expanduser('~') + '/.aws/mfa.json'

def authorize_mfa(credentials):
    try:
        arn = credentials['arn']
    except KeyError:
        arn = input('arn of the mfa device: ')

    token = input('mfa token code: ')

    try:
        result = subprocess.run(
            ['aws', 'sts', 'get-session-token', '--serial-number', arn, '--token-code', token],
            capture_output=True,
            text=True,
            check=True
        )

    except subprocess.CalledProcessError as err:
        sys.exit(err.stderr)

    credentials = json.loads(result.stdout)['Credentials']

    with open(mfa_config, 'w') as f:
        json.dump({'arn': arn, **credentials}, f, indent=2)

    return credentials


try:
    with open(mfa_config, 'r') as f:
        credentials = json.load(f)
except (FileNotFoundError, json.decoder.JSONDecodeError):
    credentials = {}

try:
    expiration = datetime.fromisoformat(credentials['Expiration'])
except KeyError:
    expiration = datetime.fromtimestamp(0).astimezone()

if datetime.now().astimezone() > expiration:
    credentials = authorize_mfa(credentials)

aws_env = {
    'AWS_ACCESS_KEY_ID': credentials['AccessKeyId'],
    'AWS_SECRET_ACCESS_KEY': credentials['SecretAccessKey'],
    'AWS_SESSION_TOKEN': credentials['SessionToken'],
}

if len(sys.argv) > 1:
    subprocess.run(sys.argv[1:], env=os.environ | aws_env)
