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)

This script simplifies updating the AWS credentials file (i.e., ~/.aws/credentials) for AWS SSO users. It will update the AWS credentials file by adding/updating the specified profile credentials using the AWS CLI v2 cached SSO login.

Setup

• Install AWS CLI v2
• Configure AWS Profile(s) for AWS SSO
• Install Python 3
• Install the Latest AWS Official Python SDK: boto3
  • The older versions of the boto3 package does not include the required SSO client.
• Save this Script Locally as aws_sso.py

I recommend aliasing or symlinking the script to make it easier to call.

How to Use

• If Not Already, Login SSO Using AWS CLI v2
• Execute Script: python3 aws_sso.py <profile-name>
  • <profile-name> defaults to default if not specified
• Profile is Ready to Use

Example

This example presumes the setup is completed with 3 profiles -- default, dev, and prod -- defined in the ~/.aws/config file. All the profiles are configured with the same SSO url and region.

# login to sso for aws cli
aws sso login
# create creds for default profile
python3 aws_sso.py
# create creds for my `dev` profile
python3 aws_sso.py dev
# create creds for my `prod` profile
python3 aws_sso.py prod

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>