AWS SSO Credentials File Updater for AWS SDKs

aws_sso.py

#!/usr/bin/env python3

import json
import os
import sys
from configparser import ConfigParser
from datetime import datetime
from pathlib import Path

import boto3

AWS_CONFIG_PATH = f"{Path.home()}/.aws/config"
AWS_CREDENTIAL_PATH = f"{Path.home()}/.aws/credentials"
AWS_SSO_CACHE_PATH = f"{Path.home()}/.aws/sso/cache"
AWS_DEFAULT_REGION = "us-east-1"


# -------------------------------------------------------------------- main ---


def set_profile_credentials(profile_name):
    profile = get_aws_profile(profile_name)
    cache_login = get_sso_cached_login(profile)
    credentials = get_sso_role_credentials(profile, cache_login)
    update_aws_credentials(profile_name, profile, credentials)


# --------------------------------------------------------------------- fns ---


def get_sso_cached_login(profile):
    file_paths = list_directory(AWS_SSO_CACHE_PATH)
    time_now = datetime.now()
    for file_path in file_paths:
        data = load_json(file_path)
        if data is None:
            continue
        if data.get("startUrl") != profile["sso_start_url"]:
            continue
        if data.get("region") != profile["sso_region"]:
            continue
        if time_now > parse_timestamp(data.get("expiresAt", "1970-01-01T00:00:00UTC")):
            continue
        return data
    raise Exception("Current cached SSO login is expired or invalid")


def get_sso_role_credentials(profile, login):
    client = boto3.client("sso", region_name=profile["sso_region"])
    response = client.get_role_credentials(
        roleName=profile["sso_role_name"],
        accountId=profile["sso_account_id"],
        accessToken=login["accessToken"],
    )
    return response["roleCredentials"]


def get_aws_profile(profile_name):
    config = read_config(AWS_CONFIG_PATH)
    profile_opts = config.items(f"profile {profile_name}")
    profile = dict(profile_opts)
    return profile


def update_aws_credentials(profile_name, profile, credentials):
    region = profile.get("region", AWS_DEFAULT_REGION)
    config = read_config(AWS_CREDENTIAL_PATH)
    if config.has_section(profile_name):
        config.remove_section(profile_name)
    config.add_section(profile_name)
    config.set(profile_name, "region", region)
    config.set(profile_name, "aws_access_key_id", credentials["accessKeyId"])
    config.set(profile_name, "aws_secret_access_key", credentials["secretAccessKey"])
    config.set(profile_name, "aws_session_token", credentials["sessionToken"])
    write_config(AWS_CREDENTIAL_PATH, config)


# ------------------------------------------------------------------- utils ---


def list_directory(path):
    file_paths = []
    if os.path.exists(path):
        file_paths = Path(path).iterdir()
    file_paths = sorted(file_paths, key=os.path.getmtime)
    file_paths.reverse()  # sort by recently updated
    return [str(f) for f in file_paths]


def load_json(path):
    try:
        with open(path) as context:
            return json.load(context)
    except ValueError:
        pass  # ignore invalid json


def parse_timestamp(value):
    return datetime.strptime(value, "%Y-%m-%dT%H:%M:%SUTC")


def read_config(path):
    config = ConfigParser()
    config.read(path)
    return config


def write_config(path, config):
    with open(path, "w") as destination:
        config.write(destination)


# ---------------------------------------------------------------- handlers ---


def script_handler(args):
    profile_name = "default"
    if len(args) == 2:
        profile_name = args[1]
    set_profile_credentials(profile_name)


if __name__ == "__main__":
    script_handler(sys.argv)

So I liked what you did here! Not sure if you know, but most SDKs and the CLI have support for a credential_process value in a Profile: https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-sourcing-external.html which your script could be easily adapted for - with the small caveat that it would need to do some file caching, because apps can ask for credentials like an insane person.

This means, for every SSO Profile we could create a "wrapper" profile, which would call your script to source credentials, and ta-da - now software with our current SDKs have a way to get credentials via SSO (until your SSO token expires of course).

I wanted a slightly more portable version of this tool, so I took a stab at writing one in golang: https://github.com/flyinprogrammer/aws-sso-fetcher - it seems to work, but it could be total garbage. Issues and feedback welcome!

Thanks for writing up this script, it made my work much simpler.

I will review the credential_process value as you outline and your golang project soon. Thank you for your comment.

This works like a charm for me, thank you!
Note that after using this once, any aws cli command with the affected profile will look at the credentials file first. This means you either need to run this script every time you login via sso or you manually delete the profile's credentials from the credentials file again.

If you just want to read the sso credentials and then pass them on, e.g. to a docker container, you could disable the "update_aws_credentials" line, print the credentials to console and then read them in a shell script, like so:

credentials=$(python3 aws_sso.py <your-profile>)
session_token=$(echo $credentials | jq -r '.sessionToken')
access_key_id=$(echo $credentials | jq -r '.accessKeyId')
secret_access_key=$(echo $credentials | jq -r '.secretAccessKey')

docker run -e AWS_SESSION_TOKEN=$session_token -e AWS_ACCESS_KEY_ID=$access_key_id -e AWS_SECRET_ACCESS_KEY=$secret_access_key -e AWS_DEFAULT_REGION=$(aws configure get region --profile <your-profile>) -t <image-tag>