AWS Command Line MFA Login

mfa-login.py

#!/usr/bin/env python3

import dateutil.parser, dateutil.tz, datetime, json, os, sys, boto3, stat
from subprocess import Popen, PIPE

# Get the profile, first argument on the command line
if (len(sys.argv) < 2):
  raise Exception('Profile name not specified on command line')
profile = sys.argv[1]

# Try to read our "~/.aws/mfa-session-<profile>.json" file containing our session
config_file = os.path.expanduser('~/.aws/mfa-session-%s.json' % profile)
try:
  with open(config_file) as json_session:
    session = json.load(json_session)
except:
  session = {}

# Parse the expiration of the session, defaulting to the epoch
expiration = dateutil.parser.parse(session.get('Expiration', '1970-01-01T00:00:00.000Z'))

# If the session is expired, contact AWS to create a new session
if (expiration < datetime.datetime.now(datetime.timezone.utc)):

  # Create a STS client with our profile
  session = boto3.Session(profile_name=profile)
  client = session.client('sts')

  # Get the caller identity, and figure out its MFA ARN
  caller = client.get_caller_identity()
  arn = caller.get('Arn')
  mfa = arn.replace(':user/', ':mfa/')

  # Use AppleScript to prompt a nice dialog, aws cli doesn't pass us stdout/stderr
  script = b'return the text returned of (display dialog "MFA Token Code" with title "AWS Authentication" default answer "" with icon caution)'
  dialog = Popen(['osascript', '-'], stdin=PIPE, stdout=PIPE, stderr=PIPE)
  stdout, stderr = dialog.communicate(script)
  if (dialog.returncode != 0):
    raise Exception(stderr.decode('utf-8'))

  # The code is returned by the dialog in stdout
  code = stdout.decode('utf-8').strip()

  # Get a session from STS, using the MFA token code we got earlier
  response = client.get_session_token(SerialNumber=mfa, TokenCode=code)
  credentials = response.get('Credentials')
  session = {
    'Version': 1,
    'AccessKeyId': credentials.get('AccessKeyId'),
    'SecretAccessKey': credentials.get('SecretAccessKey'),
    'SessionToken': credentials.get('SessionToken'),
    'Expiration': credentials.get('Expiration').isoformat()
  }

  # Save this session in our "~/.aws/mfa-session-<profile>.json" file
  with open(config_file, 'w') as json_file:
    os.fchmod(json_file.fileno(), stat.S_IREAD | stat.S_IWRITE)
    json.dump(session, json_file, indent=2)
    json_file.write('\n')
    json_file.flush()


# Let the aws cli command get the JSON
json.dump(session, sys.stdout, indent=2)
sys.stdout.write('\n')
sys.stdout.flush()