#!/usr/bin/env python3
import os
import json
import argparse
import subprocess
import configparser
import tty
from shutil import copy2
import time
import sys

import termios


def readchar():
    fd = sys.stdin.fileno()
    old_settings = termios.tcgetattr(fd)
    try:
        tty.setraw(sys.stdin.fileno())
        ch = sys.stdin.read(1)
    finally:
        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    return ch


parser = argparse.ArgumentParser(description='Update your AWS CLI Token')
parser.add_argument('--credential-path', help='path to the aws credentials file',
                    default=os.path.expanduser('~/.aws/credentials'))

args = parser.parse_args()

config = configparser.ConfigParser()
config.read(args.credential_path)

sections = config.sections()
i = 1
print("Select a profile to generate token using mfa:")
for section in sections:
    print(f"{i}. {section}")
    i = i + 1
selection = readchar()
try:
    profile = sections[int(selection) - 1]
except ValueError:
    print("Invalid input")
    exit(0)

profile_out = profile + '-mfa'

if profile_out in config.sections():
    if 'aws_arn_mfa' not in config[profile_out]:
        config[profile_out]['aws_arn_mfa'] = input("Enter mfa arn: ")
else:
    config.add_section(profile_out)
    config[profile_out]['aws_arn_mfa'] = input("Enter mfa arn: ")

token = input("Enter MFA token:")
result = subprocess.run(
    ['aws', 'sts', 'get-session-token', '--profile', profile, '--serial-number', config[profile_out]['aws_arn_mfa'],
     '--token-code', token], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if result.returncode != 0:
    parser.error(result.stderr.decode('utf-8').strip('\n'))

credentials = json.loads(result.stdout.decode('utf-8'))['Credentials']

config[profile_out]['aws_access_key_id'] = credentials['AccessKeyId']
config[profile_out]['aws_secret_access_key'] = credentials['SecretAccessKey']
config[profile_out]['aws_session_token'] = credentials['SessionToken']

print("Set as default\n Setting config as default will overwrite current default profile! \n Set as default (y/n)?")

if readchar() == 'y':
    profile_out = 'default'
    config[profile_out]['aws_arn_mfa'] = input("Enter mfa arn: ")
    config[profile_out]['aws_access_key_id'] = credentials['AccessKeyId']
    config[profile_out]['aws_secret_access_key'] = credentials['SecretAccessKey']
    config[profile_out]['aws_session_token'] = credentials['SessionToken']

copy2(args.credential_path, args.credential_path.rstrip('/') + str(time.time()) + '_bak')

with open(args.credential_path, 'w') as configFile:
    config.write(configFile)

print('Saved {} credentials to {}'.format(profile_out, args.credential_path))