#!/usr/bin/env python3

import boto3, sys, getopt, getpass, configparser
from os.path import expanduser

home = expanduser("~")
aws_creds = home + "/.aws/credentials"

def usage():
   print(
      "Usage: mfa [OPTIONS] PROFILE\n\n"
      "Uses permanent access key and secret key stored in "
      "[PROFILE] within " + aws_creds + " file to generate temporary\n"
      "credentials in '[PROFILE]-mfa' for AWSCLI/SDK access.\n\n"
      "   -h, --help             display this message\n"
      "   -u, --user             set name of AWS user associated with MFA device\n"
      "   -a, --account          set AWS account number to request temporary credentials from\n"
      "   -t, --target-account   set target AWS account to assume role in\n"
      "   -n, --role-arn        set role name to assume in target account \n"
      "   -r, --region           set default region to associate with temporary credentials\n"
   )

def main(argv):

   try:
      opts, args = getopt.getopt(
         argv,
         "hu:a:r:n:t:p:",
         [
            "user=",
            "account=",
            "region=",
            "role-name=",
            "target-account=",
            "profile="
         ]
      )
   except getopt.GetoptError:
      usage()
      sys.exit(2)

   if len(args) != 1:
      usage()
      print("ERROR: Profile not specified\n")
      sys.exit(2)

   profile = args[0]
   profile_mfa = profile + "-mfa"

   conf = {
      "username" : None,
      "account_number" : None,
      "region" : None,
      "target_account" : None,
      "role_name" : None
   }

   config = configparser.ConfigParser()
   config.read(aws_creds)
   try:
      for c in conf: 
         conf[c] = config.get(profile, c)
   except configparser.NoOptionError as e:
      print(repr(e) + ", using command line parameters")
   except configparser.NoSectionError as e:
      print(
         "\nERROR: " + repr(e) + " found in " + aws_creds +
         " or file not found, please configure profile using "
         "'aws configure --profile " + conf["profile"] + "'\n"
      )
      sys.exit(2)

   for opt, arg in opts:
      if opt in ("-h"):
         usage()
         sys.exit()
      elif opt in ("-u", "--user"):
         conf["username"] = arg
      elif opt in ("-a", "--account"):
         conf["account_number"] = str(arg)
      elif opt in ("-r", "--region"):
         conf["region"] = arg
      elif opt in ("-n", "--role-name"):
         conf["role_name"] = arg
      elif opt in ("-t", "--target-account"):
         conf["target_account"] = str(arg)

   if (conf['username'] == None) or \
         (conf['account_number'] == None) or \
         (conf['region'] == None):
      usage()
      print("\nERROR: username, account, or region not set\n")
      sys.exit(2)

   print(
      "\nINFO: Credentials in profile [" + profile_mfa +
      "] will be overwritten\n"
   )
   mfa_token = getpass.getpass("Enter 6-digit MFA token: ")
   mfa_serial_number = "arn:aws:iam::" + conf['account_number'] + ":mfa/" + conf['username']
   
   # get temporary credentials using long lived credentials and number generated from MFA device 
   session = boto3.Session(profile_name = profile)
   sts = session.client("sts")
   if conf['role_name']:
      print(
         "Using AssumeRole '" + conf['role_name'] +
         "' in remote account '" + conf['target_account'] + "'"
      )
      role_name =  "arn:aws:iam::" + conf['target_account'] + ":role/" + conf['role_name']
      role_session_name = conf['username'] + "@" + conf['target_account']
      response = sts.assume_role(
         RoleArn = role_name,
         RoleSessionName = role_session_name,
         SerialNumber = mfa_serial_number,
         TokenCode = str(mfa_token),
         DurationSeconds = 28800
      )
   else:
      print(
         "Using GetSessionToken to acquire temporary credentials in account '" + 
         conf['account_number'] + "'"
      )
      response = sts.get_session_token(
         DurationSeconds = 28800,
         SerialNumber = mfa_serial_number,
         TokenCode = mfa_token
      )
   sts_token = response['Credentials']
   
   # save config and temp credentials
   try:
      config.add_section(profile_mfa)
   except configparser.DuplicateSectionError:
      pass
   print(sts_token)
   config.set(profile_mfa, "aws_access_key_id", sts_token['AccessKeyId'])
   config.set(profile_mfa, "aws_secret_access_key", sts_token['SecretAccessKey'])
   config.set(profile_mfa, "aws_security_token", sts_token['SessionToken'])
   config.set(profile_mfa, "aws_session_token", sts_token['SessionToken'])
   config.set(profile_mfa, "expiration", str(sts_token['Expiration']))
   for section in [ profile, profile_mfa ]:
      for c in conf:
         if conf[c]:
            config.set(section, c, str(conf[c]))
   with open(aws_creds, 'w') as configfile:
      config.write(configfile)

   print(
      "Temporary credentials updated in [" + profile_mfa +
      "] profiles in ~/.aws/credentials"
   )

if __name__ == "__main__":
   main(sys.argv[1:])
