get_mfa_credentials.py

"""
This script gets a session token for AWS and automatically updates the information in your credentials file.

ATTENTION: This script changes the first three lines after the [mfa] tag in the AWS credentials file. If you have
change something specifically at your credentials file that for example added additional attributes be careful
when using this script.

To make the script work:
Copy to the folder where your aws credentials are stored.
Call with 'python get_mfa_credentials {current mfa token} {your serial number}'.
You can find your serial number by entering 'aws iam list-virtual-mfa-devices' in your terminal.
Since your serial number does not change between calls, you can also locally hardcode the assignment, and
only call the script with your 'current mfa token'.
"""
import os
import pathlib
import re
import subprocess
import sys


if __name__ == "__main__":
    # Parameters go get session token from aws
    mfa_key = input("MFA CODE: ")
    serial_number = sys.argv[1]

    # Get path to credentials file
    directory = os.path.expanduser("~/.aws")
    credentials = os.path.join(directory, "credentials")

    # Read current credentials
    with open(credentials, "r") as f:
        data = f.readlines()

    # Find mfa entry if any available
    mfa_idx = None
    for i, line in enumerate(data):
        if line.strip() == "[mfa]":
            mfa_idx = i

    # Get the session token from Aws
    command = f"aws sts get-session-token --serial-number {serial_number} --token-code {mfa_key}"
    p = subprocess.Popen(
        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )
    output = p.stdout.read()
    output = str(output)

    # Parse answer string with regex
    AccessKeyId = re.search('"AccessKeyId": "([^"]*)"', output).groups()[0]
    SecretAccessKey = re.search('"SecretAccessKey": "([^"]*)"', output).groups()[0]
    SessionToken = re.search('"SessionToken": "([^"]*)"', output).groups()[0]

    # If index did not exist yet, append to the existing file
    if mfa_idx is None:
        data.append("\n")
        data.append("[mfa]\n")
        data.append(f"aws_access_key_id = {AccessKeyId}\n")
        data.append(f"aws_secret_access_key = {SecretAccessKey}\n")
        data.append(f"aws_session_token = {SessionToken}\n")
    # Else replace the old information
    else:
        data[mfa_idx + 1] = f"aws_access_key_id = {AccessKeyId}\n"
        data[mfa_idx + 2] = f"aws_secret_access_key = {SecretAccessKey}\n"
        data[mfa_idx + 3] = f"aws_session_token = {SessionToken}\n"

    # Recreate the data
    data = "".join(data)

    # Store the credentials
    with open(credentials, "w") as f:
        f.write(data)