This script will take the output of an aws sts assume-role and automatically export it into your Env Vars.

set_sts_tokens.py

import sys
import json

def main(PIPED_IN):
    """
    """
    a = json.loads(PIPED_IN)
    env_var_values = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN"]
    sts_values = ["AccessKeyId", "SecretAccessKey", "SessionToken"]

    for key,value in zip(env_var_values, sts_values):
       print(f"export {key}={a['Credentials'][value]}")
    return

if __name__ == "__main__":
    PIPED_IN = ""
    for line in sys.stdin:
        PIPED_IN += line
    main(PIPED_IN)