Create a new token auth AWS session using an MFA device

awsmfa.py

#!/bin/env python3
"""
Create a new token auth AWS session using an MFA device.

Usage:
  $(python3 awsmfa.py 123456)
  
You can also install this into your shell profile, like ~/.bashrc
```
setawsmfa() {
        eval $(python3 ~/bin/awsmfa.py $1)
}
```
Then usage becomes `setawsmfa 123456`
"""

import os
import sys
import subprocess
import json

# Set your MFA device
mfa_arn = "arn:aws:iam::..."

def maybeUnsetEnv(key):
    try:
        del os.environ[key]
    except KeyError:
        pass

def set_env(mfaDeviceToken):
    args = [
            "aws", "sts", "get-session-token",
            "--serial-number", mfa_arn,
            "--token-code", mfaDeviceToken
            ]

    maybeUnsetEnv('AWS_ACCESS_KEY_ID')
    maybeUnsetEnv('AWS_SECRET_ACCESS_KEY')
    maybeUnsetEnv('AWS_SESSION_TOKEN ')
    cmd = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)

    creds = json.loads(cmd.stdout)['Credentials']

    print('export AWS_ACCESS_KEY_ID=' + creds['AccessKeyId'])
    print('export AWS_SECRET_ACCESS_KEY=' + creds['SecretAccessKey'])
    print('export AWS_SESSION_TOKEN=' + creds['SessionToken'])


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